dnt 0.2.1__py3-none-any.whl → 0.3.1.8__py3-none-any.whl

dnt/track/{dsort/utils → botsort/bot_tracker/tracking_utils}/evaluation.py

@@ -3,7 +3,8 @@ import numpy as np
 import copy
 import motmetrics as mm
 mm.lap.default_solver = 'lap'
-from utils.io import read_results, unzip_objs
+
+from yolox.tracking_utils.io import read_results, unzip_objs
 
 
 class Evaluator(object):
@@ -39,7 +40,6 @@ class Evaluator(object):
         ignore_objs = self.gt_ignore_frame_dict.get(frame_id, [])
         ignore_tlwhs = unzip_objs(ignore_objs)[0]
 
-
         # remove ignored results
         keep = np.ones(len(trk_tlwhs), dtype=bool)
         iou_distance = mm.distances.iou_matrix(ignore_tlwhs, trk_tlwhs, max_iou=0.5)
@@ -53,6 +53,15 @@ class Evaluator(object):
             keep[match_js] = False
             trk_tlwhs = trk_tlwhs[keep]
             trk_ids = trk_ids[keep]
+        #match_is, match_js = mm.lap.linear_sum_assignment(iou_distance)
+        #match_is, match_js = map(lambda a: np.asarray(a, dtype=int), [match_is, match_js])
+        #match_ious = iou_distance[match_is, match_js]
+
+        #match_js = np.asarray(match_js, dtype=int)
+        #match_js = match_js[np.logical_not(np.isnan(match_ious))]
+        #keep[match_js] = False
+        #trk_tlwhs = trk_tlwhs[keep]
+        #trk_ids = trk_ids[keep]
 
         # get distance matrix
         iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5)
@@ -70,7 +79,8 @@ class Evaluator(object):
         self.reset_accumulator()
 
         result_frame_dict = read_results(filename, self.data_type, is_gt=False)
-        frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys())))
+        #frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys())))
+        frames = sorted(list(set(result_frame_dict.keys())))
         for frame_id in frames:
             trk_objs = result_frame_dict.get(frame_id, [])
             trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2]
@@ -100,4 +110,4 @@ class Evaluator(object):
         import pandas as pd
         writer = pd.ExcelWriter(filename)
         summary.to_excel(writer)
-        writer.save()
+        writer.save()

dnt/track/{dsort/utils → botsort/bot_tracker/tracking_utils}/io.py

@@ -2,58 +2,34 @@ import os
 from typing import Dict
 import numpy as np
 
-# from utils.log import get_logger
 
+def write_results(filename, results_dict: Dict, data_type: str):
+    if not filename:
+        return
+    path = os.path.dirname(filename)
+    if not os.path.exists(path):
+        os.makedirs(path)
 
-def write_results(filename, results, data_type):
-    if data_type == 'mot':
-        save_format = '{frame},{id},{x1},{y1},{w},{h},-1,-1,-1,-1\n'
+    if data_type in ('mot', 'mcmot', 'lab'):
+        save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n'
     elif data_type == 'kitti':
-        save_format = '{frame} {id} pedestrian 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n'
+        save_format = '{frame} {id} pedestrian -1 -1 -10 {x1} {y1} {x2} {y2} -1 -1 -1 -1000 -1000 -1000 -10 {score}\n'
     else:
         raise ValueError(data_type)
 
     with open(filename, 'w') as f:
-        for frame_id, tlwhs, track_ids in results:
+        for frame_id, frame_data in results_dict.items():
             if data_type == 'kitti':
                 frame_id -= 1
-            for tlwh, track_id in zip(tlwhs, track_ids):
+            for tlwh, track_id in frame_data:
                 if track_id < 0:
                     continue
                 x1, y1, w, h = tlwh
                 x2, y2 = x1 + w, y1 + h
-                line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h)
+                line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h, score=1.0)
                 f.write(line)
 
 
