change query id show syntax

add expandable ReID framework

.gitignore

@@ -1,6 +1,10 @@
 # tmp files
 *__pycache__*
 *.pt
+*.jpg
 *.mp4
 *.t7
 *.pth
+*label*
+result.txt
+rank10_detection*.txt

deep_sort/deep/feature_extractor.py

@@ -1,13 +1,11 @@
 import torch
+from torch import nn
 import torchvision.transforms as transforms
 import numpy as np
 import cv2
 import logging
 
-from .model import Net
+from .model import Net,MyNet,RestNet18
-# from fastreid.config import get_cfg
-# from fastreid.engine import DefaultTrainer
-# from fastreid.utils.checkpoint import Checkpointer
 
 class Extractor(object):
     def __init__(self, model_path, use_cuda=True):
@@ -49,36 +47,40 @@ class Extractor(object):
             features = self.net(im_batch)
         return features.cpu().numpy()
 
-class FastReIDExtractor(object):
-    def __init__(self, model_config, model_path, use_cuda=True):
-        raise Exception("fastreid is unloaded")
-        cfg = get_cfg()
-        cfg.merge_from_file(model_config)
-        cfg.MODEL.BACKBONE.PRETRAIN = False
-        self.net = DefaultTrainer.build_model(cfg)
-        self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
     
-        Checkpointer(self.net).load(model_path)
+class MyExtractor(object):
+    def __init__(self, model_path, use_cuda=True):
+        self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
+        self.net = RestNet18()
+        self.net.load_state_dict(torch.load(model_path,map_location=torch.device(self.device)))
+        if self.device=="cuda":
+            self.net = self.net.to(self.device)
         logger = logging.getLogger("root.tracker")
         logger.info("Loading weights from {}... Done!".format(model_path))
