box select in a track and feat simi modify in tracker

pipeline.py

@@ -10,6 +10,7 @@ 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
 
 from tracking.dotrack.dotracks_back import doBackTracks
@@ -19,12 +20,44 @@ 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):
+def save_subimgs(imgdict, boxes, spath, ctype, featdict = None):
+    '''
+    当前 box 特征和该轨迹前一个 box 特征的相似度，可用于和跟踪序列中的相似度进行比较
+    '''
+    boxes = boxes[np.argsort(boxes[:, 7])]
     for i in range(len(boxes)):
-        fid, bid = int(boxes[i, 7]), int(boxes[i, 8])
+        simi = None
-        if f"{fid}_{bid}" in imgdict.keys():
+        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}_{fid}_{bid}.png"
+        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)
         
 
@@ -177,15 +210,18 @@ def pipeline(
             yolos = ShoppingDict["backCamera"]["yoloResnetTracker"]
             ctype = 0
             
-        imgdict = {}
+        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_pipeline_subimgs, ctype)
+                save_subimgs(imgdict, track, savepath_pipeline_subimgs, ctype, featdict)
             else:
-                save_subimgs(imgdict, track.boxes, savepath_pipeline_subimgs,  ctype)
+                save_subimgs(imgdict, track.slt_boxes, savepath_pipeline_subimgs,  ctype, featdict)
   
 
     '''轨迹显示模块'''
@@ -243,14 +279,14 @@ def main():
         if item.is_dir():            
             # item = evtdir/Path("20241209-160201-b97f7a0e-7322-4375-9f17-c475500097e9_6926265317292")
             parmDict["eventpath"] = item
-            # pipeline(**parmDict)
-            
-            try:
             pipeline(**parmDict)
-            except Exception as e:
+            
-                errEvents.append(str(item))
+            # try:
+            #     pipeline(**parmDict)
+            # except Exception as e:
+            #     errEvents.append(str(item))
         k+=1
-        if k==1:
+        if k==2:
             break
             
     errfile = os.path.join(parmDict["savepath"], f'error_events.txt')

track_reid.py

@@ -38,6 +38,7 @@ import glob
 import numpy as np
 import pickle
 import torch
+from scipy.spatial.distance import cdist
 
 FILE = Path(__file__).resolve()
 ROOT = FILE.parents[0]  # YOLOv5 root directory