-# def write_results(filename, results_dict: Dict, data_type: str):
-#     if not filename:
-#         return
-#     path = os.path.dirname(filename)
-#     if not os.path.exists(path):
-#         os.makedirs(path)
-
-#     if data_type in ('mot', 'mcmot', 'lab'):
-#         save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n'
-#     elif data_type == 'kitti':
-#         save_format = '{frame} {id} pedestrian -1 -1 -10 {x1} {y1} {x2} {y2} -1 -1 -1 -1000 -1000 -1000 -10 {score}\n'
-#     else:
-#         raise ValueError(data_type)
-
-#     with open(filename, 'w') as f:
-#         for frame_id, frame_data in results_dict.items():
-#             if data_type == 'kitti':
-#                 frame_id -= 1
-#             for tlwh, track_id in frame_data:
-#                 if track_id < 0:
-#                     continue
-#                 x1, y1, w, h = tlwh
-#                 x2, y2 = x1 + w, y1 + h
-#                 line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h, score=1.0)
-#                 f.write(line)
-#     logger.info('Save results to {}'.format(filename))
-
-
 def read_results(filename, data_type: str, is_gt=False, is_ignore=False):
     if data_type in ('mot', 'lab'):
         read_fun = read_mot_results
@@ -96,6 +72,8 @@ def read_mot_results(filename, is_gt, is_ignore):
                     continue
                 results_dict.setdefault(fid, list())
 
+                box_size = float(linelist[4]) * float(linelist[5])
+
                 if is_gt:
                     if 'MOT16-' in filename or 'MOT17-' in filename:
                         label = int(float(linelist[7]))
@@ -115,6 +93,11 @@ def read_mot_results(filename, is_gt, is_ignore):
                 else:
                     score = float(linelist[6])
 
+                #if box_size > 7000:
+                #if box_size <= 7000 or box_size >= 15000:
+                #if box_size < 15000:
+                    #continue
+
                 tlwh = tuple(map(float, linelist[2:6]))
                 target_id = int(linelist[1])
 

dnt/track/botsort/bot_tracker/tracking_utils/timer.py

@@ -0,0 +1,37 @@
+import time
+
+
+class Timer(object):
+    """A simple timer."""
+    def __init__(self):
+        self.total_time = 0.
+        self.calls = 0
+        self.start_time = 0.
+        self.diff = 0.
+        self.average_time = 0.
+
+        self.duration = 0.
+
+    def tic(self):
+        # using time.time instead of time.clock because time time.clock
+        # does not normalize for multithreading
+        self.start_time = time.time()
+
+    def toc(self, average=True):
+        self.diff = time.time() - self.start_time
+        self.total_time += self.diff
+        self.calls += 1
+        self.average_time = self.total_time / self.calls
+        if average:
+            self.duration = self.average_time
+        else:
+            self.duration = self.diff
+        return self.duration
+
+    def clear(self):
+        self.total_time = 0.
+        self.calls = 0
+        self.start_time = 0.
+        self.diff = 0.
+        self.average_time = 0.
+        self.duration = 0.

dnt/track/botsort/inference.py