-        self.net.to(self.device)
+        self.raw_transformer = transforms.Compose([
-        self.net.eval()
+        transforms.ToPILImage(),
-        height, width = cfg.INPUT.SIZE_TEST
+        transforms.Resize((128,64)),# hxw
-        self.size = (width, height)
+        transforms.ToTensor(),
+        ])   
         self.norm = transforms.Compose([
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
         
 
+
     def _preprocess(self, im_crops):
-        def _resize(im, size):
+        """
-            return cv2.resize(im.astype(np.float32)/255., size)
+        TODO:
-
+            1. to float with scale from 0 to 1
-        im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float()
+            2. resize to (64, 128) as Market1501 dataset did
+            3. concatenate to a numpy array
+            3. to torch Tensor
+            4. normalize
+        """
+        im_batch = torch.cat([self.raw_transformer(im).unsqueeze(0) for im in im_crops], dim=0)
         return im_batch
 
-
     def __call__(self, im_crops):
         im_batch = self._preprocess(im_crops)
         with torch.no_grad():
@@ -87,10 +89,3 @@ class FastReIDExtractor(object):
         return features.cpu().numpy()
 
 
-
-if __name__ == '__main__':
-    img = cv2.imread("demo.jpg")[:,:,(2,1,0)]
-    extr = Extractor("checkpoint/ckpt.t7")
-    feature = extr(img)
-    print(feature.shape)
-

deep_sort/deep/model.py

@@ -94,11 +94,135 @@ class Net(nn.Module):
         x = self.classifier(x)
         return x
     
+# Define model
+# input:128x256x3
+# output: 64
+class MyNet(nn.Module):
+    def __init__(self):
+        super(MyNet, self).__init__()
+        self.conv1 = nn.Conv2d(3,8,5,1,2)
+        self.maxpool1 = nn.MaxPool2d(2) # 64x128x8
+        self.bn1   = nn.BatchNorm2d(8)
+        self.conv2 = nn.Conv2d(8,16,5,1,2)
+        self.maxpool2 = nn.MaxPool2d(2) # 32x64x16
+        self.bn2   = nn.BatchNorm2d(16)
+        self.conv3 = nn.Conv2d(16,32,5,1,2)
+        self.conv4 = nn.Conv2d(32,64,5,1,2)
+        self.maxpool3 = nn.MaxPool2d(2) # 16x32x64
+        self.bn3   = nn.BatchNorm2d(64)
+        self.conv5 = nn.Conv2d(64,32,5,1,2) # 16x32x32
+        self.maxpool4 = nn.MaxPool2d(2) # 8x16x32
+        self.bn4   = nn.BatchNorm2d(32)
+        self.conv6 = nn.Conv2d(32,64,8,8) #1x2x64
+        self.flat = nn.Flatten()
+        self.linear = nn.Linear(2*64,64) # 64D-feature ID
+        
+        self.sigmoid = nn.Sigmoid()
+        self.lrelu = nn.LeakyReLU()
+
+    def forward(self, x:torch.Tensor):   
+        x = self.conv1(x)
+        x = self.lrelu(x)
+        x = self.maxpool1(x)
+        x = self.bn1(x)
+        x = self.conv2(x)
+        x = self.lrelu(x)
+        x = self.maxpool2(x)
+        x = self.bn2(x)
+        x = self.conv3(x)
+        x = self.lrelu(x)
+        x = self.conv4(x)
+        x = self.lrelu(x)
+        x = self.maxpool3(x)
+        x = self.bn3(x)
+        x = self.conv5(x)
+        x = self.lrelu(x)
+        x = self.maxpool4(x)
+        x = self.bn4(x)
+        x = self.conv6(x)
+        x = self.sigmoid(x)
+        x = self.flat(x)
+        x = self.linear(x)
+        
+        return x
 
 if __name__ == '__main__':
     net = Net()
     x = torch.randn(4,3,128,64)
     y = net(x)
-    import ipdb; ipdb.set_trace()
+    
+#########ResNet18
+class RestNetBasicBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride):
+        super(RestNetBasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        output = self.conv1(x)
+        output = F.relu(self.bn1(output))
+        output = self.conv2(output)
+        output = self.bn2(output)
+        return F.leaky_relu(x + output)
+
+
+class RestNetDownBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride):
+        super(RestNetDownBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.extra = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0),
+            nn.BatchNorm2d(out_channels)
+        )
+
+    def forward(self, x):
+        extra_x = self.extra(x)
+        output = self.conv1(x)
+        out = F.relu(self.bn1(output))
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        return F.leaky_relu(extra_x + out)
+
+
+class RestNet18(nn.Module):
+    def __init__(self):
+        super(RestNet18, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=(3, 3), padding=1)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.maxpool = nn.MaxPool2d(2)
+
+        self.layer1 = nn.Sequential(RestNetBasicBlock(32, 32, 1),
+                                    RestNetBasicBlock(32, 32, 1))
+
+        self.layer2 = nn.Sequential(RestNetDownBlock(32, 64, [2, 1]),
+                                    RestNetBasicBlock(64, 64, 1))
+
+        self.layer3 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]),
+                                    RestNetBasicBlock(128, 128, 1))
+
+        self.layer4 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]),
+                                    RestNetBasicBlock(256, 256, 1))
+
+        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
+
+        self.fc = nn.Linear(32768, 23)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        # x = self.avgpool(x)
+        x = x.reshape(x.shape[0], -1)
+        # print(x.size())
+        x = self.fc(x)
+        return x
 
 

deep_sort/deep_sort.py

@@ -1,7 +1,7 @@
 import numpy as np
 import torch
 
-from .deep.feature_extractor import Extractor, FastReIDExtractor
+from .deep.feature_extractor import Extractor,MyExtractor
 from .sort.nn_matching import NearestNeighborDistanceMetric
 from .sort.preprocessing import non_max_suppression
 from .sort.detection import Detection
@@ -12,14 +12,14 @@ __all__ = ['DeepSort']
 
 
 class DeepSort(object):
-    def __init__(self, model_path, model_config=None, max_dist=0.3, min_confidence=0.35, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=30, n_init=3, nn_budget=100, use_cuda=True):
+    def __init__(self, model_path, model_config=None, max_dist=1e-2, min_confidence=0.35, nms_max_overlap=0.9, max_iou_distance=0.40, max_age=10, n_init=3, nn_budget=5000, use_cuda=True):
         self.min_confidence = min_confidence
         self.nms_max_overlap = nms_max_overlap
 
         if model_config is None:
             self.extractor = Extractor(model_path, use_cuda=use_cuda)
         else:
-            self.extractor = FastReIDExtractor(model_config, model_path, use_cuda=use_cuda)
+            self.extractor = MyExtractor(model_path, use_cuda=use_cuda)
 
         max_cosine_distance = max_dist
         metric = NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
@@ -30,6 +30,7 @@ class DeepSort(object):
         # generate detections
         features = self._get_features(bbox_xywh, ori_img)
         bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)
+
         detections = [Detection(bbox_tlwh[i], conf, features[i]) for i,conf in enumerate(confidences) if conf>self.min_confidence]
 
         # run on non-maximum supression
@@ -44,16 +45,18 @@ class DeepSort(object):
         
         # output bbox identities
         outputs = []
+        output_features = []
         for track in self.tracker.tracks:
             if not track.is_confirmed() or track.time_since_update > 1:
                 continue
             box = track.to_tlwh()
             x1,y1,x2,y2 = self._tlwh_to_xyxy(box)
             track_id = track.track_id
+            output_features.append(track.features) 
             outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=np.int))
         if len(outputs) > 0:
             outputs = np.stack(outputs,axis=0)
-        return outputs
+        return outputs,output_features
 
 
     """
@@ -63,7 +66,6 @@ class DeepSort(object):
     """
     @staticmethod
     def _xywh_to_tlwh(bbox_xywh):
-        
         if isinstance(bbox_xywh, np.ndarray):
             bbox_tlwh = bbox_xywh.copy()
         elif isinstance(bbox_xywh, torch.Tensor):
@@ -98,8 +100,8 @@ class DeepSort(object):
     def _xyxy_to_tlwh(self, bbox_xyxy):
         x1,y1,x2,y2 = bbox_xyxy
 
-        t = int(x1)
+        t = x1
-        l = int(y1)
+        l = y1
         w = int(x2-x1)
         h = int(y2-y1)
         return t,l,w,h

deep_sort/sort/iou_matching.py

