232 lines
7.0 KiB
Python
232 lines
7.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Thu Mar 27 16:09:07 2025
|
|
|
|
@author: ym
|
|
"""
|
|
import os
|
|
import sys
|
|
import cv2
|
|
import pickle
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from scipy.spatial.distance import cdist
|
|
import copy
|
|
|
|
from dotrack.dotracks_back import doBackTracks
|
|
from dotrack.dotracks_front import doFrontTracks
|
|
from utils.drawtracks import plot_frameID_y2, draw_all_trajectories
|
|
from utils.read_data import read_similar
|
|
|
|
|
|
def get_trail(ShoppingDict, ppath):
|
|
|
|
evtname = ShoppingDict["eventName"]
|
|
|
|
back_yrt = ShoppingDict["backCamera"]["yoloResnetTracker"]
|
|
front_yrt = ShoppingDict["frontCamera"]["yoloResnetTracker"]
|
|
|
|
back_vts = ShoppingDict["frontCamera"]["tracking"]
|
|
front_vts = ShoppingDict["backCamera"]["tracking"]
|
|
|
|
|
|
event_tracks = [("back", back_yrt, back_vts), ("front", front_yrt, front_vts)]
|
|
|
|
|
|
savepath = ppath / "alltrail"
|
|
if not savepath.exists():
|
|
savepath.mkdir()
|
|
|
|
savepath = str(savepath)
|
|
evtime = evtname[:15]
|
|
|
|
illus = [None, None]
|
|
for camera_type, yrtOut, vts in event_tracks:
|
|
if len(vts.Residual)==1: continue
|
|
|
|
if camera_type == 'front':
|
|
edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
|
|
|
|
img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
|
|
illus[0] = img_tracking
|
|
|
|
plt = plot_frameID_y2(vts)
|
|
plt.savefig(os.path.join(savepath, f"{evtime}_front.png"))
|
|
|
|
if camera_type == 'back':
|
|
edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
|
|
|
|
img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
|
|
illus[1] = img_tracking
|
|
|
|
illus = [im for im in illus if im is not None]
|
|
if len(illus):
|
|
img_cat = np.concatenate(illus, axis = 1)
|
|
if len(illus)==2:
|
|
H, W = img_cat.shape[:2]
|
|
cv2.line(img_cat, (int(W/2), 0), (int(W/2), int(H)), (128, 128, 255), 3)
|
|
|
|
trajpath = os.path.join(savepath, f"{evtime}.png")
|
|
cv2.imwrite(trajpath, img_cat)
|
|
|
|
return evtime
|
|
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def track_opt(ShoppingDict, ppath):
|
|
'''
|
|
将一个对象读取，修改其中一个属性
|
|
|
|
'''
|
|
evtname = ShoppingDict["eventName"]
|
|
|
|
shopping = copy.deepcopy(ShoppingDict)
|
|
|
|
## only need to init item: tracking for each Camera
|
|
shopping["frontCamera"]["tracking"] = []
|
|
shopping["backCamera"]["tracking"] = []
|
|
|
|
back_camera = ShoppingDict["backCamera"]["cameraType"]
|
|
back_yrt = ShoppingDict["backCamera"]["yoloResnetTracker"]
|
|
front_camera = ShoppingDict["frontCamera"]["cameraType"]
|
|
front_yrt = ShoppingDict["frontCamera"]["yoloResnetTracker"]
|
|
yrts = [(back_camera, back_yrt), (front_camera, front_yrt)]
|
|
|
|
event_tracks = []
|
|
errtrail = ''
|
|
for camera_type, yrtOut in yrts:
|
|
'''================= 1. tracking ================='''
|
|
'''(1) 生成用于 tracking 模块的 boxes、feats'''
|
|
# bboxes = np.empty((0, 6), dtype=np.float64)
|
|
trackerboxes = np.empty((0, 9), dtype=np.float64)
|
|
trackefeats = {}
|
|
for frameDict in yrtOut:
|
|
tboxes = frameDict["tboxes"]
|
|
ffeats = frameDict["feats"]
|
|
|
|
# boxes = frameDict["bboxes"]
|
|
# bboxes = np.concatenate((bboxes, np.array(boxes)), axis=0)
|
|
trackerboxes = np.concatenate((trackerboxes, np.array(tboxes)), axis=0)
|
|
for i in range(len(tboxes)):
|
|
fid, bid = int(tboxes[i, 7]), int(tboxes[i, 8])
|
|
trackefeats.update({f"{fid}_{bid}": ffeats[f"{fid}_{bid}"]})
|
|
|
|
'''(2) tracking, 后摄'''
|
|
if camera_type == "back":
|
|
vts = doBackTracks(trackerboxes, trackefeats)
|
|
vts.classify()
|
|
shopping["backCamera"]["tracking"] = vts
|
|
|
|
if len(vts.Residual)!=1:
|
|
errtrail = evtname
|
|
|
|
'''(3) tracking, 前摄'''
|
|
if camera_type == "front":
|
|
vts = doFrontTracks(trackerboxes, trackefeats)
|
|
vts.classify()
|
|
shopping["frontCamera"]["tracking"] = vts
|
|
|
|
if len(vts.Residual)!=1:
|
|
errtrail = evtname
|
|
|
|
event_tracks.append((camera_type, yrtOut, vts))
|
|
|
|
|
|
pckpath = ppath / "track_optim"
|
|
if not pckpath.exists():
|
|
pckpath.mkdir()
|
|
|
|
fpath = pckpath / "{}_new.pickle".format(evtname)
|
|
with open(str(fpath), 'wb') as f:
|
|
pickle.dump(shopping, f)
|
|
|
|
|
|
savepath = ppath / "yolos_tracking" / evtname
|
|
|
|
illus = [None, None]
|
|
for camera_type, yrtOut, vts in event_tracks:
|
|
if len(vts.tracks)==0: continue
|
|
|
|
if camera_type == 'front':
|
|
edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
|
|
|
|
img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
|
|
illus[0] = img_tracking
|
|
|
|
plt = plot_frameID_y2(vts)
|
|
plt.savefig(os.path.join(savepath, "front_y2_new.png"))
|
|
|
|
if camera_type == 'back':
|
|
edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
|
|
|
|
img_tracking = draw_all_trajectories(vts, edgeline, savepath, camera_type, draw5p=False)
|
|
illus[1] = img_tracking
|
|
|
|
illus = [im for im in illus if im is not None]
|
|
if len(illus):
|
|
img_cat = np.concatenate(illus, axis = 1)
|
|
if len(illus)==2:
|
|
H, W = img_cat.shape[:2]
|
|
cv2.line(img_cat, (int(W/2), 0), (int(W/2), int(H)), (128, 128, 255), 3)
|
|
|
|
trajpath = os.path.join(savepath, "trajectory_new.png")
|
|
cv2.imwrite(trajpath, img_cat)
|
|
|
|
return errtrail
|
|
|
|
|
|
def main():
|
|
# evttypes = ["single_event_V10", "single_event_V5", "performence_V10", "performence_V5"]
|
|
evttypes = ["single_event_V10"]
|
|
|
|
k = 0
|
|
error_trail = []
|
|
for evttype in evttypes:
|
|
ppath = Path("/home/wqg/dataset/pipeline/yrt/{}".format(evttype))
|
|
|
|
pkpath = ppath / "shopping_pkl"
|
|
for fp in pkpath.iterdir():
|
|
# fp = pkpath / "{}.pickle".format("20250305-152917-635_6970209860221_6970209860221")
|
|
print(fp)
|
|
|
|
if fp.suffix != '.pickle': continue
|
|
with open(str(fp), 'rb') as f:
|
|
ShoppingDict = pickle.load(f)
|
|
|
|
# errtrail = track_opt(ShoppingDict, ppath)
|
|
# error_trail.append(errtrail)
|
|
|
|
errtrail = get_trail(ShoppingDict, ppath)
|
|
if errtrail is not None:
|
|
error_trail.append(errtrail)
|
|
|
|
# k+=1
|
|
# if k==100:
|
|
# break
|
|
|
|
errfile = ppath / 'error_trail.txt'
|
|
with open(errfile, 'w', encoding='utf-8') as f:
|
|
for line in error_trail:
|
|
f.write(line + '\n')
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|