@@ -0,0 +1,96 @@
+import os, sys
+import cv2
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import argparse
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
+
+from bot_tracker.mc_bot_sort import BoTSORT
+from config import Config
+
+def track(video_file:str, det_file:str, out_file:str = None, cfg:dict = Config.get_cfg_botsort(), 
+          video_index:int = None, total_videos:int = None)->pd.DataFrame:
+
+    cap = cv2.VideoCapture(video_file)
+    if not cap.isOpened():
+        raise IOError("Couldn't open webcam or video")
+    
+    if video_file is not None:
+        fps = cap.get(cv2.CAP_PROP_FPS)
+    else:
+        fps = cfg['fps']
+
+    botsort = BoTSORT(cfg, frame_rate=fps)
+    
+    detections = pd.read_csv(det_file, header=None).to_numpy()
+    start_frame = int(min(detections[:,0]))
+    end_frame = int(max(detections[:,0]))
+    tot_frames = max((end_frame - start_frame + 1), 0)
+
+    pbar = tqdm(total=tot_frames, desc='Tracking {}'.format(os.path.basename(video_file)), unit = 'frames')
+    if video_index and total_videos:
+        pbar.desc = 'Tracking {} of {}'.format(video_index, total_videos)
+
+    results = []
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+    while cap.isOpened():
+
+        pos_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+        ret, frame = cap.read()
+
+        if (not ret) or (pos_frame>end_frame):
+            break
+        
+        im = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frame_dets = detections[np.where(detections[:,0] == pos_frame)]
+        nrows, ncols = frame_dets.shape
+
+        if nrows > 0:
+            dets = []
+            for det in frame_dets:
+                dets.append([det[2], det[3], det[2]+det[4], det[3]+det[5], det[6], det[7]]) # x1, y1, x2, y2, score, class
+            dets = np.array(dets)
+            '''
+            for det in frame_dets:
+                bbox_xywh.append([det[2]+det[4]/2, det[3]+det[5]/2, det[4], det[5]])
+            bbox_xywh = np.array(bbox_xywh)
+            conf_score = np.array(frame_dets[:,6])
+            '''
+            outputs = botsort.update(dets, im)
+
+            if len(outputs) > 0:
+                for t in outputs:
+                    tlwh = t.tlwh
+                    tlbr = t.tlbr
+                    tid = t.track_id
+                    tcls = t.cls
+                    score = t.score
+                    vertical = tlwh[2] / tlwh[3] > cfg['aspect_ratio_thresh']
+                    if tlwh[2] * tlwh[3] > cfg['min_box_area'] and not vertical:
+                        results.append([pos_frame, tid, int(tlwh[0]), int(tlwh[1]), int(tlwh[2]), int(tlwh[3]), round(score,1), int(tcls), -1, -1])
+
+        pbar.update()
+    cap.release()
+    pbar.close()
+
+    df = pd.DataFrame(results)
+    if out_file:
+        if cfg['output_header']:
+            header = ['frame', 'track_id', 'x', 'y', 'w', 'h', 'score', 'class', 'r3','r4']
+        else:
+            header = None
+        df.to_csv(out_file, index=False, header=header)
+
+    return df
+
+if __name__ == "__main__":
+
+    video_file = "/mnt/d/videos/ped2stage/videos/gh1293.mp4"
+    iou_file = "/mnt/d/videos/ped2stage/dets/gh1293_iou_ped.txt"
+    out_file = "/mnt/d/videos/ped2stage/tracks/gh1293_ped_track_2.txt"
+
+    track(video_file, iou_file, out_file)
+    

dnt/track/config.py

