add yolo v10 and modify pipeline

pipeline.py

@@ -11,7 +11,7 @@ import pickle
 import numpy as np
 from pathlib import Path
 from scipy.spatial.distance import cdist
-from track_reid import yolo_resnet_tracker
+from track_reid import yolo_resnet_tracker, yolov10_resnet_tracker
 
 from tracking.dotrack.dotracks_back import doBackTracks
 from tracking.dotrack.dotracks_front import doFrontTracks 
@@ -65,22 +65,20 @@ def pipeline(
         eventpath,
         savepath,
         SourceType,
-        weights
+        weights,
+        YoloVersion="V5"
         ):
     '''
     eventpath: 单个事件的存储路径
     
     '''
+    optdict = {}
+    optdict["weights"] = weights
     
     if SourceType == "video":
         vpaths = get_video_pairs(eventpath)
     elif SourceType == "image":
         vpaths = get_image_pairs(eventpath)
-    
-    optdict = {}
-    optdict["weights"] = weights
-
-    
     event_tracks = []
     
     ## 构造购物事件字典
@@ -101,9 +99,9 @@ def pipeline(
         savepath_spdict.mkdir(parents=True, exist_ok=True)
     pf_path = Path(savepath_spdict) / Path(str(evtname)+".pickle")
 
-    if pf_path.exists():
-        print(f"Pickle file have saved: {evtname}.pickle")
-        return
+    # if pf_path.exists():
+    #     print(f"Pickle file have saved: {evtname}.pickle")
+    #     return
 
     '''====================== 构造 ShoppingDict 模块 ======================='''
     ShoppingDict = {"eventPath": eventpath, 
@@ -160,12 +158,16 @@ def pipeline(
         '''================= 3. Yolo + Resnet + Tracker ================='''
         optdict["source"] = vpath
         optdict["save_dir"] = savepath_pipeline_imgs   
-        optdict["is_save_img"] = False
+        optdict["is_save_img"] = True
         optdict["is_save_video"] = True
 
         
-        yrtOut = yolo_resnet_tracker(**optdict)
+        if YoloVersion == "V5":
+            yrtOut = yolo_resnet_tracker(**optdict)
+        elif YoloVersion == "V10":
+            yrtOut = yolov10_resnet_tracker(**optdict)
         
+
         yrtOut_save = []
         for frdict in yrtOut:
             fr_dict = {}
@@ -285,21 +287,32 @@ def pipeline(
         trajpath = os.path.join(savepath_pipeline, "trajectory.png")
         cv2.imwrite(trajpath, img_cat)
     
-def main():
+def execute_pipeline(evtdir = r"D:\datasets\ym\后台数据\unzip", 
+                     source_type = "video",   # video, image, 
+                     save_path = r"D:\work\result_pipeline", 
+                     yolo_ver = "V10",   # V10, V5
+                     
+                     weight_yolo_v5 = r'./ckpts/best_cls10_0906.pt' ,
+                     weight_yolo_v10 = r'./ckpts/best_v10s_width0375_1205.pt',
+                     k=0
+                     ):
+    '''
+    运行函数 pipeline()，遍历事件文件夹，每个文件夹是一个事件
     '''
-    函数：pipeline()，遍历事件文件夹，选择类型 image 或 video,
-    '''     
     parmDict = {}
-    evtdir = r"../dataset/backend_20250310"
-    parmDict["SourceType"] = "video"   # video, image
-    parmDict["savepath"] = r"../dataset/run"
-    parmDict["weights"] = r'./ckpts/best_cls10_0906.pt' 
+    parmDict["SourceType"] = source_type
+    parmDict["savepath"] = save_path
+    parmDict["YoloVersion"] = yolo_ver
+    if parmDict["YoloVersion"] == "V5":
+        parmDict["weights"] = weight_yolo_v5
+    elif parmDict["YoloVersion"] == "V10":
+        parmDict["weights"] = weight_yolo_v10
 
     evtdir = Path(evtdir) 
-    k, errEvents = 0, []
+    errEvents = []
     for item in evtdir.iterdir():
         if item.is_dir():            
-            # item = evtdir/Path("20250303-103058-074_6914973604223_6914973604223")
+            item = evtdir/Path("20250310-175352-741")
             parmDict["eventpath"] = item
             pipeline(**parmDict)
             # try:
@@ -307,19 +320,21 @@ def main():
             # except Exception as e:
             #     errEvents.append(str(item))
         k+=1
-        if k==5:
+        if k==1:
             break
             
-    errfile = os.path.join(parmDict["savepath"], f'error_events.txt')
+    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__": 
-    main()
+    execute_pipeline()
     
+    # spath_v10 = r"D:\work\result_pipeline_v10"
+    # spath_v5 = r"D:\work\result_pipeline_v5"
+    # execute_pipeline(save_path=spath_v10, yolo_ver="V10")
+    # execute_pipeline(save_path=spath_v5, yolo_ver="V5")