# -*- coding: utf-8 -*- """ Created on Sun Sep 29 08:59:21 2024 @author: ym """ import os # import sys import cv2 import pickle import numpy as np from pathlib import Path from scipy.spatial.distance import cdist from track_reid import yolo_resnet_tracker, yolov10_resnet_tracker from tracking.dotrack.dotracks_back import doBackTracks from tracking.dotrack.dotracks_front import doFrontTracks from tracking.utils.drawtracks import plot_frameID_y2, draw_all_trajectories from utils.getsource import get_image_pairs, get_video_pairs from tracking.utils.read_data import read_similar def save_subimgs(imgdict, boxes, spath, ctype, featdict = None): ''' 当前 box 特征和该轨迹前一个 box 特征的相似度，可用于和跟踪序列中的相似度进行比较 ''' boxes = boxes[np.argsort(boxes[:, 7])] for i in range(len(boxes)): simi = None tid, fid, bid = int(boxes[i, 4]), int(boxes[i, 7]), int(boxes[i, 8]) if i>0: _, fid0, bid0 = int(boxes[i-1, 4]), int(boxes[i-1, 7]), int(boxes[i-1, 8]) if f"{fid0}_{bid0}" in featdict.keys() and f"{fid}_{bid}" in featdict.keys(): feat0 = featdict[f"{fid0}_{bid0}"] feat1 = featdict[f"{fid}_{bid}"] simi = 1 - np.maximum(0.0, cdist(feat0[None, :], feat1[None, :], "cosine"))[0][0] img = imgdict[f"{fid}_{bid}"] imgpath = spath / f"{ctype}_tid{tid}-{fid}-{bid}.png" if simi is not None: imgpath = spath / f"{ctype}_tid{tid}-{fid}-{bid}_sim{simi:.2f}.png" cv2.imwrite(imgpath, img) def save_subimgs_1(imgdict, boxes, spath, ctype, simidict = None): ''' 当前 box 特征和该轨迹 smooth_feat 特征的相似度, yolo_resnet_tracker 函数中， 采用该方式记录特征相似度 ''' for i in range(len(boxes)): tid, fid, bid = int(boxes[i, 4]), int(boxes[i, 7]), int(boxes[i, 8]) key = f"{fid}_{bid}" img = imgdict[key] imgpath = spath / f"{ctype}_tid{tid}-{fid}-{bid}.png" if simidict is not None and key in simidict.keys(): imgpath = spath / f"{ctype}_tid{tid}-{fid}-{bid}_sim{simidict[key]:.2f}.png" cv2.imwrite(imgpath, img) def show_result(event_tracks, yrtDict, savepath_pipe): '''保存 Tracking 输出的运动轨迹子图，并记录相似度''' savepath_pipe_subimgs = savepath_pipe / Path("subimgs") if not savepath_pipe_subimgs.exists(): savepath_pipe_subimgs.mkdir(parents=True, exist_ok=True) for CamerType, vts in event_tracks: if len(vts.tracks)==0: continue if CamerType == 'front': # yolos = ShoppingDict["frontCamera"]["yoloResnetTracker"] yolos = yrtDict["frontyrt"] ctype = 1 if CamerType == 'back': # yolos = ShoppingDict["backCamera"]["yoloResnetTracker"] yolos = yrtDict["backyrt"] ctype = 0 imgdict, featdict, simidict = {}, {}, {} for y in yolos: imgdict.update(y["imgs"]) featdict.update(y["feats"]) simidict.update(y["featsimi"]) for track in vts.Residual: if isinstance(track, np.ndarray): save_subimgs(imgdict, track, savepath_pipe_subimgs, ctype, featdict) else: save_subimgs(imgdict, track.slt_boxes, savepath_pipe_subimgs, ctype, featdict) '''(3) 轨迹显示与保存''' illus = [None, None] for CamerType, vts in event_tracks: if len(vts.tracks)==0: continue if CamerType == 'front': edgeline = cv2.imread("./tracking/shopcart/cart_tempt/board_ftmp_line.png") h, w = edgeline.shape[:2] # nh, nw = h//2, w//2 # edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA) img_tracking = draw_all_trajectories(vts, edgeline, savepath_pipe, CamerType, draw5p=True) illus[0] = img_tracking plt = plot_frameID_y2(vts) plt.savefig(os.path.join(savepath_pipe, "front_y2.png")) if CamerType == 'back': edgeline = cv2.imread("./tracking/shopcart/cart_tempt/edgeline.png") h, w = edgeline.shape[:2] # nh, nw = h//2, w//2 # edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA) img_tracking = draw_all_trajectories(vts, edgeline, savepath_pipe, CamerType, draw5p=True) 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_pipe, "trajectory.png") cv2.imwrite(trajpath, img_cat) def pipeline(eventpath, SourceType, weights, DataType = "raw", #raw, pkl: images or videos, pkl, pickle file YoloVersion="V5", savepath = None, saveimages = True ): ## 构造购物事件字典 evtname = Path(eventpath).stem barcode = evtname.split('_')[-1] if len(evtname.split('_'))>=2 \ and len(evtname.split('_')[-1])>=8 \ and evtname.split('_')[-1].isdigit() else '' '''事件结果存储文件夹: savepath_pipe, savepath_pkl''' if not savepath: savepath = Path(__file__).resolve().parents[0] / "events_result" savepath_pipe = Path(savepath) / Path("yolos_tracking") / evtname savepath_pkl = Path(savepath) / "shopping_pkl" if not savepath_pkl.exists(): savepath_pkl.mkdir(parents=True, exist_ok=True) pklpath = Path(savepath_pkl) / Path(str(evtname)+".pickle") yrt_out = [] if DataType == "raw": ### 不重复执行已经过yolo-resnet-tracker if pklpath.exists(): print(f"Pickle file have saved: {evtname}.pickle") return if SourceType == "video": vpaths = get_video_pairs(eventpath) elif SourceType == "image": vpaths = get_image_pairs(eventpath) for vpath in vpaths: '''================= 2. 事件结果存储文件夹 =================''' if isinstance(vpath, list): savepath_pipe_imgs = savepath_pipe / Path("images") else: savepath_pipe_imgs = savepath_pipe / Path(str(Path(vpath).stem)) if not savepath_pipe_imgs.exists(): savepath_pipe_imgs.mkdir(parents=True, exist_ok=True) optdict = {} optdict["weights"] = weights optdict["source"] = vpath optdict["save_dir"] = savepath_pipe_imgs optdict["is_save_img"] = saveimages optdict["is_save_video"] = True if YoloVersion == "V5": yrtOut = yolo_resnet_tracker(**optdict) elif YoloVersion == "V10": yrtOut = yolov10_resnet_tracker(**optdict) yrt_out.append((vpath, yrtOut)) elif DataType == "pkl": pass else: return '''====================== 构造 ShoppingDict 模块 =======================''' ShoppingDict = {"eventPath": eventpath, "eventName": evtname, "barcode": barcode, "eventType": '', # "input", "output", "other" "frontCamera": {}, "backCamera": {}, "one2n": [] # } procpath = Path(eventpath).joinpath('process.data') if procpath.is_file(): SimiDict = read_similar(procpath) ShoppingDict["one2n"] = SimiDict['one2n'] yrtDict = {} event_tracks = [] for vpath, yrtOut in yrt_out: '''================= 1. 构造相机事件字典 =================''' CameraEvent = {"cameraType": '', # "front", "back" "videoPath": '', "imagePaths": [], "yoloResnetTracker": [], "tracking": [], } if isinstance(vpath, list): CameraEvent["imagePaths"] = vpath bname = os.path.basename(vpath[0]) if not isinstance(vpath, list): CameraEvent["videoPath"] = vpath bname = os.path.basename(vpath).split('.')[0] if bname.split('_')[0] == "0" or bname.find('back')>=0: CameraEvent["cameraType"] = "back" if bname.split('_')[0] == "1" or bname.find('front')>=0: CameraEvent["cameraType"] = "front" '''2种保存方式: (1) no save subimg, (2) save img''' ###(1) save images yrtOut_save = [] for frdict in yrtOut: fr_dict = {} for k, v in frdict.items(): if k != "imgs": fr_dict[k]=v yrtOut_save.append(fr_dict) CameraEvent["yoloResnetTracker"] = yrtOut_save ###(2) no save images # CameraEvent["yoloResnetTracker"] = yrtOut '''================= 4. 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 CameraEvent["cameraType"] == "back": vts = doBackTracks(trackerboxes, trackefeats) vts.classify() event_tracks.append(("back", vts)) CameraEvent["tracking"] = vts ShoppingDict["backCamera"] = CameraEvent yrtDict["backyrt"] = yrtOut '''(2) tracking, 前摄''' if CameraEvent["cameraType"] == "front": vts = doFrontTracks(trackerboxes, trackefeats) vts.classify() event_tracks.append(("front", vts)) CameraEvent["tracking"] = vts ShoppingDict["frontCamera"] = CameraEvent yrtDict["frontyrt"] = yrtOut '''========================== 保存模块 =================================''' # 保存 ShoppingDict with open(str(pklpath), 'wb') as f: pickle.dump(ShoppingDict, f) # 绘制并保存轨迹图 show_result(event_tracks, yrtDict, savepath_pipe) def execute_pipeline(evtdir = r"D:\datasets\ym\后台数据\unzip", DataType = "raw", # raw, pkl save_path = r"D:\work\result_pipeline", kk=1, source_type = "video", # video, image, yolo_ver = "V10", # V10, V5 weight_yolo_v5 = r'./ckpts/best_cls10_0906.pt' , weight_yolo_v10 = r'./ckpts/best_v10s_width0375_1205.pt', saveimages = True ): ''' 运行函数 pipeline()，遍历事件文件夹，每个文件夹是一个事件 ''' parmDict = {} parmDict["DataType"] = DataType parmDict["savepath"] = save_path parmDict["SourceType"] = source_type parmDict["YoloVersion"] = yolo_ver if parmDict["YoloVersion"] == "V5": parmDict["weights"] = weight_yolo_v5 elif parmDict["YoloVersion"] == "V10": parmDict["weights"] = weight_yolo_v10 parmDict["saveimages"] = saveimages evtdir = Path(evtdir) errEvents = [] k = 0 for item in evtdir.iterdir(): if item.is_dir(): # item = evtdir/Path("20241212-171505-f0afe929-fdfe-4efa-94d0-2fa748d65fbb_6907992518930") parmDict["eventpath"] = item pipeline(**parmDict) # try: # pipeline(**parmDict) # except Exception as e: # errEvents.append(str(item)) k+=1 if kk is not None and k==kk: break errfile = os.path.join(parmDict["savepath"], 'error_events.txt') with open(errfile, 'w', encoding='utf-8') as f: for line in errEvents: f.write(line + '\n') if __name__ == "__main__": datapath = r'/home/wqg/dataset/test_dataset/base_dataset/single_event/source/' savepath = r'/home/wqg/dataset/pipeline/test_result/single_event_V10' execute_pipeline(evtdir = datapath, DataType = "raw", # raw, pkl kk=1, source_type = "video", # video, image, save_path = savepath, yolo_ver = "V10", # V10, V5 weight_yolo_v5 = r'./ckpts/best_cls10_0906.pt' , weight_yolo_v10 = r'./ckpts/best_v10s_width0375_1205.pt', saveimages = False )