@@ -0,0 +1,120 @@
+import yaml
+import os, io
+
+class Config:
+    def __init__(self, cfg:dict):
+        pass
+
+    @staticmethod
+    def get_cfg_sort()->dict:
+        '''
+        method: str - Tracking method, 'sort'
+        max_age: int - Maximum age of a track, default is 1
+        min_inits: int - Number of frames to wait before activate a track, default is 3
+        iou_threshold: float - IOU threshold, default is 0.3
+        '''
+        return {'method':'sort', 
+                'max_age':1, 
+                'min_inits':3, 
+                'iou_threshold':0.3}
+
+    @staticmethod
+    def get_cfg_dsort(model:str="default")->dict:
+        '''
+        Parameters:
+            model: str - DeepSORT model, 'default', 'bot_R50', 'bot_R50-ibn', 'bot_S50', 'bot_R101-ibn', 'veh_bot_R50-ibn'
+        Contents:
+            method: str - Tracking method, 'dsort'
+            reid_cfg: str - ReID config file path
+            reid_ckpt: str - ReID weights file path
+            max_dist: float - Maximum cosine distance, default is 0.2
+            min_confidence: float - Minimum detection confidence, default is 0.3
+            nms_max_overlap: float - Maximum overlap in NMS, default is 0.5
+            max_iou_distance: float - Maximum IOU distance, default is 0.7
+            max_age: int - Maximum age of a track, default is 70
+            n_init: int - Number of frames to wait before activate a track, default is 3
+            nn_budget: int - Maximum number of nearest neighbors, default is 100
+        '''
+        cfg = { 'method': 'dsort',
+                'reid_cfg': None,
+                'reid_ckpt': 'ckpt.t7', 
+                'max_dist': 0.2, 
+                'min_confidence': 0.3,
+                'nms_max_overlap': 0.5, 
+                'max_iou_distance': 0.7, 
+                'max_age': 70, 
+                'n_init': 3, 
+                'nn_budget': 100}
+
+        if model == 'bot_R50':
+            cfg['reid_cfg'] = 'Market1501/bagtricks_R50.yml'
+            cfg['reid_ckpt'] = 'market_bot_R50.pth'
+        elif model == 'bot_R50-ibn':
+            cfg['reid_cfg'] = 'Market1501/bagtricks_R50-ibn.yml'
+            cfg['reid_ckpt'] = 'market_bot_R50-ibn.pth'
+        elif model == 'bot_S50':
+            cfg['reid_cfg'] = 'Market1501/bagtricks_S50.yml'
+            cfg['reid_ckpt'] = 'market_bot_S50.pth'
+        elif model == 'bot_R101-ibn':
+            cfg['reid_cfg'] = 'Market1501/bagtricks_R101-ibn.yml'
+            cfg['reid_ckpt'] = 'market_bot_R101-ibn.pth'
+        elif model == 'veh_bot_R50-ibn':
+            cfg['reid_cfg'] = 'VehicleID/bagtricks_R50-ibn.yml'
+            cfg['reid_ckpt'] = 'vehicleid_bot_R50-ibn.pth'
+            
+        return cfg
+
+    @staticmethod
+    def get_cfg_botsort()->dict:
+        '''
+        method: str - Tracking method, 'botsort'
+        device: str - Calucation device, default is 'cuda'
+        half: bool - Half precision, default is False
+        track_high_thresh: float - Tracking confidence threshold, default is 0.5
+        track_low_thresh: float - Lowest detection threshold, default is 0.1
+        new_track_thresh: float - New track thresh, default is 0.6
+        track_buffer: int - The frames for keep lost tracks, default is 30
+        fps: int - Video fps, default is 30
+        match_thresh: float - Matching threshold for tracking, default is 0.8
+        aspect_ratio_thresh: float - Threshold for filtering out boxes of which aspect ratio are above the given value, default is 1.6
+        min_box_area: int - Filter out tiny boxes, default is 10
+        fuse_score: bool - Fuse score and iou for association, default is False
+        mot20: str - MOT20 dataset, default is None
+        cmc_method: str - CMC method, 'sparseOptFlow', 'orb', 'ecc', default is 'sparseOptFlow'
+        ablation: bool - Ablation, default is False
+        name: str - Benchmark name, default is 'None'
+        with_reid: bool - With ReID module, default is True
+        fast_reid_config: str - ReID config file path, default is 'Market1501/sbs_S50.yml'
+        fast_reid_weights: str - ReID config file path, default is 'market_bot_S50.pth'
+        proximity_thresh: float - Threshold for rejecting low overlap reid matches, default is 0.5
+        appearance_thresh: float - Threshold for rejecting low appearance similarity reid matches, default is 0.25
+        '''
+        return { 'method': 'botsort',
+                'device': 'cuda',               # calucation device: cpu, cuda
+                'half': False,                  # half precision: True, False
+                'output_header': False,         # add header in the output file
+
+                 # Track
+                'track_high_thresh': 0.5,       # tracking confidence threshold
+                'track_low_thresh': 0.1,        # lowest detection threshold
+                'new_track_thresh': 0.6,        # new track thresh
+                'track_buffer': 30,             # the frames for keep lost tracks
+                'fps': 30,                      # video fps
+                'match_thresh': 0.8,            # matching threshold for tracking
+                'aspect_ratio_thresh': 1.6,     # threshold for filtering out boxes of which aspect ratio are above the given value.
+                'min_box_area': 10,             # filter out tiny boxes
+                'fuse_score': False,            # fuse score and iou for association
+                'mot20': None,                  # mot20 dataset
+                
+                # CMC
+                'cmc_method': 'sparseOptFlow',  # cmc method: files (Vidstab GMC) | sparseOptFlow |orb | ecc
+                'ablation': False,              # ablation
+                'name': 'None',                 # benchmark name
+             
+                # ReID  
+                'with_reid': True,                             # with ReID module.
+                'fast_reid_config': 'Market1501/sbs_S50.yml',        # reid config file path
+                'fast_reid_weights': 'market_bot_S50.pth',       # reid config file path
+                'proximity_thresh': 0.5,                        # threshold for rejecting low overlap reid matches
+                'appearance_thresh': 0.25                       # threshold for rejecting low appearance similarity reid matches
+            }

