modify pipeline.py

track_reid.py

@@ -62,6 +62,7 @@ from hands.hand_inference import hand_pose
 
 from contrast.feat_extract.config import config as conf
 from contrast.feat_extract.inference import FeatsInterface
+
 ReIDEncoder = FeatsInterface(conf)
 
 IMG_FORMATS = '.bmp', '.dng', '.jpeg', '.jpg', '.mpo', '.png', '.tif', '.tiff', '.webp', '.pfm'  # include image suffixes
@@ -139,12 +140,8 @@ def init_trackers(tracker_yaml = None, bs=1):
 def yolo_resnet_tracker(
         weights=ROOT / 'yolov5s.pt',  # model path or triton URL
         source=ROOT / 'data/images',  # file/dir/URL/glob/screen/0(webcam)
-        
-        project=ROOT / 'runs/detect',  # save results to project/name
-        name='exp',  # save results to project/name
         save_dir = '',
-        
-        is_save_img = False,
+        is_save_img = True,
         is_save_video = True,
         
         tracker_yaml = "./tracking/trackers/cfg/botsort.yaml",
@@ -154,19 +151,9 @@ def yolo_resnet_tracker(
         max_det=1000,  # maximum detections per image
         device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
 
-        view_img=False,  # show results
-        save_txt=False,  # save results to *.txt
-        save_csv=False,  # save results in CSV format
-        save_conf=False,  # save confidences in --save-txt labels
-        save_crop=False,  # save cropped prediction boxes
-        nosave=False,  # do not save images/videos
-        update=False,  # update all models
-        exist_ok=False,  # existing project/name ok, do not increment
-
         classes=None,  # filter by class: --class 0, or --class 0 2 3
         agnostic_nms=False,  # class-agnostic NMS
         augment=False,  # augmented inference
-        visualize=False,  # visualize features
         
         line_thickness=3,  # bounding box thickness (pixels)
         hide_labels=False,  # hide labels
@@ -208,7 +195,7 @@ def yolo_resnet_tracker(
 
         # Inference
         with dt[1]:
-            visualize = increment_path(project / Path(path).stem, mkdir=True) if visualize else False
+            # visualize = increment_path(project / Path(path).stem, mkdir=True) if visualize else False
             pred = model(im, augment=augment, visualize=False)
 
         # NMS
@@ -242,15 +229,16 @@ def yolo_resnet_tracker(
                     # trackerBoxes = np.concatenate([trackerBoxes, tracks], axis=0)
                     '''================== 1. 存储 dets/subimgs/features Dict ============='''
                     imgs, features = ReIDEncoder.inference(im0, tracks)
-                    featdict = {}
+                    imgdict, featdict = {}, {}
                     for ii, bid in enumerate(tracks[:, 8]):
                         featdict.update({f"{int(frameId)}_{int(bid)}": features[ii, :]})            # [f"feat_{int(bid)}"] = features[i, :]
+                        imgdict.update({f"{int(frameId)}_{int(bid)}": imgs[ii]})
     
                     frameDict = {"path": path, 
                                  "fid": int(frameId),
                                  "bboxes": det,
                                  "tboxes": tracks,
-                                 "imgs": imgs,
+                                 "imgs": imgdict,
                                  "feats": featdict} 
                     yoloResnetTracker.append(frameDict)