@@ -23,6 +23,8 @@ def iou(bbox, candidates):
         occluded by the candidate.
 
     """
+    length = len(candidates)
+
     bbox_tl, bbox_br = bbox[:2], bbox[:2] + bbox[2:]
     candidates_tl = candidates[:, :2]
     candidates_br = candidates[:, :2] + candidates[:, 2:]
@@ -36,7 +38,14 @@ def iou(bbox, candidates):
     area_intersection = wh.prod(axis=1)
     area_bbox = bbox[2:].prod()
     area_candidates = candidates[:, 2:].prod(axis=1)
-    return area_intersection / (area_bbox + area_candidates - area_intersection)
+    
+    # should be consious
+    gious = []
+    for i in range(length):
+        gious.append(float((area_bbox + area_candidates[i] - area_intersection[i])/(max(bbox[0]+bbox[2],candidates[i][0]+candidates[i][2])-min(bbox[0],candidates[i][0]))/(max(bbox[1]+bbox[3],candidates[i][1]+candidates[i][3])-min(bbox[1],candidates[i][1]))) )
+        
+    # return area_intersection / (area_bbox + area_candidates - area_intersection)
+    return np.array(gious)
 
 
 def iou_cost(tracks, detections, track_indices=None,

deep_sort/sort/tracker.py

@@ -86,7 +86,8 @@ class Tracker:
                 continue
             features += track.features
             targets += [track.track_id for _ in track.features]
-            track.features = []
+            # track.features = []
+            
         self.metric.partial_fit(
             np.asarray(features), np.asarray(targets), active_targets)
 
@@ -99,7 +100,6 @@ class Tracker:
             cost_matrix = linear_assignment.gate_cost_matrix(
                 self.kf, cost_matrix, tracks, dets, track_indices,
                 detection_indices)
-
             return cost_matrix
 
         # Split track set into confirmed and unconfirmed tracks.

detect.py

@@ -1,10 +1,7 @@
-import cv2
 import torch
-import numpy as np
 
-from torchvision import transforms
 from models.experimental import attempt_load
-from utils.general import check_img_size,non_max_suppression
+from utils.general import non_max_suppression
 
 import param
 
@@ -13,7 +10,7 @@ class YOLOv7(object):
     def __init__(self,option:param.Parameters) -> None:
         # load FP32 model
         self.opt = option
-        self.model = attempt_load(weights=self.opt.weights, device=self.opt.device)
+        self.model = attempt_load(weights=self.opt.weights_yolo, device=self.opt.device)
 
     def detect(self,imgs:torch.Tensor) -> list:
         # img.size should be (Batch,3,H,W),every value should be range of [0,1]

example_deepsort.py

@@ -1,84 +0,0 @@
-import torch
-from torchvision import transforms
-import numpy as np
-import cv2
-
-import param
-from detect import YOLOv7
-from deep_sort import deep_sort as dsort
-
-def get_img_test(raw_img:np.array):
-    raw_transform = transforms.Compose([transforms.ToPILImage(),
-                                        transforms.Resize((360,640)),
-                                        transforms.Pad((0,(640-360)//2)),])
-    return raw_transform(raw_img)
-
-def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
-    pred_img = pred_img.astype(np.uint)
-    print(type(img))
-    for info in pred_img:
-        x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
-        cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
-    print(type(img))
-    cv2.imshow("test",img)
-    cv2.waitKey(0)
-    cv2.imwrite("result.png",img)
-
-routine = {}
-    
-def opencv_sort_plot(img:cv2.Mat,pred_yolo:np.array,pred_sort:np.array):
-    for info in pred_yolo:
-        x1,y1,x2,y2 = int(info[0]),int(info[1]),int(info[2]),int(info[3])
-        cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),1)    
-    for info in pred_sort:
-        x1,y1,x2,y2,id = info[0],info[1],info[2],info[3],info[4]
-        cxy = (int((x1+x2)/2),int((y1+y2)/2))
-        color = (int(id%3*100),int(id%4*75),int(id%5*50))
-        cv2.putText(img,str(id),cxy, cv2.FONT_HERSHEY_PLAIN, 1.0, color, 2)
-        if routine.get(id) is None:
-            routine[id] = [cxy]
-        else:
-            for i in range(1,len(routine[id])):
-                cv2.line(img,routine[id][i-1],routine[id][i],color,1)
-            cv2.line(img,routine[id][-1],cxy,color,1)
-            routine[id].append(cxy)
-    cv2.imshow("test",img)
-    cv2.waitKey(1)
-    return img
-
-
-
-# example        
-if __name__ == '__main__':
-    dsort_path = "./deep_sort/deep/checkpoint/market_agw_R50.pth"
-    my_option = param.Parameters()
-    model = YOLOv7(my_option)                       # load model in head
-    deepsort = dsort.DeepSort(model_path= dsort_path,model_config=None,use_cuda=(torch.device("cuda:0") == my_option.device))
-    with torch.no_grad():
-        sources = cv2.VideoCapture("test.mp4")
-        target = cv2.VideoWriter("output.mp4",cv2.VideoWriter_fourcc('M', 'P', '4', '2'),24,(640,640))
-        while True:
-            ret,frame = sources.read()
-            if ret is False:
-                break
-            img    = get_img_test(frame)                # get image as torch.Tensor with size of [1,1,640,640]
-            img_tensor = transforms.ToTensor()(img).unsqueeze(dim=0).to(my_option.device)     
-            result = model.detect(img_tensor)                  # get the sequence of result   
-            result = result[0].detach().cpu().numpy()   # the single img is index 0
-            bbox_xywhs = []
-            confs      = []
-            for xyxycc in result: 
-                xywh   =  deepsort._xyxy_to_xywh(xyxycc[0:4])
-                conf   =  xyxycc[4]
-                clas   =  int(xyxycc[5])
-                bbox_xywhs.append(xywh[:])
-                confs.append(conf)
-            dpsort = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(img))
-            print(f"result:{result}")
-            print(f"dpsort:{dpsort}")
-            frame_t = opencv_sort_plot(np.array(img),result,dpsort)
-            target.write(frame_t)
-        target.release()
-            # img_src = cv2.imread("test.png")            # re-read for imshow
-            # img_src = cv2.resize(img_src,(640,360))
-            # opencv_box_plot(img_src,result)             # show the result in picture

example_reid.py

@@ -0,0 +1,185 @@
+import torch
+from torchvision import transforms
+import numpy as np
+import cv2
+import os
+
+import param
+from detect import YOLOv7
+from deep_sort import deep_sort as dsort
+
+def get_img_test(raw_img:np.array):
+    raw_transform = transforms.Compose([transforms.ToPILImage(),
+                                        transforms.Resize((360,640)),
+                                        transforms.Pad((0,(640-360)//2)),])
+    return raw_transform(raw_img)
+
+def reverse_box_get_from_yolo(detection):
+    x1 = int(detection[0]*1920/640)
+    x2 = int(detection[2]*1920/640)
+    y1 = int((detection[1]-140)*1080/360)
+    y2 = int((detection[3]-140)*1080/360)
+    return x1,y1,x2,y2
+
+def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
+    pred_img = pred_img.astype(np.uint)
+    for info in pred_img:
+        x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
+        cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
+
+    cv2.imshow("test",img)
+    cv2.waitKey(0)
+    cv2.imwrite("result.png",img)
+    
+def opencv_match_pointer_plot(img:cv2.Mat,start_base:tuple,text:str,color:tuple):
+    SIZE = 50
+    brush = [0,0]
+    brush[0]=start_base[0]
+    brush[1]=start_base[1]
+    brush_push = [brush[0]+SIZE,brush[1]-SIZE]
+    cv2.line(img,brush,brush_push,color,1)
+    # brush=brush_push
+    # brush_push[0]+=SIZE*2
+    # brush_push[1]+=0
+    # cv2.line(img,brush,brush_push,color,1)    
+    cv2.putText(img,text,brush_push, cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
+    return img
+
+routine = {}
+    
+def opencv_sort_plot(img:cv2.Mat,pred_yolo:np.array,pred_sort:np.array):
+    for info in pred_yolo:
+        x1,y1,x2,y2 = int(info[0]),int(info[1]),int(info[2]),int(info[3])
+        cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),1)    
+    for info in pred_sort:
+        x1,y1,x2,y2,id = info[0],info[1],info[2],info[3],info[4]
+        cxy = (int((x1+x2)/2),int((y1+y2)/2))
+        color = (int(id%3*100),int(id%4*75),int(id%5*50))
+        cv2.putText(img,str(id),cxy, cv2.FONT_HERSHEY_PLAIN, 1.0, color, 2)
+        if routine.get(id) is None:
+            routine[id] = [cxy]
+        else:
+            for i in range(1,len(routine[id])):
+                cv2.line(img,routine[id][i-1],routine[id][i],color,1)
+            cv2.line(img,routine[id][-1],cxy,color,1)
+            routine[id].append(cxy)
+    return img
+
+def extractidfeature(id_path:str,extractor):
+    img_id = cv2.imread(id_path)
+    tensor_id = extractor([img_id])
+    return tensor_id
+
+# example        
+if __name__ == '__main__':
+    MAX_BUFFLEN        = 10
+    DISTANCE_THRESHOLD = 0.4
+    SAVE_TXT_FLAG = True
+    
+    my_option = param.Parameters()
+    
+    print(torch.__version__)
+    print(my_option.device)
+    
+    query_path = f"./mydataset/query/cam{my_option.query_index}/"                   # query from .jpg photos
+    gallary_path = f"mydataset/video/cam{my_option.gallary_index}.mp4"              # source from mp4 via yolo
+    output_video_path = f"output/output{my_option.gallary_index}_{my_option.query_index}.mp4"    # output file path
+    output_txt_path = f"output/rank10_detection_{my_option.gallary_index}_{my_option.query_index}.txt"
+    
+    query_features = []
+    paths = os.listdir(query_path)
+    query_match_buff = [[] for p in paths]
+    query_matched = [-1 for p in paths]
+    
+    # load detection model and deepsort model
+    model = YOLOv7(my_option)         
+    deepsort = dsort.DeepSort(model_path=my_option.weights_reid,model_config="self",use_cuda=(torch.device("cuda:0") == my_option.device))
+    
+    # get query feature id
+    for p in paths:
+        query_features.append(extractidfeature(query_path+p,deepsort.extractor))
+        
+    with torch.no_grad():
+        sources = cv2.VideoCapture(gallary_path)
+        target = cv2.VideoWriter(output_video_path,cv2.VideoWriter_fourcc('m', 'p', '4', 'v'),24,(1920,1080))
+        frame_counter = 0
+        while True:
+            ret,frame = sources.read()
+            if ret is False:
+                break
+            
+            img    = get_img_test(frame)                        # get image as torch.Tensor with size of [1,1,640,640]
+            img_tensor = transforms.ToTensor()(img).unsqueeze(dim=0).to(my_option.device)     
+            detections = model.detect(img_tensor)               # get the sequence of result   
+            detections = detections[0].detach().cpu().numpy()   # the single img is index 0
+            
+            for detection in detections:
+                detection[0],detection[1],detection[2],detection[3] = reverse_box_get_from_yolo(detection[0:4])
+            
+            bbox_xywhs = []
+            confs      = []
+            for xyxycc in detections: 
+                xywh   =  deepsort._xyxy_to_xywh(xyxycc[0:4])
+                conf   =  xyxycc[4]
+                clas   =  int(xyxycc[5])
+                bbox_xywhs.append(xywh[:])
+                confs.append(conf)
+            dpsort, features = deepsort.update(bbox_xywh=np.array(bbox_xywhs),confidences=np.array(confs),ori_img=np.array(frame))
+            
+            img_drawing = np.array(frame)
+            img_drawing = opencv_sort_plot(img_drawing,detections,dpsort)
+            
+            # ReID: compute features
+            for i in range(len(dpsort)):
+                track = dpsort[i]
+                MOT_id = track[-1]
+                feature_distances = []
+                for j in range(len(query_features)):
+                    feature_distance = np.mean(np.power(features[i]-query_features[j],2))
+                    feature_distances.append(feature_distance)
+                    
+                # ReID: rank distances for every feature   
+                feature_distances_sorted = sorted(feature_distances)
+                for distance in feature_distances_sorted:
+                    if distance < DISTANCE_THRESHOLD:
+                        index_d = feature_distances.index(distance)
+                        query_match_buff[index_d].append((MOT_id,distance,(frame_counter,track[0],track[1],track[2],track[3]))) 
+                        query_match_buff[index_d].sort(key=lambda x:x[1],reverse=True)
+                        if len(query_match_buff[index_d]) > MAX_BUFFLEN:
+                            query_match_buff[index_d].pop()
+                            count = 0
+                            for term in query_match_buff[index_d]:
+                                if term[0] == MOT_id:
+                                    count+=1
+                            if count > MAX_BUFFLEN//3:
+                                cx = int((track[0]+track[2])/2)
+                                cy = int((track[1]+track[3])/2)
+                                color = (int(index_d%2*100),int(index_d%3*75),int(index_d%4*50))
+                                
+                                # img_drawing=opencv_match_pointer_plot(img_drawing,(cx,track[1]-10),f"{paths[index_d]}",color)
+                                # cv2.rectangle(img_drawing,(track[i][0],track[i][1]),(track[i][2],track[i][3]),color,1)                
+                                cv2.putText(img_drawing,f"q_{paths[index_d]}".split(".")[0],(track[0],track[1]), cv2.FONT_HERSHEY_PLAIN, 1, color, 2)
+                            break
+                    else:
+                        break
+            
+            # Final plot and show video
+            cv2.imshow("test",img_drawing)
+            cv2.waitKey(1)
+            target.write(img_drawing)
+            frame_counter += 1
+        
+        # save the video and end the program
+        target.release()
+        
+        # rank10 output as txt
+        if SAVE_TXT_FLAG is True:
+            with open(output_txt_path,mode="w+") as f:
+                result = []
+                for i in range(len(query_match_buff)):
+                    for raw_info in query_match_buff[i]:
+                        detect_info = raw_info[-1]
+                        result.append((detect_info[0],i,detect_info[1],detect_info[2],detect_info[3],detect_info[4]))
+                        result.sort(key=lambda x:x[0])
+                for info in result:
+                    f.write(str(info[0])+","+str(info[1])+","+str(info[2])+","+str(info[3])+","+str(info[4])+","+str(info[5])+"\n")

example_yolov7.py

@@ -1,40 +0,0 @@
-import torch
-from torchvision import transforms
-import numpy as np
-import cv2
-
-import param
-from detect import YOLOv7
-
-def get_img_test(path_src:str):
-    raw_img = cv2.imread(path_src)
-    raw_transform = transforms.Compose([transforms.ToPILImage(),
-                                        transforms.Resize((360,640)),
-                                        transforms.Pad((0,(640-360)//2)),
-                                        transforms.ToTensor()])
-    return raw_transform(raw_img).unsqueeze(dim=0)
-
-def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
-    pred_img = pred_img.astype(np.uint)
-    print(type(img))
-    for info in pred_img:
-        x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
-        cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
-    print(type(img))
-    cv2.imshow("test",img)
-    cv2.waitKey(0)
-    cv2.imwrite("result.png",img)
-
-# example        
-if __name__ == '__main__':
-    my_option = param.Parameters()
-    model = YOLOv7(my_option)                       # load model in head
-    with torch.no_grad():
-        img = get_img_test("test.png")              # get image as torch.Tensor with size of [1,1,640,640]
-        print(f"img size is {img.size()}")        
-        result = model.detect(img)                  # get the sequence of result
-        print(f"result is :{result}")
-        result = result[0].detach().cpu().numpy()   # the single img is index 0
-        img_src = cv2.imread("test.png")            # re-read for imshow
-        img_src = cv2.resize(img_src,(640,360))
-        opencv_box_plot(img_src,result)             # show the result in picture

param.py

@@ -2,13 +2,16 @@ import torch
 
 class Parameters(object):
     def __init__(self) -> None:
-        self.weights  = "./best.pt"
+        self.weights_yolo = "./best.pt"
-        self.conf_thres = 0.35
+        self.weights_reid = "./resnet18.pth"
-        self.iou_thres  = 0.45
+        self.conf_thres   = 0.40
-        self.device     = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.iou_thres    = 0.60
-        self.nosave     = False
+        self.device       = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.classes    = [0]
+        self.nosave       = False
+        self.classes      = [0]
         self.agnostic_nms = None
         self.augment      = None
+        self.query_index  = 1
+        self.gallary_index= 2
         pass
     

requirements.txt

@@ -0,0 +1,40 @@
+# Usage: pip install -r requirements.txt
+
+# Base ----------------------------------------
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python>=4.1.1
+Pillow>=7.1.2
+PyYAML>=5.3.1
+requests>=2.23.0
+scipy>=1.4.1
+torch>=1.7.0,!=1.12.0
+torchvision>=0.8.1,!=0.13.0
+tqdm>=4.41.0
+protobuf<4.21.3
+
+# Logging -------------------------------------
+tensorboard>=2.4.1
+# wandb
+
+# Plotting ------------------------------------
+pandas>=1.1.4
+seaborn>=0.11.0
+
+# Export --------------------------------------
+# coremltools>=4.1  # CoreML export
+# onnx>=1.9.0  # ONNX export
+# onnx-simplifier>=0.3.6  # ONNX simplifier
+# scikit-learn==0.19.2  # CoreML quantization
+# tensorflow>=2.4.1  # TFLite export
+# tensorflowjs>=3.9.0  # TF.js export
+# openvino-dev  # OpenVINO export
+
+# Extras --------------------------------------
+ipython  # interactive notebook
+psutil  # system utilization
+thop  # FLOPs computation
+# albumentations>=1.0.3
+# pycocotools>=2.0  # COCO mAP
+# roboflow
+yacs

utils/torch_utils.py

@@ -129,6 +129,7 @@ def profile(x, ops, n=100, device=None):
         s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
         s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
         p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0  # parameters
+        raise Exception("thop package is discarded")
         print(f'{p:12}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
 
 
@@ -213,6 +214,7 @@ def model_info(model, verbose=False, img_size=640):
                   (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
 
     try:  # FLOPS
+        raise Exception("package thop is discarded")
         from thop import profile
         stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32
         img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device)  # input