dnt/track/dsort/configs/bagtricks_R50.yml

@@ -0,0 +1,7 @@
+_BASE_: ../Base-bagtricks.yml
+
+DATASETS:
+  NAMES: ("Market1501",)
+  TESTS: ("Market1501",)
+
+OUTPUT_DIR: logs/market1501/bagtricks_R50

dnt/track/dsort/configs/fastreid.yaml

@@ -1,3 +1,3 @@
 FASTREID:
-  CFG: "thirdparty/fast-reid/configs/Market1501/bagtricks_R50.yml"
+  CFG: "bagtricks_R50.yml"
   CHECKPOINT: "deep_sort/deep/checkpoint/market_bot_R50.pth"

dnt/track/dsort/deep_sort/deep/feature_extractor.py

@@ -3,13 +3,29 @@ import torchvision.transforms as transforms
 import numpy as np
 import cv2
 import logging
+import os, sys, requests
+from tqdm import tqdm
 
-from .model import Net
+sys.path.append(os.path.join(os.path.dirname(__file__)))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../third_party/fast-reid/")))
+
+from model import Net
+from fastreid.config import get_cfg # type: ignore
+from fastreid.engine import DefaultTrainer # type: ignore
+from fastreid.utils.checkpoint import Checkpointer # type: ignore
 
 class Extractor(object):
-    def __init__(self, model_path, use_cuda=True):
+    def __init__(self, model_path, device:str='auto', half:bool=False):
         self.net = Net(reid=True)
-        self.device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu"
+        if device == 'auto':
+            self.device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+        elif device == 'cuda' and torch.cuda.is_available():
+            self.device = 'cuda'
+        elif device == 'mps' and torch.backends.mps.is_available():
+            self.device = 'mps'
+        else:
+            self.device = 'cpu'        
+        self.half = half
         state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)['net_dict']
         self.net.load_state_dict(state_dict)
         logger = logging.getLogger("root.tracker")
@@ -44,12 +60,75 @@ class Extractor(object):
             features = self.net(im_batch)
         return features.cpu().numpy()
 
+
 class FastReIDExtractor(object):