@@ -222,7 +223,19 @@ def yolo_resnet_tracker(
                           这里，frame_index 也可以用视频的 帧ID 代替, box_index 保持不变
                 '''
                 det_tracking = Boxes(det, im0.shape).cpu().numpy()
-                tracks = tracker.update(det_tracking, im0)
+                tracks, outfeats = tracker.update(det_tracking, im0)
+                
+                simdict, simdict1  = {}, {}
+                for fid, bid, mfeat, cfeat, features in outfeats:
+                    if mfeat is not None and cfeat is not None:
+                        simi = 1 - np.maximum(0.0, cdist(mfeat[None, :], cfeat[None, :], "cosine"))[0][0]
+                        simdict.update({f"{int(frameId)}_{int(bid)}":simi})
+                        
+                    if cfeat is not None and len(features)>=2:
+                        mfeat = features[-2]
+                        simi = 1 - np.maximum(0.0, cdist(mfeat[None, :], cfeat[None, :], "cosine"))[0][0]
+                        simdict1.update({f"{int(frameId)}_{int(bid)}":simi})
+                    
 
                 if len(tracks) > 0:
                     tracks[:, 7] = frameId
@@ -239,7 +252,10 @@ def yolo_resnet_tracker(
                                  "bboxes": det,
                                  "tboxes": tracks,
                                  "imgs": imgdict,
-                                 "feats": featdict} 
+                                 "feats": featdict,
+                                 "featsimi": simdict,    # 当前 box 特征和该轨迹 smooth_feat 特征的相似度
+                                 "featsimi1": simdict1  # 当前 box 特征和该轨迹前一个 box 特征的相似度
+                                 } 
                     yoloResnetTracker.append(frameDict)
     
                     # imgs, features = inference_image(im0, tracks)
@@ -248,7 +264,14 @@ def yolo_resnet_tracker(
                     '''================== 2. 提取手势位置 ==================='''
                     for *xyxy, id, conf, cls, fid, bid in reversed(tracks):
                         name = ('' if id==-1 else f'id:{int(id)} ') + names[int(cls)]  
-                        label = None if hide_labels else (name if hide_conf else f'{name} {conf:.2f}')
+                        if f"{int(frameId)}_{int(bid)}" in simdict.keys():
+                            sim = simdict[f"{int(frameId)}_{int(bid)}"]
+                            label = f"{name} {sim:.2f}"
+                        else:
+                            label = None if hide_labels else name
+                            
+                            
+                        # label = None if hide_labels else (name if hide_conf else f'{name} {conf:.1f}')
                         
                         if id >=0 and cls==0:
                             color = colors(int(cls), True)
@@ -489,7 +512,7 @@ def run(
                 '''
 
                 det_tracking = Boxes(det, im0.shape).cpu().numpy()
-                tracks = tracker.update(det_tracking, im0)
+                tracks, outfeats = tracker.update(det_tracking, im0)
                 if len(tracks) == 0:
                     continue
                 

tracking/dotrack/dotracks.py

@@ -22,6 +22,37 @@ class MoveState:
     FreeMove = 3
     Unknown = -1
     
+def bbox_ioa(box1, box2, iou=False, eps=1e-7):
+    """
+    Calculate the intersection over box2 area given box1 and box2. Boxes are in x1y1x2y2 format.
+
+    Args:
+        box1 (np.array): A numpy array of shape (n, 4) representing n bounding boxes.
+        box2 (np.array): A numpy array of shape (m, 4) representing m bounding boxes.
+        iou (bool): Calculate the standard iou if True else return inter_area/box2_area.
+        eps (float, optional): A small value to avoid division by zero. Defaults to 1e-7.
+
+    Returns:
+        (np.array): A numpy array of shape (n, m) representing the intersection over box2 area.
+    """
+
+    # Get the coordinates of bounding boxes
+    b1_x1, b1_y1, b1_x2, b1_y2 = box1.T
+    b2_x1, b2_y1, b2_x2, b2_y2 = box2.T
+
+    # Intersection area
+    inter_area = (np.minimum(b1_x2[:, None], b2_x2) - np.maximum(b1_x1[:, None], b2_x1)).clip(0) * \
+                 (np.minimum(b1_y2[:, None], b2_y2) - np.maximum(b1_y1[:, None], b2_y1)).clip(0)
+
+    # box2 area
+    area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
+    if iou:
+        box1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
+        area = area + box1_area[:, None] - inter_area
+
+    # Intersection over box2 area
+    return inter_area / (area + eps)
+    
 class  ShoppingCart:
     
     def __init__(self, bboxes):
@@ -90,6 +121,7 @@ class Track:
         
         self.boxes = boxes
         self.features = features
+        self.slt_boxes = self.select_boxes()
         
         self.tid = int(boxes[0, 4])
         self.cls = int(boxes[0, 6])
@@ -139,6 +171,43 @@ class Track:
         if self.cls == 0:
             self.extract_hand_features()
             
+    def select_boxes(self):
+        
+        slt_boxes = []
+        idx = np.argsort(self.boxes[:, 7])
+        boxes = self.boxes[idx] 
+        features = self.features[idx]
+        
+        for i in range(len(boxes)):
+            simi = None
+            box, tid, fid, bid = boxes[i, :4], int(boxes[i, 4]), int(boxes[i, 7]), int(boxes[i, 8])
+            
+            if i == 0:
+                slt_boxes.append(boxes[i, :])
+                continue
+            
+            if len(boxes)!=len(features):
+                print("check!")
+                continue
+
+            box0, tid0, fid0, bid0 =  boxes[i-1, :4], int(boxes[i-1, 4]), int(boxes[i-1, 7]), int(boxes[i-1, 8])
+            
+            # 当前 box 和轨迹上一个 box 的iou
+            iou = bbox_ioa(box[None, :], box0[None, :])
+
+            # 当前 box 和轨迹上一个 box 的 feat similarity
+            feat0 = features[i, :][None, :]
+            feat1 = features[i-1, :][None, :]
+            simi = 1 - np.maximum(0.0, cdist(feat0, feat1, "cosine"))[0][0]
+            
+            if iou > 0.85 and simi>0.85:
+                continue 
+            
+            slt_boxes.append(boxes[i, :])
+
+        
+        return np.array(slt_boxes)
+        
         
     def compute_cornpoints(self):
         '''
@@ -417,6 +486,8 @@ class doTracks:
         self.FreeMove = []    # subset of self.Residual
         
 
+
+    
     def array2list(self):
         '''
             将 bboxes 变换为 track 列表

tracking/trackers/bot_sort.py

@@ -172,7 +172,7 @@ class BOTSORT(BYTETracker):
             '''1. reid 相似度阈值，低于该值的两 boxes 图像不可能是同一对象，需要确定一个合理的可信阈值
                2. iou 的约束为若约束，故 iou_dists 应设置为较大的值
             '''
-            emb_dists_mask = (emb_dists > 0.9)
+            emb_dists_mask = (emb_dists > 0.8)
             iou_dists[emb_dists_mask] = 1
             emb_dists[iou_dists_mask] = 1
   

tracking/trackers/byte_tracker.py

@@ -375,7 +375,12 @@ class BYTETracker:
         
         output = np.asarray(output1 + output2, dtype=np.float32)
         
-        return output
+        
+        out_feat1 = [(x.frame_id, x.idx, x.smooth_feat, x.curr_feat,  x.features) for x in self.tracked_stracks if x.is_activated]
+        out_feat2 = [(x.frame_id, x.idx, x.smooth_feat, x.curr_feat, x.features) for x in first_finded if x.first_find]
+        
+
+        return output, out_feat1 + out_feat2
     
     
     def get_result(self):