-    pass
+    def __init__(self, model_config, model_path, device:str='cuda', half:bool=False):
+        cfg = get_cfg()
+        cfg.set_new_allowed(True)
+        cfg.merge_from_file(model_config)
+        cfg.MODEL.BACKBONE.PRETRAIN = False
+        self.net = DefaultTrainer.build_model(cfg)
+        if device == 'auto':
+            self.device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+        elif device == 'cuda' and torch.cuda.is_available():
+            self.device = 'cuda'
+        elif device == 'mps' and torch.backends.mps.is_available():
+            self.device = 'mps'
+        else:
+            self.device = 'cpu'        
+        self.half = half
+
+        if not os.path.exists(model_path):
+            try:
+                url = 'https://github.com/JDAI-CV/fast-reid/releases/download/v0.1.1/'+os.path.basename(model_path)
+                response = requests.get(url, stream=True, allow_redirects=True)
+                with open(model_path, mode="wb") as file:
+                    pbar = tqdm(unit="B", total=int(response.headers['Content-Length']), 
+                                desc="Downloading "+os.path.basename(model_path)+" ...")
+                    for chunk in response.iter_content(chunk_size=10 * 1024):
+                        file.write(chunk)
+                        pbar.update(len(chunk))
+            except ValueError:
+                raise FileNotFoundError(f"File {model_path} not found and cannot be downloaded from {url}!")
+            
+        Checkpointer(self.net).load(model_path)
+        logger = logging.getLogger("root.tracker")
+        logger.info("Loading weights from {}... Done!".format(model_path))
+        self.net.to(self.device)
+        self.net.eval()
+        height, width = cfg.INPUT.SIZE_TEST
+        self.size = (width, height)
+        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)
+
+        im_batch = torch.cat([self.norm(_resize(im, self.size)).unsqueeze(0) for im in im_crops], dim=0).float()
+        return im_batch
+
+
+    def __call__(self, im_crops):
+        im_batch = self._preprocess(im_crops)
+        with torch.no_grad():
+            im_batch = im_batch.to(self.device)
+            features = self.net(im_batch)
+        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)
+    dirname = os.path.dirname(__file__)
+    img_file = os.path.join(dirname, "demo.jpg")
+    imgs = []
+    imgs.append(cv2.imread(img_file)[:,:,(2,1,0)])
+    cfg_file = os.path.join(dirname, "fastreid.yaml")
+    model_1 = os.path.join(dirname, "checkpoint/ckpt.t7")
+    model_2 = os.path.join(dirname, "checkpoint/market_bot_R50.pth")
+    extr1 = Extractor(model_1, use_cuda=True)
+    extr2 = FastReIDExtractor(cfg_file, model_2)
+    feature1 = extr1(imgs)
+    feature2 = extr2(imgs)
+    print(feature1.shape, feature2.shape)
 

dnt/track/dsort/deep_sort/deep_sort.py

@@ -1,3 +1,4 @@
+import sys, os
 import numpy as np
 import torch
 
@@ -7,30 +8,40 @@ from .sort.preprocessing import non_max_suppression
 from .sort.detection import Detection
 from .sort.tracker import Tracker
 
-
 __all__ = ['DeepSort']
 
-
 class DeepSort(object):
-    def __init__(self, model_path, model_config=None, max_dist=0.2, min_confidence=0.3, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=70, n_init=3, nn_budget=100, 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)
+    def __init__(self, cfg:dict, device:str='auto', half:bool=False):
+        
+        self.model_path = cfg['reid_ckpt']              # path to the reid checkpoint (model weights)
+        self.model_config= cfg['reid_cfg']              # path to the reid config file (yaml)
+        self.max_cosine_distance = cfg['max_dist']
+        self.min_confidence = cfg['min_confidence']     # ignore detections if confidence lower than this value
+        self.nms_max_overlap = cfg['nms_max_overlap']
+        self.max_iou_distance = cfg['max_iou_distance'] 
+        self.max_age = cfg['max_age']           # if a track is not matched in max_age frames, it will be deleted
+        self.n_init = cfg['n_init']             # number of detections before creating a track
+        self.nn_budget = cfg['nn_budget']
+        
+        if self.model_config is None:
+            self.extractor = Extractor(os.path.join(os.path.dirname(__file__), 'deep/checkpoint/', self.model_path)
+                                       , device=device, half=half)   # default extractor
         else:
-            self.extractor = FastReIDExtractor(model_config, model_path, use_cuda=use_cuda)
+            cfg_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../third_party/fast-reid/configs/', self.model_config))
+            ckpt_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../third_party/fast-reid/checkpoint/', self.model_path))
+            self.extractor = FastReIDExtractor(cfg_path, ckpt_path, device=device, half=half)    # fast-reid extractor
 
-        max_cosine_distance = max_dist
-        metric = NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
-        self.tracker = Tracker(metric, max_iou_distance=max_iou_distance, max_age=max_age, n_init=n_init)
+        metric = NearestNeighborDistanceMetric("cosine", self.max_cosine_distance, self.nn_budget)
+        self.tracker = Tracker(metric, max_iou_distance=self.max_iou_distance, max_age=self.max_age, n_init=self.n_init)
 
-    def update(self, bbox_xywh, confidences, ori_img):
+    def update(self, bbox_xywh, confidences, classes, ori_img):
         self.height, self.width = ori_img.shape[:2]
+
         # 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]
+        features = self._get_features(bbox_xywh, ori_img)   # extract features for bboxes
+        bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)           # convert bbox from xc_yc_w_h to left top width height/width
+        detections = [Detection(bbox_tlwh[i], conf, classes[i], features[i]) 
+                        for i, conf in enumerate(confidences) if conf>self.min_confidence] # ignore low confidence bboxes
 
         # run on non-maximum supression
         boxes = np.array([d.tlwh for d in detections])
@@ -39,8 +50,8 @@ class DeepSort(object):
         detections = [detections[i] for i in indices]
 
         # update tracker
-        self.tracker.predict()
-        self.tracker.update(detections)
+        self.tracker.predict()              # predict the location of the bboxes
+        self.tracker.update(detections)     # update tracks by location and appearance
 
         # output bbox identities
         outputs = []
@@ -50,12 +61,12 @@ class DeepSort(object):
             box = track.to_tlwh()
             x1,y1,x2,y2 = self._tlwh_to_xyxy(box)
             track_id = track.track_id
-            outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=int))
+            track_cls = track.cls
+            outputs.append(np.array([x1,y1,x2,y2,track_id,track_cls], dtype=int))
         if len(outputs) > 0:
             outputs = np.stack(outputs,axis=0)
         return outputs
 
-
     """
     TODO:
         Convert bbox from xc_yc_w_h to xtl_ytl_w_h
@@ -71,7 +82,6 @@ class DeepSort(object):
         bbox_tlwh[:,1] = bbox_xywh[:,1] - bbox_xywh[:,3]/2.
         return bbox_tlwh
 
-
     def _xywh_to_xyxy(self, bbox_xywh):
         x,y,w,h = bbox_xywh
         x1 = max(int(x-w/2),0)

dnt/track/dsort/deep_sort/sort/detection.py

@@ -26,10 +26,11 @@ class Detection(object):
 
     """
 
-    def __init__(self, tlwh, confidence, feature):
+    def __init__(self, tlwh, confidence, cls, feature):
         self.tlwh = np.asarray(tlwh, dtype=float)
         self.confidence = float(confidence)
         self.feature = np.asarray(feature, dtype=float)
+        self.cls = cls
 
     def to_tlbr(self):
         """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,

dnt/track/dsort/deep_sort/sort/iou_matching.py

@@ -3,7 +3,6 @@ from __future__ import absolute_import
 import numpy as np
 from . import linear_assignment
 
-
 def iou(bbox, candidates):
     """Computer intersection over union.
 
@@ -38,7 +37,6 @@ def iou(bbox, candidates):
     area_candidates = candidates[:, 2:].prod(axis=1)
     return area_intersection / (area_bbox + area_candidates - area_intersection)
 
-
 def iou_cost(tracks, detections, track_indices=None,
              detection_indices=None):
     """An intersection over union distance metric.