Added tracker

app.py

@@ -11,6 +11,7 @@ from utils.plots import plot_one_box
 from utils.torch_utils import time_synchronized
 import time
 from ultralytics import YOLO
+from track import MOT
 
 def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):
     # Resize and pad image while meeting stride-multiple constraints
@@ -178,6 +179,10 @@ def inference_comp(image,iou_threshold,confidence_threshold):
     v7_out, v7_fps = inference(image, "yolov7",iou_threshold,confidence_threshold)
     return v7_out,v8_out,v7_fps,v8_fps
 
+def MODT(sourceVideo, model_link, trackingmethod):
+    model_path = 'weights/'+str(model_link)+'.pt'
+    return MOT(model_path, trackingmethod, sourceVideo), 30
+
 examples_images = ['data/images/1.jpg',
             'data/images/2.jpg',
             'data/images/bus.jpg',
@@ -185,6 +190,7 @@ examples_images = ['data/images/1.jpg',
 examples_videos = ['data/video/1.mp4','data/video/2.mp4'] 
 
 models = ['yolov8m','yolov7','yolov7t']
+trackers = ['strongsort', 'bytetrack', 'ocsort']
 
 with gr.Blocks() as demo:
     gr.Markdown("## IDD Inference on Yolo V7 and V8 ")
@@ -205,11 +211,14 @@ with gr.Blocks() as demo:
             video_input = gr.Video(type='pil', label="Input Video", source="upload")
             video_output = gr.Video(type="pil", label="Output Video",format="mp4")
         fps_video = gr.Number(0,label='FPS')
-        video_drop = gr.Dropdown(choices=models,value=models[0])
+        video_drop = gr.Dropdown(label="Model", choices=models,value=models[0])
+        tracking_drop = gr.Dropdown(label="Tracker", choices=trackers,value=trackers[0])
         video_iou_threshold = gr.Slider(label="IOU Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.45)
         video_conf_threshold = gr.Slider(label="Confidence Threshold",interactive=True, minimum=0.0, maximum=1.0, value=0.25)
         gr.Examples(examples=examples_videos,inputs=video_input,outputs=video_output)
-        video_button = gr.Button("Detect")
+        with gr.Row():
+            video_button_detect = gr.Button("Detect")
+            video_button_track = gr.Button("Track")
     
     with gr.Tab("Compare Models"):
         gr.Markdown("## YOLOv7 vs YOLOv8 Object detection comparision")
@@ -231,12 +240,14 @@ with gr.Blocks() as demo:
     text_button.click(inference, inputs=[image_input,image_drop,
                                          image_iou_threshold,image_conf_threshold],
                                         outputs=[image_output,fps_image])
-    video_button.click(inference2, inputs=[video_input,video_drop,
+    video_button_detect.click(inference2, inputs=[video_input,video_drop,
                                            video_iou_threshold,video_conf_threshold],            
                                         outputs=[video_output,fps_video])
     text_comp_button.click(inference_comp,inputs=[image_comp_input,
                                             image_comp_iou_threshold,
                                             image_comp_conf_threshold],
                                             outputs=[image_comp_output_v7,image_comp_output_v8,v7_fps_image,v8_fps_image])
+    video_button_track.click(MODT,inputs=[video_input,video_drop, tracking_drop],          
+                             outputs=[video_output, fps_video])
 
 demo.launch(debug=True,enable_queue=True)

requirements.txt

@@ -1,6 +1,18 @@
-numpy>=1.18.5
+# pip install -r requirements.txt
+
+# Base ----------------------------------------
+gitpython
+ipython  # interactive notebook
+matplotlib>=3.2.2
+numpy==1.23.1
 opencv-python>=4.1.1
 torch>=1.7.0,!=1.12.0
+Pillow>=7.1.2
+psutil  # system resources
+PyYAML>=5.3.1
+requests>=2.23.0
+scipy>=1.4.1
+thop>=0.1.1  # FLOPs computation
 torchvision>=0.8.1,!=0.13.0
 gradio>=3.9.1
 tqdm>=4.64.0
@@ -8,4 +20,38 @@ seaborn>=0.11.0
 scipy>=1.4.1
 Pillow>=7.1.2
 huggingface-hub >= 0.11.0
-ultralytics >=8.0.34
+ultralytics >=8.0.34
+
+# Logging ---------------------------------------------------------------------
+tensorboard>=2.4.1
+# clearml>=1.2.0
+# comet
+
+# Plotting --------------------------------------------------------------------
+pandas>=1.1.4
+seaborn>=0.11.0
+
+# StrongSORT ------------------------------------------------------------------
+easydict
+
+# torchreid -------------------------------------------------------------------
+gdown
+
+# ByteTrack -------------------------------------------------------------------
+lap
+
+# OCSORT ----------------------------------------------------------------------
+filterpy
+
+# Export ----------------------------------------------------------------------
+# onnx>=1.9.0               # ONNX export
+# onnx-simplifier>=0.4.1    # ONNX simplifier
+# nvidia-pyindex            # TensorRT export
+# nvidia-tensorrt           # TensorRT export
+# openvino-dev              # OpenVINO export
+
+# Hyperparam search -----------------------------------------------------------
+# optuna
+# plotly                    # for hp importance and pareto front plots
+# kaleido
+# joblib

track.py

@@ -0,0 +1,397 @@
+import argparse
+import cv2
+import os
+# limit the number of cpus used by high performance libraries
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+
+import sys
+import platform
+import numpy as np
+from pathlib import Path
+import torch
+import torch.backends.cudnn as cudnn
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]  # yolov5 strongsort root directory
+WEIGHTS = ROOT / 'weights'
+
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+if str(ROOT / 'yolov8') not in sys.path:
+    sys.path.append(str(ROOT / 'yolov8'))  # add yolov5 ROOT to PATH
+if str(ROOT / 'trackers' / 'strongsort') not in sys.path:
+    sys.path.append(str(ROOT / 'trackers' / 'strongsort'))  # add strong_sort ROOT to PATH
+
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+import logging
+#from yolov8.ultralytics.nn.autobackend import AutoBackend
+from ultralytics.nn.autobackend import AutoBackend
+#from yolov8.ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages, LoadStreams
+from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages, LoadStreams
+#from yolov8.ultralytics.yolo.data.utils import IMG_FORMATS, VID_FORMATS
+from ultralytics.yolo.data.utils import IMG_FORMATS, VID_FORMATS
+#from yolov8.ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, SETTINGS, callbacks, colorstr, ops
+from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, SETTINGS, callbacks, colorstr, ops
+
+#from yolov8.ultralytics.yolo.utils.checks import check_file, check_imgsz, check_imshow, print_args, check_requirements
+from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_imshow, print_args, check_requirements
+from ultralytics.yolo.utils.files import increment_path
+from ultralytics.yolo.utils.torch_utils import select_device
+from ultralytics.yolo.utils.ops import Profile, non_max_suppression, scale_boxes, process_mask, process_mask_native
+from ultralytics.yolo.utils.plotting import Annotator, colors, save_one_box
+
+from trackers.multi_tracker_zoo import create_tracker
+
+
+@torch.no_grad()
+def run(
+        source='0',
+        yolo_weights=WEIGHTS / 'yolov5m.pt',  # model.pt path(s),
+        reid_weights=WEIGHTS / 'osnet_x0_25_msmt17.pt',  # model.pt path,
+        tracking_method='strongsort',
+        tracking_config=None,
+        imgsz=(640, 640),  # inference size (height, width)
+        conf_thres=0.25,  # confidence threshold
+        iou_thres=0.45,  # NMS IOU threshold
+        max_det=1000,  # maximum detections per image
+        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
+        show_vid=False,  # show results
+        save_txt=False,  # save results to *.txt
+        save_conf=False,  # save confidences in --save-txt labels
+        save_crop=False,  # save cropped prediction boxes
+        save_trajectories=False,  # save trajectories for each track
+        save_vid=True,  # save confidences in --save-txt labels
+        nosave=False,  # do not save images/videos
+        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
+        update=False,  # update all models
+        #project=ROOT / 'runs' / 'track',  # save results to project/name
+        project=ROOT ,# save results to project/name
+        name='exp',  # save results to project/name
+        exist_ok=True,  # existing project/name ok, do not increment
+        line_thickness=2,  # bounding box thickness (pixels)
+        hide_labels=False,  # hide labels
+        hide_conf=False,  # hide confidences
+        hide_class=False,  # hide IDs
+        half=False,  # use FP16 half-precision inference
+        dnn=False,  # use OpenCV DNN for ONNX inference
+        vid_stride=1,  # video frame-rate stride
+        retina_masks=False,
+):
+
+    source = str(source)
+    save_img = not nosave and not source.endswith('.txt')  # save inference images
+    is_file = Path(source).suffix[1:] in (VID_FORMATS)
+    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
+    webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
+    if is_url and is_file:
+        source = check_file(source)  # download
+
+    # Directories
+    if not isinstance(yolo_weights, list):  # single yolo model
+        exp_name = yolo_weights.stem
+    elif type(yolo_weights) is list and len(yolo_weights) == 1:  # single models after --yolo_weights
+        exp_name = Path(yolo_weights[0]).stem
+    else:  # multiple models after --yolo_weights
+        exp_name = 'ensemble'
+    exp_name = name if name else exp_name + "_" + reid_weights.stem
+    save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok)  # increment run
+    (save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+
+    # Load model
+    device = select_device(device)
+    is_seg = '-seg' in str(yolo_weights)
+    model = AutoBackend(yolo_weights, device=device, dnn=dnn, fp16=half)
+    stride, names, pt = model.stride, model.names, model.pt
+    imgsz = check_imgsz(imgsz, stride=stride)  # check image size
+
+    # Dataloader
+    bs = 1
+    if webcam:
+        show_vid = check_imshow(warn=True)
+        dataset = LoadStreams(
+            source,
+            imgsz=imgsz,
+            stride=stride,
+            auto=pt,
+            transforms=getattr(model.model, 'transforms', None),
+            vid_stride=vid_stride
+        )
+        bs = len(dataset)
+    else:
+        dataset = LoadImages(
+            source,
+            imgsz=imgsz,
+            stride=stride,
+            auto=pt,
+            transforms=getattr(model.model, 'transforms', None),
+            vid_stride=vid_stride
+        )
+    vid_path, vid_writer, txt_path = [None] * bs, [None] * bs, [None] * bs
+    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
+
+    # Create as many strong sort instances as there are video sources
+    tracker_list = []
+    for i in range(bs):
+        tracker = create_tracker(tracking_method, tracking_config, reid_weights, device, half)
+        tracker_list.append(tracker, )
+        if hasattr(tracker_list[i], 'model'):
+            if hasattr(tracker_list[i].model, 'warmup'):
+                tracker_list[i].model.warmup()
+    outputs = [None] * bs
+
+    # Run tracking
+    #model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
+    seen, windows, dt = 0, [], (Profile(), Profile(), Profile(), Profile())
+    curr_frames, prev_frames = [None] * bs, [None] * bs
+    for frame_idx, batch in enumerate(dataset):
+        path, im, im0s, vid_cap, s = batch
+        visualize = increment_path(save_dir / Path(path[0]).stem, mkdir=True) if visualize else False
+        with dt[0]:
+            im = torch.from_numpy(im).to(device)
+            im = im.half() if half else im.float()  # uint8 to fp16/32
+            im /= 255.0  # 0 - 255 to 0.0 - 1.0
+            if len(im.shape) == 3:
+                im = im[None]  # expand for batch dim
+
+        # Inference
+        with dt[1]:
+            preds = model(im, augment=augment, visualize=visualize)
+
+        # Apply NMS
+        with dt[2]:
+            if is_seg:
+                masks = []
+                p = non_max_suppression(preds[0], conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det, nm=32)
+                proto = preds[1][-1]
+            else:
+                p = non_max_suppression(preds, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
+            
+        # Process detections
+        filename = 'out.mp4'
+        for i, det in enumerate(p):  # detections per image
+            seen += 1
+            if webcam:  # bs >= 1
+                p, im0, _ = path[i], im0s[i].copy(), dataset.count
+                p = Path(p)  # to Path
+                s += f'{i}: '
+                txt_file_name = p.name
+                save_path = str(save_dir / filename)  # im.jpg, vid.mp4, ...
+                
+            else:
+                p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
+                p = Path(p)  # to Path
+                # video file
+                if source.endswith(VID_FORMATS):
+                    txt_file_name = p.stem
+                    save_path = str(save_dir / filename)  # im.jpg, vid.mp4, ...
+                    LOGGER.info(f"p.name is {p.name}, save_path value is {save_path}")
+                # folder with imgs
+                else:
+                    txt_file_name = p.parent.name  # get folder name containing current img
+                    save_path = str(save_dir / p.parent.name)  # im.jpg, vid.mp4, ...
+            curr_frames[i] = im0
+
+            txt_path = str(save_dir / 'tracks' / txt_file_name)  # im.txt
+            s += '%gx%g ' % im.shape[2:]  # print string
+            imc = im0.copy() if save_crop else im0  # for save_crop
+
+            annotator = Annotator(im0, line_width=line_thickness, example=str(names))
+            
+            if hasattr(tracker_list[i], 'tracker') and hasattr(tracker_list[i].tracker, 'camera_update'):
+                if prev_frames[i] is not None and curr_frames[i] is not None:  # camera motion compensation
+                    tracker_list[i].tracker.camera_update(prev_frames[i], curr_frames[i])
+
+            if det is not None and len(det):
+                if is_seg:
+                    shape = im0.shape
+                    # scale bbox first the crop masks
+                    if retina_masks:
+                        det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], shape).round()  # rescale boxes to im0 size
+                        masks.append(process_mask_native(proto[i], det[:, 6:], det[:, :4], im0.shape[:2]))  # HWC
+                    else:
+                        masks.append(process_mask(proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True))  # HWC
+                        det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], shape).round()  # rescale boxes to im0 size
+                else:
+                    det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()  # rescale boxes to im0 size
+
+                # Print results
+                for c in det[:, 5].unique():
+                    n = (det[:, 5] == c).sum()  # detections per class
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+
+                # pass detections to strongsort
+                with dt[3]:
+                    outputs[i] = tracker_list[i].update(det.cpu(), im0)
+                
+                # draw boxes for visualization
+                if len(outputs[i]) > 0:
+                    
+                    if is_seg:
+                        # Mask plotting
+                        annotator.masks(
+                            masks[i],
+                            colors=[colors(x, True) for x in det[:, 5]],
+                            im_gpu=torch.as_tensor(im0, dtype=torch.float16).to(device).permute(2, 0, 1).flip(0).contiguous() /
+                            255 if retina_masks else im[i]
+                        )
+                    
+                    for j, (output) in enumerate(outputs[i]):
+                        
+                        bbox = output[0:4]
+                        id = output[4]
+                        cls = output[5]
+                        conf = output[6]
+
+                        if save_txt:
+                            # to MOT format
+                            bbox_left = output[0]
+                            bbox_top = output[1]
+                            bbox_w = output[2] - output[0]
+                            bbox_h = output[3] - output[1]
+                            # Write MOT compliant results to file
+                            with open(txt_path + '.txt', 'a') as f:
+                                f.write(('%g ' * 10 + '\n') % (frame_idx + 1, id, bbox_left,  # MOT format
+                                                               bbox_top, bbox_w, bbox_h, -1, -1, -1, i))
+
+                        if save_vid or save_crop or show_vid:  # Add bbox/seg to image
+                            c = int(cls)  # integer class
+                            id = int(id)  # integer id
+                            label = None if hide_labels else (f'{id} {names[c]}' if hide_conf else \
+                                (f'{id} {conf:.2f}' if hide_class else f'{id} {names[c]} {conf:.2f}'))
+                            color = colors(c, True)
+                            annotator.box_label(bbox, label, color=color)
+                            
+                            if save_trajectories and tracking_method == 'strongsort':
+                                q = output[7]
+                                tracker_list[i].trajectory(im0, q, color=color)
+                            if save_crop:
+                                txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''
+                                save_one_box(np.array(bbox, dtype=np.int16), imc, file=save_dir / 'crops' / txt_file_name / names[c] / f'{id}' / f'{p.stem}.jpg', BGR=True)
+                            
+            else:
+                pass
+                #tracker_list[i].tracker.pred_n_update_all_tracks()
+                
+            # Stream results
+            im0 = annotator.result()
+            if show_vid:
+                if platform.system() == 'Linux' and p not in windows:
+                    windows.append(p)
+                    cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
+                    cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
+                cv2.imshow(str(p), im0)
+                if cv2.waitKey(1) == ord('q'):  # 1 millisecond
+                    exit()
+
+            # Save results (image with detections)
+            if save_vid:
+                LOGGER.info(f"vid_path, save_path {vid_path[i]}{save_path}")
+                if vid_path[i] != save_path:  # new video
+                    vid_path[i] = save_path
+                    if isinstance(vid_writer[i], cv2.VideoWriter):
+                        vid_writer[i].release()  # release previous video writer
+                    if vid_cap:  # video
+                        fps = vid_cap.get(cv2.CAP_PROP_FPS)
+                        w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+                        h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                    else:  # stream
+                        fps, w, h = 30, im0.shape[1], im0.shape[0]
+                    save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
+                    LOGGER.info(f"test Results saved to {colorstr('bold', save_path)}")
+                    vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+                vid_writer[i].write(im0)
+
+            prev_frames[i] = curr_frames[i]
+            
+        # Print total time (preprocessing + inference + NMS + tracking)
+        LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{sum([dt.dt for dt in dt if hasattr(dt, 'dt')]) * 1E3:.1f}ms")
+
+    # Print results
+    t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
+    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms {tracking_method} update per image at shape {(1, 3, *imgsz)}' % t)
+    if save_txt or save_vid:
+        s = f"\n{len(list((save_dir / 'tracks').glob('*.txt')))} tracks saved to {save_dir / 'tracks'}" if save_txt else ''
+        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
+    if update:
+        strip_optimizer(yolo_weights)  # update model (to fix SourceChangeWarning)
+
+
+def parse_opt():
+    parser = argparse.ArgumentParser()
+    #parser.add_argument('--yolo-weights', nargs='+', type=Path, default=WEIGHTS / 'yolov8s-seg.pt', help='model.pt path(s)')
+    parser.add_argument('--reid-weights', type=Path, default=WEIGHTS / 'osnet_x0_25_msmt17.pt')
+    #parser.add_argument('--tracking-method', type=str, default='bytetrack', help='strongsort, ocsort, bytetrack')
+    parser.add_argument('--tracking-config', type=Path, default=None)
+    #parser.add_argument('--source', type=str, default='0', help='file/dir/URL/glob, 0 for webcam')  
+    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
+    parser.add_argument('--conf-thres', type=float, default=0.5, help='confidence threshold')
+    parser.add_argument('--iou-thres', type=float, default=0.5, help='NMS IoU threshold')
+    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
+    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--show-vid', action='store_true', help='display tracking video results')
+    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
+    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
+    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
+    parser.add_argument('--save-trajectories', action='store_true', help='save trajectories for each track')
+    parser.add_argument('--save-vid', action='store_true',default=True, help='save video tracking results')
+    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
+    # class 0 is person, 1 is bycicle, 2 is car... 79 is oven
+    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
+    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
+    parser.add_argument('--augment', action='store_true', help='augmented inference')
+    parser.add_argument('--visualize', action='store_true', help='visualize features')
+    parser.add_argument('--update', action='store_true', help='update all models')
+    parser.add_argument('--project', default=ROOT , help='save results to project/name')
+    parser.add_argument('--name', default='exp', help='save results to ROOT')
+    parser.add_argument('--exist-ok', default='True', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument('--line-thickness', default=2, type=int, help='bounding box thickness (pixels)')
+    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
+    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
+    parser.add_argument('--hide-class', default=False, action='store_true', help='hide IDs')
+    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
+    parser.add_argument('--retina-masks', action='store_true', help='whether to plot masks in native resolution')
+    #opt = parser.parse_args()
+    #opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+    #opt.tracking_config = ROOT / 'trackers' / opt.tracking_method / 'configs' / (opt.tracking_method + '.yaml')
+    #print_args(vars(opt))
+    #return opt
+    return parser
+
+
+def main(opt):
+    check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    run(**vars(opt))
+
+
+#if __name__ == "__main__":
+#    opt = parse_opt()
+#    main(opt)
+
+def MOT(yoloweights, trackingmethod, sourceVideo):
+    parser = parse_opt()
+    parser.add_argument('--yolo-weights', nargs='+', type=Path, default= yoloweights, help='model.pt path(s)')
+    parser.add_argument('--tracking-method', type=str, default= trackingmethod, help='strongsort, ocsort, bytetrack')
+    parser.add_argument('--source', type=str, default=sourceVideo, help='file/dir/URL/glob, 0 for webcam')  
+    opt = parser.parse_args()
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+    opt.tracking_config = ROOT / 'trackers' / opt.tracking_method / 'configs' / (opt.tracking_method + '.yaml')
+    print_args(vars(opt))
+    main(opt)
+    save_dir = increment_path('exp', exist_ok=True) 
+    input = os.path.join(save_dir,'out.mp4')
+    outpath = 'output.mp4'  #'output/'+ 'output.mp4' 
+    command = f"ffmpeg -i {input} -vf fps=30 -vcodec libx264 {outpath}"
+    print(command)
+    os.system(command)
+    #!ffmpeg -i $input -vf fps=30 -vcodec libx264 $outpath tbd
+    return outpath

trackers/bytetrack/basetrack.py

@@ -0,0 +1,52 @@
+import numpy as np
+from collections import OrderedDict
+
+
+class TrackState(object):
+    New = 0
+    Tracked = 1
+    Lost = 2
+    Removed = 3
+
+
+class BaseTrack(object):
+    _count = 0
+
+    track_id = 0
+    is_activated = False
+    state = TrackState.New
+
+    history = OrderedDict()
+    features = []
+    curr_feature = None
+    score = 0
+    start_frame = 0
+    frame_id = 0
+    time_since_update = 0
+
+    # multi-camera
+    location = (np.inf, np.inf)
+
+    @property
+    def end_frame(self):
+        return self.frame_id
+
+    @staticmethod
+    def next_id():
+        BaseTrack._count += 1
+        return BaseTrack._count
+
+    def activate(self, *args):
+        raise NotImplementedError
+
+    def predict(self):
+        raise NotImplementedError
+
+    def update(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def mark_lost(self):
+        self.state = TrackState.Lost
+
+    def mark_removed(self):
+        self.state = TrackState.Removed

trackers/bytetrack/byte_tracker.py

@@ -0,0 +1,353 @@
+import numpy as np
+from collections import deque
+import os
+import os.path as osp
+import copy
+import torch
+import torch.nn.functional as F
+
+from ultralytics.yolo.utils.ops import xywh2xyxy, xyxy2xywh
+
+
+from trackers.bytetrack.kalman_filter import KalmanFilter
+from trackers.bytetrack import matching
+from trackers.bytetrack.basetrack import BaseTrack, TrackState
+
+class STrack(BaseTrack):
+    shared_kalman = KalmanFilter()
+    def __init__(self, tlwh, score, cls):
+
+        # wait activate
+        self._tlwh = np.asarray(tlwh, dtype=np.float32)
+        self.kalman_filter = None
+        self.mean, self.covariance = None, None
+        self.is_activated = False
+
+        self.score = score
+        self.tracklet_len = 0
+        self.cls = cls
+
+    def predict(self):
+        mean_state = self.mean.copy()
+        if self.state != TrackState.Tracked:
+            mean_state[7] = 0
+        self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)
+
+    @staticmethod
+    def multi_predict(stracks):
+        if len(stracks) > 0:
+            multi_mean = np.asarray([st.mean.copy() for st in stracks])
+            multi_covariance = np.asarray([st.covariance for st in stracks])
+            for i, st in enumerate(stracks):
+                if st.state != TrackState.Tracked:
+                    multi_mean[i][7] = 0
+            multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(multi_mean, multi_covariance)
+            for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+                stracks[i].mean = mean
+                stracks[i].covariance = cov
+
+    def activate(self, kalman_filter, frame_id):
+        """Start a new tracklet"""
+        self.kalman_filter = kalman_filter
+        self.track_id = self.next_id()
+        self.mean, self.covariance = self.kalman_filter.initiate(self.tlwh_to_xyah(self._tlwh))
+
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        if frame_id == 1:
+            self.is_activated = True
+        # self.is_activated = True
+        self.frame_id = frame_id
+        self.start_frame = frame_id
+
+    def re_activate(self, new_track, frame_id, new_id=False):
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
+        )
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        self.is_activated = True
+        self.frame_id = frame_id
+        if new_id:
+            self.track_id = self.next_id()
+        self.score = new_track.score
+        self.cls = new_track.cls
+
+    def update(self, new_track, frame_id):
+        """
+        Update a matched track
+        :type new_track: STrack
+        :type frame_id: int
+        :type update_feature: bool
+        :return:
+        """
+        self.frame_id = frame_id
+        self.tracklet_len += 1
+        # self.cls = cls
+
+        new_tlwh = new_track.tlwh
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh))
+        self.state = TrackState.Tracked
+        self.is_activated = True
+
+        self.score = new_track.score
+
+    @property
+    # @jit(nopython=True)
+    def tlwh(self):
+        """Get current position in bounding box format `(top left x, top left y,
+                width, height)`.
+        """
+        if self.mean is None:
+            return self._tlwh.copy()
+        ret = self.mean[:4].copy()
+        ret[2] *= ret[3]
+        ret[:2] -= ret[2:] / 2
+        return ret
+
+    @property
+    # @jit(nopython=True)
+    def tlbr(self):
+        """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
+        `(top left, bottom right)`.
+        """
+        ret = self.tlwh.copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    @staticmethod
+    # @jit(nopython=True)
+    def tlwh_to_xyah(tlwh):
+        """Convert bounding box to format `(center x, center y, aspect ratio,
+        height)`, where the aspect ratio is `width / height`.
+        """
+        ret = np.asarray(tlwh).copy()
+        ret[:2] += ret[2:] / 2
+        ret[2] /= ret[3]
+        return ret
+
+    def to_xyah(self):
+        return self.tlwh_to_xyah(self.tlwh)
+
+    @staticmethod
+    # @jit(nopython=True)
+    def tlbr_to_tlwh(tlbr):
+        ret = np.asarray(tlbr).copy()
+        ret[2:] -= ret[:2]
+        return ret
+
+    @staticmethod
+    # @jit(nopython=True)
+    def tlwh_to_tlbr(tlwh):
+        ret = np.asarray(tlwh).copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    def __repr__(self):
+        return 'OT_{}_({}-{})'.format(self.track_id, self.start_frame, self.end_frame)
+
+
+class BYTETracker(object):
+    def __init__(self, track_thresh=0.45, match_thresh=0.8, track_buffer=25, frame_rate=30):
+        self.tracked_stracks = []  # type: list[STrack]
+        self.lost_stracks = []  # type: list[STrack]
+        self.removed_stracks = []  # type: list[STrack]
+
+        self.frame_id = 0
+        self.track_buffer=track_buffer
+        
+        self.track_thresh = track_thresh
+        self.match_thresh = match_thresh
+        self.det_thresh = track_thresh + 0.1
+        self.buffer_size = int(frame_rate / 30.0 * track_buffer)
+        self.max_time_lost = self.buffer_size
+        self.kalman_filter = KalmanFilter()
+
+    def update(self, dets, _):
+        self.frame_id += 1
+        activated_starcks = []
+        refind_stracks = []
+        lost_stracks = []
+        removed_stracks = []
+
+        xyxys = dets[:, 0:4]
+        xywh = xyxy2xywh(xyxys)
+        confs = dets[:, 4]
+        clss = dets[:, 5]
+        
+        classes = clss.numpy()
+        xyxys = xyxys.numpy()
+        confs = confs.numpy()
+
+        remain_inds = confs > self.track_thresh
+        inds_low = confs > 0.1
+        inds_high = confs < self.track_thresh
+
+        inds_second = np.logical_and(inds_low, inds_high)
+        
+        dets_second = xywh[inds_second]
+        dets = xywh[remain_inds]
+        
+        scores_keep = confs[remain_inds]
+        scores_second = confs[inds_second]
+        
+        clss_keep = classes[remain_inds]
+        clss_second = classes[inds_second]
+        
+
+        if len(dets) > 0:
+            '''Detections'''
+            detections = [STrack(xyxy, s, c) for (xyxy, s, c) in zip(dets, scores_keep, clss_keep)]
+        else:
+            detections = []
+
+        ''' Add newly detected tracklets to tracked_stracks'''
+        unconfirmed = []
+        tracked_stracks = []  # type: list[STrack]
+        for track in self.tracked_stracks:
+            if not track.is_activated:
+                unconfirmed.append(track)
+            else:
+                tracked_stracks.append(track)
+
+        ''' Step 2: First association, with high score detection boxes'''
+        strack_pool = joint_stracks(tracked_stracks, self.lost_stracks)
+        # Predict the current location with KF
+        STrack.multi_predict(strack_pool)
+        dists = matching.iou_distance(strack_pool, detections)
+        #if not self.args.mot20:
+        dists = matching.fuse_score(dists, detections)
+        matches, u_track, u_detection = matching.linear_assignment(dists, thresh=self.match_thresh)
+
+        for itracked, idet in matches:
+            track = strack_pool[itracked]
+            det = detections[idet]
+            if track.state == TrackState.Tracked:
+                track.update(detections[idet], self.frame_id)
+                activated_starcks.append(track)
+            else:
+                track.re_activate(det, self.frame_id, new_id=False)
+                refind_stracks.append(track)
+
+        ''' Step 3: Second association, with low score detection boxes'''
+        # association the untrack to the low score detections
+        if len(dets_second) > 0:
+            '''Detections'''
+            detections_second = [STrack(xywh, s, c) for (xywh, s, c) in zip(dets_second, scores_second, clss_second)]
+        else:
+            detections_second = []
+        r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state == TrackState.Tracked]
+        dists = matching.iou_distance(r_tracked_stracks, detections_second)
+        matches, u_track, u_detection_second = matching.linear_assignment(dists, thresh=0.5)
+        for itracked, idet in matches:
+            track = r_tracked_stracks[itracked]
+            det = detections_second[idet]
+            if track.state == TrackState.Tracked:
+                track.update(det, self.frame_id)
+                activated_starcks.append(track)
+            else:
+                track.re_activate(det, self.frame_id, new_id=False)
+                refind_stracks.append(track)
+
+        for it in u_track:
+            track = r_tracked_stracks[it]
+            if not track.state == TrackState.Lost:
+                track.mark_lost()
+                lost_stracks.append(track)
+
+        '''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
+        detections = [detections[i] for i in u_detection]
+        dists = matching.iou_distance(unconfirmed, detections)
+        #if not self.args.mot20:
+        dists = matching.fuse_score(dists, detections)
+        matches, u_unconfirmed, u_detection = matching.linear_assignment(dists, thresh=0.7)
+        for itracked, idet in matches:
+            unconfirmed[itracked].update(detections[idet], self.frame_id)
+            activated_starcks.append(unconfirmed[itracked])
+        for it in u_unconfirmed:
+            track = unconfirmed[it]
+            track.mark_removed()
+            removed_stracks.append(track)
+
+        """ Step 4: Init new stracks"""
+        for inew in u_detection:
+            track = detections[inew]
+            if track.score < self.det_thresh:
+                continue
+            track.activate(self.kalman_filter, self.frame_id)
+            activated_starcks.append(track)
+        """ Step 5: Update state"""
+        for track in self.lost_stracks:
+            if self.frame_id - track.end_frame > self.max_time_lost:
+                track.mark_removed()
+                removed_stracks.append(track)
+
+        # print('Ramained match {} s'.format(t4-t3))
+
+        self.tracked_stracks = [t for t in self.tracked_stracks if t.state == TrackState.Tracked]
+        self.tracked_stracks = joint_stracks(self.tracked_stracks, activated_starcks)
+        self.tracked_stracks = joint_stracks(self.tracked_stracks, refind_stracks)
+        self.lost_stracks = sub_stracks(self.lost_stracks, self.tracked_stracks)
+        self.lost_stracks.extend(lost_stracks)
+        self.lost_stracks = sub_stracks(self.lost_stracks, self.removed_stracks)
+        self.removed_stracks.extend(removed_stracks)
+        self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)
+        # get scores of lost tracks
+        output_stracks = [track for track in self.tracked_stracks if track.is_activated]
+        outputs = []
+        for t in output_stracks:
+            output= []
+            tlwh = t.tlwh
+            tid = t.track_id
+            tlwh = np.expand_dims(tlwh, axis=0)
+            xyxy = xywh2xyxy(tlwh)
+            xyxy = np.squeeze(xyxy, axis=0)
+            output.extend(xyxy)
+            output.append(tid)
+            output.append(t.cls)
+            output.append(t.score)
+            outputs.append(output)
+
+        return outputs
+#track_id, class_id, conf
+
+def joint_stracks(tlista, tlistb):
+    exists = {}
+    res = []
+    for t in tlista:
+        exists[t.track_id] = 1
+        res.append(t)
+    for t in tlistb:
+        tid = t.track_id
+        if not exists.get(tid, 0):
+            exists[tid] = 1
+            res.append(t)
+    return res
+
+
+def sub_stracks(tlista, tlistb):
+    stracks = {}
+    for t in tlista:
+        stracks[t.track_id] = t
+    for t in tlistb:
+        tid = t.track_id
+        if stracks.get(tid, 0):
+            del stracks[tid]
+    return list(stracks.values())
+
+
+def remove_duplicate_stracks(stracksa, stracksb):
+    pdist = matching.iou_distance(stracksa, stracksb)
+    pairs = np.where(pdist < 0.15)
+    dupa, dupb = list(), list()
+    for p, q in zip(*pairs):
+        timep = stracksa[p].frame_id - stracksa[p].start_frame
+        timeq = stracksb[q].frame_id - stracksb[q].start_frame
+        if timep > timeq:
+            dupb.append(q)
+        else:
+            dupa.append(p)
+    resa = [t for i, t in enumerate(stracksa) if not i in dupa]
+    resb = [t for i, t in enumerate(stracksb) if not i in dupb]
+    return resa, resb

trackers/bytetrack/configs/bytetrack.yaml

@@ -0,0 +1,7 @@
+bytetrack:
+  track_thresh: 0.6  # tracking confidence threshold
+  track_buffer: 30   # the frames for keep lost tracks
+  match_thresh: 0.8  # matching threshold for tracking
+  frame_rate: 30     # FPS
+  conf_thres: 0.5122620708221085
+  

trackers/bytetrack/kalman_filter.py

@@ -0,0 +1,270 @@
+# vim: expandtab:ts=4:sw=4
+import numpy as np
+import scipy.linalg
+
+
+"""
+Table for the 0.95 quantile of the chi-square distribution with N degrees of
+freedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv
+function and used as Mahalanobis gating threshold.
+"""
+chi2inv95 = {
+    1: 3.8415,
+    2: 5.9915,
+    3: 7.8147,
+    4: 9.4877,
+    5: 11.070,
+    6: 12.592,
+    7: 14.067,
+    8: 15.507,
+    9: 16.919}
+
+
+class KalmanFilter(object):
+    """
+    A simple Kalman filter for tracking bounding boxes in image space.
+
+    The 8-dimensional state space
+
+        x, y, a, h, vx, vy, va, vh
+
+    contains the bounding box center position (x, y), aspect ratio a, height h,
+    and their respective velocities.
+
+    Object motion follows a constant velocity model. The bounding box location
+    (x, y, a, h) is taken as direct observation of the state space (linear
+    observation model).
+
+    """
+
+    def __init__(self):
+        ndim, dt = 4, 1.
+
+        # Create Kalman filter model matrices.
+        self._motion_mat = np.eye(2 * ndim, 2 * ndim)
+        for i in range(ndim):
+            self._motion_mat[i, ndim + i] = dt
+        self._update_mat = np.eye(ndim, 2 * ndim)
+
+        # Motion and observation uncertainty are chosen relative to the current
+        # state estimate. These weights control the amount of uncertainty in
+        # the model. This is a bit hacky.
+        self._std_weight_position = 1. / 20
+        self._std_weight_velocity = 1. / 160
+
+    def initiate(self, measurement):
+        """Create track from unassociated measurement.
+
+        Parameters
+        ----------
+        measurement : ndarray
+            Bounding box coordinates (x, y, a, h) with center position (x, y),
+            aspect ratio a, and height h.
+
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the mean vector (8 dimensional) and covariance matrix (8x8
+            dimensional) of the new track. Unobserved velocities are initialized
+            to 0 mean.
+
+        """
+        mean_pos = measurement
+        mean_vel = np.zeros_like(mean_pos)
+        mean = np.r_[mean_pos, mean_vel]
+
+        std = [
+            2 * self._std_weight_position * measurement[3],
+            2 * self._std_weight_position * measurement[3],
+            1e-2,
+            2 * self._std_weight_position * measurement[3],
+            10 * self._std_weight_velocity * measurement[3],
+            10 * self._std_weight_velocity * measurement[3],
+            1e-5,
+            10 * self._std_weight_velocity * measurement[3]]
+        covariance = np.diag(np.square(std))
+        return mean, covariance
+
+    def predict(self, mean, covariance):
+        """Run Kalman filter prediction step.
+
+        Parameters
+        ----------
+        mean : ndarray
+            The 8 dimensional mean vector of the object state at the previous
+            time step.
+        covariance : ndarray
+            The 8x8 dimensional covariance matrix of the object state at the
+            previous time step.
+
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the mean vector and covariance matrix of the predicted
+            state. Unobserved velocities are initialized to 0 mean.
+
+        """
+        std_pos = [
+            self._std_weight_position * mean[3],
+            self._std_weight_position * mean[3],
+            1e-2,
+            self._std_weight_position * mean[3]]
+        std_vel = [
+            self._std_weight_velocity * mean[3],
+            self._std_weight_velocity * mean[3],
+            1e-5,
+            self._std_weight_velocity * mean[3]]
+        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
+
+        #mean = np.dot(self._motion_mat, mean)
+        mean = np.dot(mean, self._motion_mat.T)
+        covariance = np.linalg.multi_dot((
+            self._motion_mat, covariance, self._motion_mat.T)) + motion_cov
+
+        return mean, covariance
+
+    def project(self, mean, covariance):
+        """Project state distribution to measurement space.
+
+        Parameters
+        ----------
+        mean : ndarray
+            The state's mean vector (8 dimensional array).
+        covariance : ndarray
+            The state's covariance matrix (8x8 dimensional).
+
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the projected mean and covariance matrix of the given state
+            estimate.
+
+        """
+        std = [
+            self._std_weight_position * mean[3],
+            self._std_weight_position * mean[3],
+            1e-1,
+            self._std_weight_position * mean[3]]
+        innovation_cov = np.diag(np.square(std))
+
+        mean = np.dot(self._update_mat, mean)
+        covariance = np.linalg.multi_dot((
+            self._update_mat, covariance, self._update_mat.T))
+        return mean, covariance + innovation_cov
+
+    def multi_predict(self, mean, covariance):
+        """Run Kalman filter prediction step (Vectorized version).
+        Parameters
+        ----------
+        mean : ndarray
+            The Nx8 dimensional mean matrix of the object states at the previous
+            time step.
+        covariance : ndarray
+            The Nx8x8 dimensional covariance matrics of the object states at the
+            previous time step.
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the mean vector and covariance matrix of the predicted
+            state. Unobserved velocities are initialized to 0 mean.
+        """
+        std_pos = [
+            self._std_weight_position * mean[:, 3],
+            self._std_weight_position * mean[:, 3],
+            1e-2 * np.ones_like(mean[:, 3]),
+            self._std_weight_position * mean[:, 3]]
+        std_vel = [
+            self._std_weight_velocity * mean[:, 3],
+            self._std_weight_velocity * mean[:, 3],
+            1e-5 * np.ones_like(mean[:, 3]),
+            self._std_weight_velocity * mean[:, 3]]
+        sqr = np.square(np.r_[std_pos, std_vel]).T
+
+        motion_cov = []
+        for i in range(len(mean)):
+            motion_cov.append(np.diag(sqr[i]))
+        motion_cov = np.asarray(motion_cov)
+
+        mean = np.dot(mean, self._motion_mat.T)
+        left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
+        covariance = np.dot(left, self._motion_mat.T) + motion_cov
+
+        return mean, covariance
+
+    def update(self, mean, covariance, measurement):
+        """Run Kalman filter correction step.
+
+        Parameters
+        ----------
+        mean : ndarray
+            The predicted state's mean vector (8 dimensional).
+        covariance : ndarray
+            The state's covariance matrix (8x8 dimensional).
+        measurement : ndarray
+            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
+            is the center position, a the aspect ratio, and h the height of the
+            bounding box.
+
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the measurement-corrected state distribution.
+
+        """
+        projected_mean, projected_cov = self.project(mean, covariance)
+
+        chol_factor, lower = scipy.linalg.cho_factor(
+            projected_cov, lower=True, check_finite=False)
+        kalman_gain = scipy.linalg.cho_solve(
+            (chol_factor, lower), np.dot(covariance, self._update_mat.T).T,
+            check_finite=False).T
+        innovation = measurement - projected_mean
+
+        new_mean = mean + np.dot(innovation, kalman_gain.T)
+        new_covariance = covariance - np.linalg.multi_dot((
+            kalman_gain, projected_cov, kalman_gain.T))
+        return new_mean, new_covariance
+
+    def gating_distance(self, mean, covariance, measurements,
+                        only_position=False, metric='maha'):
+        """Compute gating distance between state distribution and measurements.
+        A suitable distance threshold can be obtained from `chi2inv95`. If
+        `only_position` is False, the chi-square distribution has 4 degrees of
+        freedom, otherwise 2.
+        Parameters
+        ----------
+        mean : ndarray
+            Mean vector over the state distribution (8 dimensional).
+        covariance : ndarray
+            Covariance of the state distribution (8x8 dimensional).
+        measurements : ndarray
+            An Nx4 dimensional matrix of N measurements, each in
+            format (x, y, a, h) where (x, y) is the bounding box center
+            position, a the aspect ratio, and h the height.
+        only_position : Optional[bool]
+            If True, distance computation is done with respect to the bounding
+            box center position only.
+        Returns
+        -------
+        ndarray
+            Returns an array of length N, where the i-th element contains the
+            squared Mahalanobis distance between (mean, covariance) and
+            `measurements[i]`.
+        """
+        mean, covariance = self.project(mean, covariance)
+        if only_position:
+            mean, covariance = mean[:2], covariance[:2, :2]
+            measurements = measurements[:, :2]
+
+        d = measurements - mean
+        if metric == 'gaussian':
+            return np.sum(d * d, axis=1)
+        elif metric == 'maha':
+            cholesky_factor = np.linalg.cholesky(covariance)
+            z = scipy.linalg.solve_triangular(
+                cholesky_factor, d.T, lower=True, check_finite=False,
+                overwrite_b=True)
+            squared_maha = np.sum(z * z, axis=0)
+            return squared_maha
+        else:
+            raise ValueError('invalid distance metric')

trackers/bytetrack/matching.py

@@ -0,0 +1,219 @@
+import cv2
+import numpy as np
+import scipy
+import lap
+from scipy.spatial.distance import cdist
+
+from trackers.bytetrack import kalman_filter
+import time
+
+def merge_matches(m1, m2, shape):
+    O,P,Q = shape
+    m1 = np.asarray(m1)
+    m2 = np.asarray(m2)
+
+    M1 = scipy.sparse.coo_matrix((np.ones(len(m1)), (m1[:, 0], m1[:, 1])), shape=(O, P))
+    M2 = scipy.sparse.coo_matrix((np.ones(len(m2)), (m2[:, 0], m2[:, 1])), shape=(P, Q))
+
+    mask = M1*M2
+    match = mask.nonzero()
+    match = list(zip(match[0], match[1]))
+    unmatched_O = tuple(set(range(O)) - set([i for i, j in match]))
+    unmatched_Q = tuple(set(range(Q)) - set([j for i, j in match]))
+
+    return match, unmatched_O, unmatched_Q
+
+
+def _indices_to_matches(cost_matrix, indices, thresh):
+    matched_cost = cost_matrix[tuple(zip(*indices))]
+    matched_mask = (matched_cost <= thresh)
+
+    matches = indices[matched_mask]
+    unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
+    unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
+
+    return matches, unmatched_a, unmatched_b
+
+
+def linear_assignment(cost_matrix, thresh):
+    if cost_matrix.size == 0:
+        return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
+    matches, unmatched_a, unmatched_b = [], [], []
+    cost, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
+    for ix, mx in enumerate(x):
+        if mx >= 0:
+            matches.append([ix, mx])
+    unmatched_a = np.where(x < 0)[0]
+    unmatched_b = np.where(y < 0)[0]
+    matches = np.asarray(matches)
+    return matches, unmatched_a, unmatched_b
+
+
+def ious(atlbrs, btlbrs):
+    """
+    Compute cost based on IoU
+    :type atlbrs: list[tlbr] | np.ndarray
+    :type atlbrs: list[tlbr] | np.ndarray
+
+    :rtype ious np.ndarray
+    """
+    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+    if ious.size == 0:
+        return ious
+
+    ious = bbox_ious(
+        np.ascontiguousarray(atlbrs, dtype=np.float32),
+        np.ascontiguousarray(btlbrs, dtype=np.float32)
+    )
+
+    return ious
+
+
+def iou_distance(atracks, btracks):
+    """
+    Compute cost based on IoU
+    :type atracks: list[STrack]
+    :type btracks: list[STrack]
+
+    :rtype cost_matrix np.ndarray
+    """
+
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlbr for track in atracks]
+        btlbrs = [track.tlbr for track in btracks]
+    _ious = ious(atlbrs, btlbrs)
+    cost_matrix = 1 - _ious
+
+    return cost_matrix
+
+def v_iou_distance(atracks, btracks):
+    """
+    Compute cost based on IoU
+    :type atracks: list[STrack]
+    :type btracks: list[STrack]
+
+    :rtype cost_matrix np.ndarray
+    """
+
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlwh_to_tlbr(track.pred_bbox) for track in atracks]
+        btlbrs = [track.tlwh_to_tlbr(track.pred_bbox) for track in btracks]
+    _ious = ious(atlbrs, btlbrs)
+    cost_matrix = 1 - _ious
+
+    return cost_matrix
+
+def embedding_distance(tracks, detections, metric='cosine'):
+    """
+    :param tracks: list[STrack]
+    :param detections: list[BaseTrack]
+    :param metric:
+    :return: cost_matrix np.ndarray
+    """
+
+    cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float32)
+    if cost_matrix.size == 0:
+        return cost_matrix
+    det_features = np.asarray([track.curr_feat for track in detections], dtype=np.float32)
+    #for i, track in enumerate(tracks):
+        #cost_matrix[i, :] = np.maximum(0.0, cdist(track.smooth_feat.reshape(1,-1), det_features, metric))
+    track_features = np.asarray([track.smooth_feat for track in tracks], dtype=np.float32)
+    cost_matrix = np.maximum(0.0, cdist(track_features, det_features, metric))  # Nomalized features
+    return cost_matrix
+
+
+def gate_cost_matrix(kf, cost_matrix, tracks, detections, only_position=False):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    gating_dim = 2 if only_position else 4
+    gating_threshold = kalman_filter.chi2inv95[gating_dim]
+    measurements = np.asarray([det.to_xyah() for det in detections])
+    for row, track in enumerate(tracks):
+        gating_distance = kf.gating_distance(
+            track.mean, track.covariance, measurements, only_position)
+        cost_matrix[row, gating_distance > gating_threshold] = np.inf
+    return cost_matrix
+
+
+def fuse_motion(kf, cost_matrix, tracks, detections, only_position=False, lambda_=0.98):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    gating_dim = 2 if only_position else 4
+    gating_threshold = kalman_filter.chi2inv95[gating_dim]
+    measurements = np.asarray([det.to_xyah() for det in detections])
+    for row, track in enumerate(tracks):
+        gating_distance = kf.gating_distance(
+            track.mean, track.covariance, measurements, only_position, metric='maha')
+        cost_matrix[row, gating_distance > gating_threshold] = np.inf
+        cost_matrix[row] = lambda_ * cost_matrix[row] + (1 - lambda_) * gating_distance
+    return cost_matrix
+
+
+def fuse_iou(cost_matrix, tracks, detections):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    reid_sim = 1 - cost_matrix
+    iou_dist = iou_distance(tracks, detections)
+    iou_sim = 1 - iou_dist
+    fuse_sim = reid_sim * (1 + iou_sim) / 2
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    #fuse_sim = fuse_sim * (1 + det_scores) / 2
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost
+
+
+def fuse_score(cost_matrix, detections):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    iou_sim = 1 - cost_matrix
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    fuse_sim = iou_sim * det_scores
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost
+
+
+def bbox_ious(boxes, query_boxes):
+    """
+    Parameters
+    ----------
+    boxes: (N, 4) ndarray of float
+    query_boxes: (K, 4) ndarray of float
+    Returns
+    -------
+    overlaps: (N, K) ndarray of overlap between boxes and query_boxes
+    """
+    N = boxes.shape[0]
+    K = query_boxes.shape[0]
+    overlaps = np.zeros((N, K), dtype=np.float32)
+    
+    for k in range(K):
+        box_area = (
+            (query_boxes[k, 2] - query_boxes[k, 0] + 1) *
+            (query_boxes[k, 3] - query_boxes[k, 1] + 1)
+        )
+        for n in range(N):
+            iw = (
+                min(boxes[n, 2], query_boxes[k, 2]) -
+                max(boxes[n, 0], query_boxes[k, 0]) + 1
+            )
+            if iw > 0:
+                ih = (
+                    min(boxes[n, 3], query_boxes[k, 3]) -
+                    max(boxes[n, 1], query_boxes[k, 1]) + 1
+                )
+                if ih > 0:
+                    ua = float(
+                        (boxes[n, 2] - boxes[n, 0] + 1) *
+                        (boxes[n, 3] - boxes[n, 1] + 1) +
+                        box_area - iw * ih
+                    )
+                    overlaps[n, k] = iw * ih / ua
+    return overlaps

trackers/multi_tracker_zoo.py

@@ -0,0 +1,52 @@
+from trackers.strongsort.utils.parser import get_config
+
+
+def create_tracker(tracker_type, tracker_config, reid_weights, device, half):
+    
+    cfg = get_config()
+    cfg.merge_from_file(tracker_config)
+    
+    if tracker_type == 'strongsort':
+        from trackers.strongsort.strong_sort import StrongSORT
+        strongsort = StrongSORT(
+            reid_weights,
+            device,
+            half,
+            max_dist=cfg.strongsort.max_dist,
+            max_iou_dist=cfg.strongsort.max_iou_dist,
+            max_age=cfg.strongsort.max_age,
+            max_unmatched_preds=cfg.strongsort.max_unmatched_preds,
+            n_init=cfg.strongsort.n_init,
+            nn_budget=cfg.strongsort.nn_budget,
+            mc_lambda=cfg.strongsort.mc_lambda,
+            ema_alpha=cfg.strongsort.ema_alpha,
+
+        )
+        return strongsort
+    
+    elif tracker_type == 'ocsort':
+        from trackers.ocsort.ocsort import OCSort
+        ocsort = OCSort(
+            det_thresh=cfg.ocsort.det_thresh,
+            max_age=cfg.ocsort.max_age,
+            min_hits=cfg.ocsort.min_hits,
+            iou_threshold=cfg.ocsort.iou_thresh,
+            delta_t=cfg.ocsort.delta_t,
+            asso_func=cfg.ocsort.asso_func,
+            inertia=cfg.ocsort.inertia,
+            use_byte=cfg.ocsort.use_byte,
+        )
+        return ocsort
+    
+    elif tracker_type == 'bytetrack':
+        from trackers.bytetrack.byte_tracker import BYTETracker
+        bytetracker = BYTETracker(
+            track_thresh=cfg.bytetrack.track_thresh,
+            match_thresh=cfg.bytetrack.match_thresh,
+            track_buffer=cfg.bytetrack.track_buffer,
+            frame_rate=cfg.bytetrack.frame_rate
+        )
+        return bytetracker
+    else:
+        print('No such tracker')
+        exit()

trackers/ocsort/association.py

@@ -0,0 +1,377 @@
+import os
+import numpy as np
+
+
+def iou_batch(bboxes1, bboxes2):
+    """
+    From SORT: Computes IOU between two bboxes in the form [x1,y1,x2,y2]
+    """
+    bboxes2 = np.expand_dims(bboxes2, 0)
+    bboxes1 = np.expand_dims(bboxes1, 1)
+    
+    xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])
+    yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])
+    xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])
+    yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])
+    w = np.maximum(0., xx2 - xx1)
+    h = np.maximum(0., yy2 - yy1)
+    wh = w * h
+    o = wh / ((bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1])                                      
+        + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - wh)                                              
+    return(o)  
+
+
+def giou_batch(bboxes1, bboxes2):
+    """
+    :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2)
+    :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2)
+    :return:
+    """
+    # for details should go to https://arxiv.org/pdf/1902.09630.pdf
+    # ensure predict's bbox form
+    bboxes2 = np.expand_dims(bboxes2, 0)
+    bboxes1 = np.expand_dims(bboxes1, 1)
+
+    xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])
+    yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])
+    xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])
+    yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])
+    w = np.maximum(0., xx2 - xx1)
+    h = np.maximum(0., yy2 - yy1)
+    wh = w * h
+    iou = wh / ((bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1])                                      
+        + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - wh)  
+
+    xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0])
+    yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])
+    xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2])
+    yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])
+    wc = xxc2 - xxc1 
+    hc = yyc2 - yyc1 
+    assert((wc > 0).all() and (hc > 0).all())
+    area_enclose = wc * hc 
+    giou = iou - (area_enclose - wh) / area_enclose
+    giou = (giou + 1.)/2.0 # resize from (-1,1) to (0,1)
+    return giou
+
+
+def diou_batch(bboxes1, bboxes2):
+    """
+    :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2)
+    :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2)
+    :return:
+    """
+    # for details should go to https://arxiv.org/pdf/1902.09630.pdf
+    # ensure predict's bbox form
+    bboxes2 = np.expand_dims(bboxes2, 0)
+    bboxes1 = np.expand_dims(bboxes1, 1)
+
+    # calculate the intersection box
+    xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])
+    yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])
+    xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])
+    yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])
+    w = np.maximum(0., xx2 - xx1)
+    h = np.maximum(0., yy2 - yy1)
+    wh = w * h
+    iou = wh / ((bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1])                                      
+        + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - wh) 
+
+    centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0
+    centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0
+    centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0
+    centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0
+
+    inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2
+
+    xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0])
+    yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])
+    xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2])
+    yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])
+
+    outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2
+    diou = iou - inner_diag / outer_diag
+
+    return (diou + 1) / 2.0 # resize from (-1,1) to (0,1)
+
+def ciou_batch(bboxes1, bboxes2):
+    """
+    :param bbox_p: predict of bbox(N,4)(x1,y1,x2,y2)
+    :param bbox_g: groundtruth of bbox(N,4)(x1,y1,x2,y2)
+    :return:
+    """
+    # for details should go to https://arxiv.org/pdf/1902.09630.pdf
+    # ensure predict's bbox form
+    bboxes2 = np.expand_dims(bboxes2, 0)
+    bboxes1 = np.expand_dims(bboxes1, 1)
+
+    # calculate the intersection box
+    xx1 = np.maximum(bboxes1[..., 0], bboxes2[..., 0])
+    yy1 = np.maximum(bboxes1[..., 1], bboxes2[..., 1])
+    xx2 = np.minimum(bboxes1[..., 2], bboxes2[..., 2])
+    yy2 = np.minimum(bboxes1[..., 3], bboxes2[..., 3])
+    w = np.maximum(0., xx2 - xx1)
+    h = np.maximum(0., yy2 - yy1)
+    wh = w * h
+    iou = wh / ((bboxes1[..., 2] - bboxes1[..., 0]) * (bboxes1[..., 3] - bboxes1[..., 1])                                      
+        + (bboxes2[..., 2] - bboxes2[..., 0]) * (bboxes2[..., 3] - bboxes2[..., 1]) - wh) 
+
+    centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0
+    centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0
+    centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0
+    centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0
+
+    inner_diag = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2
+
+    xxc1 = np.minimum(bboxes1[..., 0], bboxes2[..., 0])
+    yyc1 = np.minimum(bboxes1[..., 1], bboxes2[..., 1])
+    xxc2 = np.maximum(bboxes1[..., 2], bboxes2[..., 2])
+    yyc2 = np.maximum(bboxes1[..., 3], bboxes2[..., 3])
+
+    outer_diag = (xxc2 - xxc1) ** 2 + (yyc2 - yyc1) ** 2
+    
+    w1 = bboxes1[..., 2] - bboxes1[..., 0]
+    h1 = bboxes1[..., 3] - bboxes1[..., 1]
+    w2 = bboxes2[..., 2] - bboxes2[..., 0]
+    h2 = bboxes2[..., 3] - bboxes2[..., 1]
+
+    # prevent dividing over zero. add one pixel shift
+    h2 = h2 + 1.
+    h1 = h1 + 1.
+    arctan = np.arctan(w2/h2) - np.arctan(w1/h1)
+    v = (4 / (np.pi ** 2)) * (arctan ** 2)
+    S = 1 - iou 
+    alpha = v / (S+v)
+    ciou = iou - inner_diag / outer_diag - alpha * v
+    
+    return (ciou + 1) / 2.0 # resize from (-1,1) to (0,1)
+
+
+def ct_dist(bboxes1, bboxes2):
+    """
+        Measure the center distance between two sets of bounding boxes,
+        this is a coarse implementation, we don't recommend using it only
+        for association, which can be unstable and sensitive to frame rate
+        and object speed.
+    """
+    bboxes2 = np.expand_dims(bboxes2, 0)
+    bboxes1 = np.expand_dims(bboxes1, 1)
+
+    centerx1 = (bboxes1[..., 0] + bboxes1[..., 2]) / 2.0
+    centery1 = (bboxes1[..., 1] + bboxes1[..., 3]) / 2.0
+    centerx2 = (bboxes2[..., 0] + bboxes2[..., 2]) / 2.0
+    centery2 = (bboxes2[..., 1] + bboxes2[..., 3]) / 2.0
+
+    ct_dist2 = (centerx1 - centerx2) ** 2 + (centery1 - centery2) ** 2
+
+    ct_dist = np.sqrt(ct_dist2)
+
+    # The linear rescaling is a naive version and needs more study
+    ct_dist = ct_dist / ct_dist.max()
+    return ct_dist.max() - ct_dist # resize to (0,1)
+
+
+
+def speed_direction_batch(dets, tracks):
+    tracks = tracks[..., np.newaxis]
+    CX1, CY1 = (dets[:,0] + dets[:,2])/2.0, (dets[:,1]+dets[:,3])/2.0
+    CX2, CY2 = (tracks[:,0] + tracks[:,2]) /2.0, (tracks[:,1]+tracks[:,3])/2.0
+    dx = CX1 - CX2 
+    dy = CY1 - CY2 
+    norm = np.sqrt(dx**2 + dy**2) + 1e-6
+    dx = dx / norm 
+    dy = dy / norm
+    return dy, dx # size: num_track x num_det
+
+
+def linear_assignment(cost_matrix):
+    try:
+        import lap
+        _, x, y = lap.lapjv(cost_matrix, extend_cost=True)
+        return np.array([[y[i],i] for i in x if i >= 0]) #
+    except ImportError:
+        from scipy.optimize import linear_sum_assignment
+        x, y = linear_sum_assignment(cost_matrix)
+        return np.array(list(zip(x, y)))
+
+
+def associate_detections_to_trackers(detections,trackers, iou_threshold = 0.3):
+    """
+    Assigns detections to tracked object (both represented as bounding boxes)
+    Returns 3 lists of matches, unmatched_detections and unmatched_trackers
+    """
+    if(len(trackers)==0):
+        return np.empty((0,2),dtype=int), np.arange(len(detections)), np.empty((0,5),dtype=int)
+
+    iou_matrix = iou_batch(detections, trackers)
+
+    if min(iou_matrix.shape) > 0:
+        a = (iou_matrix > iou_threshold).astype(np.int32)
+        if a.sum(1).max() == 1 and a.sum(0).max() == 1:
+            matched_indices = np.stack(np.where(a), axis=1)
+        else:
+            matched_indices = linear_assignment(-iou_matrix)
+    else:
+        matched_indices = np.empty(shape=(0,2))
+
+    unmatched_detections = []
+    for d, det in enumerate(detections):
+        if(d not in matched_indices[:,0]):
+            unmatched_detections.append(d)
+    unmatched_trackers = []
+    for t, trk in enumerate(trackers):
+        if(t not in matched_indices[:,1]):
+            unmatched_trackers.append(t)
+
+    #filter out matched with low IOU
+    matches = []
+    for m in matched_indices:
+        if(iou_matrix[m[0], m[1]]<iou_threshold):
+            unmatched_detections.append(m[0])
+            unmatched_trackers.append(m[1])
+        else:
+            matches.append(m.reshape(1,2))
+    if(len(matches)==0):
+        matches = np.empty((0,2),dtype=int)
+    else:
+        matches = np.concatenate(matches,axis=0)
+
+    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
+
+
+def associate(detections, trackers, iou_threshold, velocities, previous_obs, vdc_weight):    
+    if(len(trackers)==0):
+        return np.empty((0,2),dtype=int), np.arange(len(detections)), np.empty((0,5),dtype=int)
+
+    Y, X = speed_direction_batch(detections, previous_obs)
+    inertia_Y, inertia_X = velocities[:,0], velocities[:,1]
+    inertia_Y = np.repeat(inertia_Y[:, np.newaxis], Y.shape[1], axis=1)
+    inertia_X = np.repeat(inertia_X[:, np.newaxis], X.shape[1], axis=1)
+    diff_angle_cos = inertia_X * X + inertia_Y * Y
+    diff_angle_cos = np.clip(diff_angle_cos, a_min=-1, a_max=1)
+    diff_angle = np.arccos(diff_angle_cos)
+    diff_angle = (np.pi /2.0 - np.abs(diff_angle)) / np.pi
+
+    valid_mask = np.ones(previous_obs.shape[0])
+    valid_mask[np.where(previous_obs[:,4]<0)] = 0
+    
+    iou_matrix = iou_batch(detections, trackers)
+    scores = np.repeat(detections[:,-1][:, np.newaxis], trackers.shape[0], axis=1)
+    # iou_matrix = iou_matrix * scores # a trick sometiems works, we don't encourage this
+    valid_mask = np.repeat(valid_mask[:, np.newaxis], X.shape[1], axis=1)
+
+    angle_diff_cost = (valid_mask * diff_angle) * vdc_weight
+    angle_diff_cost = angle_diff_cost.T
+    angle_diff_cost = angle_diff_cost * scores
+
+    if min(iou_matrix.shape) > 0:
+        a = (iou_matrix > iou_threshold).astype(np.int32)
+        if a.sum(1).max() == 1 and a.sum(0).max() == 1:
+            matched_indices = np.stack(np.where(a), axis=1)
+        else:
+            matched_indices = linear_assignment(-(iou_matrix+angle_diff_cost))
+    else:
+        matched_indices = np.empty(shape=(0,2))
+
+    unmatched_detections = []
+    for d, det in enumerate(detections):
+        if(d not in matched_indices[:,0]):
+            unmatched_detections.append(d)
+    unmatched_trackers = []
+    for t, trk in enumerate(trackers):
+        if(t not in matched_indices[:,1]):
+            unmatched_trackers.append(t)
+
+    # filter out matched with low IOU
+    matches = []
+    for m in matched_indices:
+        if(iou_matrix[m[0], m[1]]<iou_threshold):
+            unmatched_detections.append(m[0])
+            unmatched_trackers.append(m[1])
+        else:
+            matches.append(m.reshape(1,2))
+    if(len(matches)==0):
+        matches = np.empty((0,2),dtype=int)
+    else:
+        matches = np.concatenate(matches,axis=0)
+
+    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)
+
+
+def associate_kitti(detections, trackers, det_cates, iou_threshold, 
+        velocities, previous_obs, vdc_weight):
+    if(len(trackers)==0):
+        return np.empty((0,2),dtype=int), np.arange(len(detections)), np.empty((0,5),dtype=int)
+
+    """
+        Cost from the velocity direction consistency
+    """
+    Y, X = speed_direction_batch(detections, previous_obs)
+    inertia_Y, inertia_X = velocities[:,0], velocities[:,1]
+    inertia_Y = np.repeat(inertia_Y[:, np.newaxis], Y.shape[1], axis=1)
+    inertia_X = np.repeat(inertia_X[:, np.newaxis], X.shape[1], axis=1)
+    diff_angle_cos = inertia_X * X + inertia_Y * Y
+    diff_angle_cos = np.clip(diff_angle_cos, a_min=-1, a_max=1)
+    diff_angle = np.arccos(diff_angle_cos)
+    diff_angle = (np.pi /2.0 - np.abs(diff_angle)) / np.pi
+
+    valid_mask = np.ones(previous_obs.shape[0])
+    valid_mask[np.where(previous_obs[:,4]<0)]=0  
+    valid_mask = np.repeat(valid_mask[:, np.newaxis], X.shape[1], axis=1)
+
+    scores = np.repeat(detections[:,-1][:, np.newaxis], trackers.shape[0], axis=1)
+    angle_diff_cost = (valid_mask * diff_angle) * vdc_weight
+    angle_diff_cost = angle_diff_cost.T
+    angle_diff_cost = angle_diff_cost * scores
+
+    """
+        Cost from IoU
+    """
+    iou_matrix = iou_batch(detections, trackers)
+    
+
+    """
+        With multiple categories, generate the cost for catgory mismatch
+    """
+    num_dets = detections.shape[0]
+    num_trk = trackers.shape[0]
+    cate_matrix = np.zeros((num_dets, num_trk))
+    for i in range(num_dets):
+            for j in range(num_trk):
+                if det_cates[i] != trackers[j, 4]:
+                        cate_matrix[i][j] = -1e6
+    
+    cost_matrix = - iou_matrix -angle_diff_cost - cate_matrix
+
+    if min(iou_matrix.shape) > 0:
+        a = (iou_matrix > iou_threshold).astype(np.int32)
+        if a.sum(1).max() == 1 and a.sum(0).max() == 1:
+            matched_indices = np.stack(np.where(a), axis=1)
+        else:
+            matched_indices = linear_assignment(cost_matrix)
+    else:
+        matched_indices = np.empty(shape=(0,2))
+
+    unmatched_detections = []
+    for d, det in enumerate(detections):
+        if(d not in matched_indices[:,0]):
+            unmatched_detections.append(d)
+    unmatched_trackers = []
+    for t, trk in enumerate(trackers):
+        if(t not in matched_indices[:,1]):
+            unmatched_trackers.append(t)
+
+    #filter out matched with low IOU
+    matches = []
+    for m in matched_indices:
+        if(iou_matrix[m[0], m[1]]<iou_threshold):
+            unmatched_detections.append(m[0])
+            unmatched_trackers.append(m[1])
+        else:
+            matches.append(m.reshape(1,2))
+    if(len(matches)==0):
+        matches = np.empty((0,2),dtype=int)
+    else:
+        matches = np.concatenate(matches,axis=0)
+
+    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

trackers/ocsort/configs/ocsort.yaml

@@ -0,0 +1,12 @@
+# Trial number:      137
+# HOTA, MOTA, IDF1:  [55.567]
+ocsort:
+  asso_func: giou
+  conf_thres: 0.5122620708221085
+  delta_t: 1
+  det_thresh: 0
+  inertia: 0.3941737016672115
+  iou_thresh: 0.22136877277096445
+  max_age: 50
+  min_hits: 1
+  use_byte: false

trackers/ocsort/kalmanfilter.py

@@ -0,0 +1,1581 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=invalid-name, too-many-arguments, too-many-branches,
+# pylint: disable=too-many-locals, too-many-instance-attributes, too-many-lines
+
+"""
+This module implements the linear Kalman filter in both an object
+oriented and procedural form. The KalmanFilter class implements
+the filter by storing the various matrices in instance variables,
+minimizing the amount of bookkeeping you have to do.
+All Kalman filters operate with a predict->update cycle. The
+predict step, implemented with the method or function predict(),
+uses the state transition matrix F to predict the state in the next
+time period (epoch). The state is stored as a gaussian (x, P), where
+x is the state (column) vector, and P is its covariance. Covariance
+matrix Q specifies the process covariance. In Bayesian terms, this
+prediction is called the *prior*, which you can think of colloquially
+as the estimate prior to incorporating the measurement.
+The update step, implemented with the method or function `update()`,
+incorporates the measurement z with covariance R, into the state
+estimate (x, P). The class stores the system uncertainty in S,
+the innovation (residual between prediction and measurement in
+measurement space) in y, and the Kalman gain in k. The procedural
+form returns these variables to you. In Bayesian terms this computes
+the *posterior* - the estimate after the information from the
+measurement is incorporated.
+Whether you use the OO form or procedural form is up to you. If
+matrices such as H, R, and F are changing each epoch, you'll probably
+opt to use the procedural form. If they are unchanging, the OO
+form is perhaps easier to use since you won't need to keep track
+of these matrices. This is especially useful if you are implementing
+banks of filters or comparing various KF designs for performance;
+a trivial coding bug could lead to using the wrong sets of matrices.
+This module also offers an implementation of the RTS smoother, and
+other helper functions, such as log likelihood computations.
+The Saver class allows you to easily save the state of the
+KalmanFilter class after every update
+This module expects NumPy arrays for all values that expect
+arrays, although in a few cases, particularly method parameters,
+it will accept types that convert to NumPy arrays, such as lists
+of lists. These exceptions are documented in the method or function.
+Examples
+--------
+The following example constructs a constant velocity kinematic
+filter, filters noisy data, and plots the results. It also demonstrates
+using the Saver class to save the state of the filter at each epoch.
+.. code-block:: Python
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from filterpy.kalman import KalmanFilter
+    from filterpy.common import Q_discrete_white_noise, Saver
+    r_std, q_std = 2., 0.003
+    cv = KalmanFilter(dim_x=2, dim_z=1)
+    cv.x = np.array([[0., 1.]]) # position, velocity
+    cv.F = np.array([[1, dt],[ [0, 1]])
+    cv.R = np.array([[r_std^^2]])
+    f.H = np.array([[1., 0.]])
+    f.P = np.diag([.1^^2, .03^^2)
+    f.Q = Q_discrete_white_noise(2, dt, q_std**2)
+    saver = Saver(cv)
+    for z in range(100):
+        cv.predict()
+        cv.update([z + randn() * r_std])
+        saver.save() # save the filter's state
+    saver.to_array()
+    plt.plot(saver.x[:, 0])
+    # plot all of the priors
+    plt.plot(saver.x_prior[:, 0])
+    # plot mahalanobis distance
+    plt.figure()
+    plt.plot(saver.mahalanobis)
+This code implements the same filter using the procedural form
+    x = np.array([[0., 1.]]) # position, velocity
+    F = np.array([[1, dt],[ [0, 1]])
+    R = np.array([[r_std^^2]])
+    H = np.array([[1., 0.]])
+    P = np.diag([.1^^2, .03^^2)
+    Q = Q_discrete_white_noise(2, dt, q_std**2)
+    for z in range(100):
+        x, P = predict(x, P, F=F, Q=Q)
+        x, P = update(x, P, z=[z + randn() * r_std], R=R, H=H)
+        xs.append(x[0, 0])
+    plt.plot(xs)
+For more examples see the test subdirectory, or refer to the
+book cited below. In it I both teach Kalman filtering from basic
+principles, and teach the use of this library in great detail.
+FilterPy library.
+http://github.com/rlabbe/filterpy
+Documentation at:
+https://filterpy.readthedocs.org
+Supporting book at:
+https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python
+This is licensed under an MIT license. See the readme.MD file
+for more information.
+Copyright 2014-2018 Roger R Labbe Jr.
+"""
+
+from __future__ import absolute_import, division
+
+from copy import deepcopy
+from math import log, exp, sqrt
+import sys
+import numpy as np
+from numpy import dot, zeros, eye, isscalar, shape
+import numpy.linalg as linalg
+from filterpy.stats import logpdf
+from filterpy.common import pretty_str, reshape_z
+
+
+class KalmanFilterNew(object):
+    """ Implements a Kalman filter. You are responsible for setting the
+    various state variables to reasonable values; the defaults  will
+    not give you a functional filter.
+    For now the best documentation is my free book Kalman and Bayesian
+    Filters in Python [2]_. The test files in this directory also give you a
+    basic idea of use, albeit without much description.
+    In brief, you will first construct this object, specifying the size of
+    the state vector with dim_x and the size of the measurement vector that
+    you will be using with dim_z. These are mostly used to perform size checks
+    when you assign values to the various matrices. For example, if you
+    specified dim_z=2 and then try to assign a 3x3 matrix to R (the
+    measurement noise matrix you will get an assert exception because R
+    should be 2x2. (If for whatever reason you need to alter the size of
+    things midstream just use the underscore version of the matrices to
+    assign directly: your_filter._R = a_3x3_matrix.)
+    After construction the filter will have default matrices created for you,
+    but you must specify the values for each. It’s usually easiest to just
+    overwrite them rather than assign to each element yourself. This will be
+    clearer in the example below. All are of type numpy.array.
+    Examples
+    --------
+    Here is a filter that tracks position and velocity using a sensor that only
+    reads position.
+    First construct the object with the required dimensionality. Here the state
+    (`dim_x`) has 2 coefficients (position and velocity), and the measurement
+    (`dim_z`) has one. In FilterPy `x` is the state, `z` is the measurement.
+    .. code::
+        from filterpy.kalman import KalmanFilter
+        f = KalmanFilter (dim_x=2, dim_z=1)
+    Assign the initial value for the state (position and velocity). You can do this
+    with a two dimensional array like so:
+        .. code::
+            f.x = np.array([[2.],    # position
+                            [0.]])   # velocity
+    or just use a one dimensional array, which I prefer doing.
+    .. code::
+        f.x = np.array([2., 0.])
+    Define the state transition matrix:
+        .. code::
+            f.F = np.array([[1.,1.],
+                            [0.,1.]])
+    Define the measurement function. Here we need to convert a position-velocity
+    vector into just a position vector, so we use:
+        .. code::
+        f.H = np.array([[1., 0.]])
+    Define the state's covariance matrix P. 
+    .. code::
+        f.P = np.array([[1000.,    0.],
+                        [   0., 1000.] ])
+    Now assign the measurement noise. Here the dimension is 1x1, so I can
+    use a scalar
+    .. code::
+        f.R = 5
+    I could have done this instead:
+    .. code::
+        f.R = np.array([[5.]])
+    Note that this must be a 2 dimensional array.
+    Finally, I will assign the process noise. Here I will take advantage of
+    another FilterPy library function:
+    .. code::
+        from filterpy.common import Q_discrete_white_noise
+        f.Q = Q_discrete_white_noise(dim=2, dt=0.1, var=0.13)
+    Now just perform the standard predict/update loop:
+    .. code::
+        while some_condition_is_true:
+            z = get_sensor_reading()
+            f.predict()
+            f.update(z)
+            do_something_with_estimate (f.x)
+    **Procedural Form**
+    This module also contains stand alone functions to perform Kalman filtering.
+    Use these if you are not a fan of objects.
+    **Example**
+    .. code::
+        while True:
+            z, R = read_sensor()
+            x, P = predict(x, P, F, Q)
+            x, P = update(x, P, z, R, H)
+    See my book Kalman and Bayesian Filters in Python [2]_.
+    You will have to set the following attributes after constructing this
+    object for the filter to perform properly. Please note that there are
+    various checks in place to ensure that you have made everything the
+    'correct' size. However, it is possible to provide incorrectly sized
+    arrays such that the linear algebra can not perform an operation.
+    It can also fail silently - you can end up with matrices of a size that
+    allows the linear algebra to work, but are the wrong shape for the problem
+    you are trying to solve.
+    Parameters
+    ----------
+    dim_x : int
+        Number of state variables for the Kalman filter. For example, if
+        you are tracking the position and velocity of an object in two
+        dimensions, dim_x would be 4.
+        This is used to set the default size of P, Q, and u
+    dim_z : int
+        Number of of measurement inputs. For example, if the sensor
+        provides you with position in (x,y), dim_z would be 2.
+    dim_u : int (optional)
+        size of the control input, if it is being used.
+        Default value of 0 indicates it is not used.
+    compute_log_likelihood : bool (default = True)
+        Computes log likelihood by default, but this can be a slow
+        computation, so if you never use it you can turn this computation
+        off.
+    Attributes
+    ----------
+    x : numpy.array(dim_x, 1)
+        Current state estimate. Any call to update() or predict() updates
+        this variable.
+    P : numpy.array(dim_x, dim_x)
+        Current state covariance matrix. Any call to update() or predict()
+        updates this variable.
+    x_prior : numpy.array(dim_x, 1)
+        Prior (predicted) state estimate. The *_prior and *_post attributes
+        are for convenience; they store the  prior and posterior of the
+        current epoch. Read Only.
+    P_prior : numpy.array(dim_x, dim_x)
+        Prior (predicted) state covariance matrix. Read Only.
+    x_post : numpy.array(dim_x, 1)
+        Posterior (updated) state estimate. Read Only.
+    P_post : numpy.array(dim_x, dim_x)
+        Posterior (updated) state covariance matrix. Read Only.
+    z : numpy.array
+        Last measurement used in update(). Read only.
+    R : numpy.array(dim_z, dim_z)
+        Measurement noise covariance matrix. Also known as the
+        observation covariance.
+    Q : numpy.array(dim_x, dim_x)
+        Process noise covariance matrix. Also known as the transition
+        covariance.
+    F : numpy.array()
+        State Transition matrix. Also known as `A` in some formulation.
+    H : numpy.array(dim_z, dim_x)
+        Measurement function. Also known as the observation matrix, or as `C`.
+    y : numpy.array
+        Residual of the update step. Read only.
+    K : numpy.array(dim_x, dim_z)
+        Kalman gain of the update step. Read only.
+    S :  numpy.array
+        System uncertainty (P projected to measurement space). Read only.
+    SI :  numpy.array
+        Inverse system uncertainty. Read only.
+    log_likelihood : float
+        log-likelihood of the last measurement. Read only.
+    likelihood : float
+        likelihood of last measurement. Read only.
+        Computed from the log-likelihood. The log-likelihood can be very
+        small,  meaning a large negative value such as -28000. Taking the
+        exp() of that results in 0.0, which can break typical algorithms
+        which multiply by this value, so by default we always return a
+        number >= sys.float_info.min.
+    mahalanobis : float
+        mahalanobis distance of the innovation. Read only.
+    inv : function, default numpy.linalg.inv
+        If you prefer another inverse function, such as the Moore-Penrose
+        pseudo inverse, set it to that instead: kf.inv = np.linalg.pinv
+        This is only used to invert self.S. If you know it is diagonal, you
+        might choose to set it to filterpy.common.inv_diagonal, which is
+        several times faster than numpy.linalg.inv for diagonal matrices.
+    alpha : float
+        Fading memory setting. 1.0 gives the normal Kalman filter, and
+        values slightly larger than 1.0 (such as 1.02) give a fading
+        memory effect - previous measurements have less influence on the
+        filter's estimates. This formulation of the Fading memory filter
+        (there are many) is due to Dan Simon [1]_.
+    References
+    ----------
+    .. [1] Dan Simon. "Optimal State Estimation." John Wiley & Sons.
+       p. 208-212. (2006)
+    .. [2] Roger Labbe. "Kalman and Bayesian Filters in Python"
+       https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python
+    """
+
+    def __init__(self, dim_x, dim_z, dim_u=0):
+        if dim_x < 1:
+            raise ValueError('dim_x must be 1 or greater')
+        if dim_z < 1:
+            raise ValueError('dim_z must be 1 or greater')
+        if dim_u < 0:
+            raise ValueError('dim_u must be 0 or greater')
+
+        self.dim_x = dim_x
+        self.dim_z = dim_z
+        self.dim_u = dim_u
+
+        self.x = zeros((dim_x, 1))        # state
+        self.P = eye(dim_x)               # uncertainty covariance
+        self.Q = eye(dim_x)               # process uncertainty
+        self.B = None                     # control transition matrix
+        self.F = eye(dim_x)               # state transition matrix
+        self.H = zeros((dim_z, dim_x))    # measurement function
+        self.R = eye(dim_z)               # measurement uncertainty
+        self._alpha_sq = 1.               # fading memory control
+        self.M = np.zeros((dim_x, dim_z)) # process-measurement cross correlation
+        self.z = np.array([[None]*self.dim_z]).T
+
+        # gain and residual are computed during the innovation step. We
+        # save them so that in case you want to inspect them for various
+        # purposes
+        self.K = np.zeros((dim_x, dim_z)) # kalman gain
+        self.y = zeros((dim_z, 1))
+        self.S = np.zeros((dim_z, dim_z)) # system uncertainty
+        self.SI = np.zeros((dim_z, dim_z)) # inverse system uncertainty
+
+        # identity matrix. Do not alter this.
+        self._I = np.eye(dim_x)
+
+        # these will always be a copy of x,P after predict() is called
+        self.x_prior = self.x.copy()
+        self.P_prior = self.P.copy()
+
+        # these will always be a copy of x,P after update() is called
+        self.x_post = self.x.copy()             
+        self.P_post = self.P.copy()
+
+        # Only computed only if requested via property
+        self._log_likelihood = log(sys.float_info.min)
+        self._likelihood = sys.float_info.min
+        self._mahalanobis = None
+
+        # keep all observations 
+        self.history_obs = []
+
+        self.inv = np.linalg.inv
+
+        self.attr_saved = None
+        self.observed = False 
+
+
+    def predict(self, u=None, B=None, F=None, Q=None):
+        """
+        Predict next state (prior) using the Kalman filter state propagation
+        equations.
+        Parameters
+        ----------
+        u : np.array, default 0
+            Optional control vector.
+        B : np.array(dim_x, dim_u), or None
+            Optional control transition matrix; a value of None
+            will cause the filter to use `self.B`.
+        F : np.array(dim_x, dim_x), or None
+            Optional state transition matrix; a value of None
+            will cause the filter to use `self.F`.
+        Q : np.array(dim_x, dim_x), scalar, or None
+            Optional process noise matrix; a value of None will cause the
+            filter to use `self.Q`.
+        """
+
+        if B is None:
+            B = self.B
+        if F is None:
+            F = self.F
+        if Q is None:
+            Q = self.Q
+        elif isscalar(Q):
+            Q = eye(self.dim_x) * Q
+
+
+        # x = Fx + Bu
+        if B is not None and u is not None:
+            self.x = dot(F, self.x) + dot(B, u)
+        else:
+            self.x = dot(F, self.x)
+
+        # P = FPF' + Q
+        self.P = self._alpha_sq * dot(dot(F, self.P), F.T) + Q
+
+        # save prior
+        self.x_prior = self.x.copy()
+        self.P_prior = self.P.copy()
+
+
+
+    def freeze(self):
+        """
+            Save the parameters before non-observation forward
+        """
+        self.attr_saved = deepcopy(self.__dict__)
+
+
+    def unfreeze(self):
+        if self.attr_saved is not None:
+            new_history = deepcopy(self.history_obs)
+            self.__dict__ = self.attr_saved
+            # self.history_obs = new_history 
+            self.history_obs = self.history_obs[:-1]
+            occur = [int(d is None) for d in new_history]
+            indices = np.where(np.array(occur)==0)[0]
+            index1 = indices[-2]
+            index2 = indices[-1]
+            box1 = new_history[index1]
+            x1, y1, s1, r1 = box1 
+            w1 = np.sqrt(s1 * r1)
+            h1 = np.sqrt(s1 / r1)
+            box2 = new_history[index2]
+            x2, y2, s2, r2 = box2 
+            w2 = np.sqrt(s2 * r2)
+            h2 = np.sqrt(s2 / r2)
+            time_gap = index2 - index1
+            dx = (x2-x1)/time_gap
+            dy = (y2-y1)/time_gap 
+            dw = (w2-w1)/time_gap 
+            dh = (h2-h1)/time_gap
+            for i in range(index2 - index1):
+                """
+                    The default virtual trajectory generation is by linear
+                    motion (constant speed hypothesis), you could modify this 
+                    part to implement your own. 
+                """
+                x = x1 + (i+1) * dx 
+                y = y1 + (i+1) * dy 
+                w = w1 + (i+1) * dw 
+                h = h1 + (i+1) * dh
+                s = w * h 
+                r = w / float(h)
+                new_box = np.array([x, y, s, r]).reshape((4, 1))
+                """
+                    I still use predict-update loop here to refresh the parameters,
+                    but this can be faster by directly modifying the internal parameters
+                    as suggested in the paper. I keep this naive but slow way for 
+                    easy read and understanding
+                """
+                self.update(new_box)
+                if not i == (index2-index1-1):
+                    self.predict()
+
+
+    def update(self, z, R=None, H=None):
+        """
+        Add a new measurement (z) to the Kalman filter.
+        If z is None, nothing is computed. However, x_post and P_post are
+        updated with the prior (x_prior, P_prior), and self.z is set to None.
+        Parameters
+        ----------
+        z : (dim_z, 1): array_like
+            measurement for this update. z can be a scalar if dim_z is 1,
+            otherwise it must be convertible to a column vector.
+            If you pass in a value of H, z must be a column vector the
+            of the correct size.
+        R : np.array, scalar, or None
+            Optionally provide R to override the measurement noise for this
+            one call, otherwise  self.R will be used.
+        H : np.array, or None
+            Optionally provide H to override the measurement function for this
+            one call, otherwise self.H will be used.
+        """
+
+        # set to None to force recompute
+        self._log_likelihood = None
+        self._likelihood = None
+        self._mahalanobis = None
+
+        # append the observation
+        self.history_obs.append(z)
+        
+        if z is None:
+            if self.observed:
+                """
+                    Got no observation so freeze the current parameters for future
+                    potential online smoothing.
+                """
+                self.freeze()
+            self.observed = False 
+            self.z = np.array([[None]*self.dim_z]).T
+            self.x_post = self.x.copy()
+            self.P_post = self.P.copy()
+            self.y = zeros((self.dim_z, 1))
+            return
+        
+        # self.observed = True
+        if not self.observed:
+            """
+                Get observation, use online smoothing to re-update parameters
+            """
+            self.unfreeze()
+        self.observed = True
+
+        if R is None:
+            R = self.R
+        elif isscalar(R):
+            R = eye(self.dim_z) * R
+
+        if H is None:
+            z = reshape_z(z, self.dim_z, self.x.ndim)
+            H = self.H
+
+        # y = z - Hx
+        # error (residual) between measurement and prediction
+        self.y = z - dot(H, self.x)
+
+        # common subexpression for speed
+        PHT = dot(self.P, H.T)
+
+        # S = HPH' + R
+        # project system uncertainty into measurement space
+        self.S = dot(H, PHT) + R
+        self.SI = self.inv(self.S)
+        # K = PH'inv(S)
+        # map system uncertainty into kalman gain
+        self.K = dot(PHT, self.SI)
+
+        # x = x + Ky
+        # predict new x with residual scaled by the kalman gain
+        self.x = self.x + dot(self.K, self.y)
+
+        # P = (I-KH)P(I-KH)' + KRK'
+        # This is more numerically stable
+        # and works for non-optimal K vs the equation
+        # P = (I-KH)P usually seen in the literature.
+
+        I_KH = self._I - dot(self.K, H)
+        self.P = dot(dot(I_KH, self.P), I_KH.T) + dot(dot(self.K, R), self.K.T)
+
+        # save measurement and posterior state
+        self.z = deepcopy(z)
+        self.x_post = self.x.copy()
+        self.P_post = self.P.copy()
+
+    def predict_steadystate(self, u=0, B=None):
+        """
+        Predict state (prior) using the Kalman filter state propagation
+        equations. Only x is updated, P is left unchanged. See
+        update_steadstate() for a longer explanation of when to use this
+        method.
+        Parameters
+        ----------
+        u : np.array
+            Optional control vector. If non-zero, it is multiplied by B
+            to create the control input into the system.
+        B : np.array(dim_x, dim_u), or None
+            Optional control transition matrix; a value of None
+            will cause the filter to use `self.B`.
+        """
+
+        if B is None:
+            B = self.B
+
+        # x = Fx + Bu
+        if B is not None:
+            self.x = dot(self.F, self.x) + dot(B, u)
+        else:
+            self.x = dot(self.F, self.x)
+
+        # save prior
+        self.x_prior = self.x.copy()
+        self.P_prior = self.P.copy()
+
+    def update_steadystate(self, z):
+        """
+        Add a new measurement (z) to the Kalman filter without recomputing
+        the Kalman gain K, the state covariance P, or the system
+        uncertainty S.
+        You can use this for LTI systems since the Kalman gain and covariance
+        converge to a fixed value. Precompute these and assign them explicitly,
+        or run the Kalman filter using the normal predict()/update(0 cycle
+        until they converge.
+        The main advantage of this call is speed. We do significantly less
+        computation, notably avoiding a costly matrix inversion.
+        Use in conjunction with predict_steadystate(), otherwise P will grow
+        without bound.
+        Parameters
+        ----------
+        z : (dim_z, 1): array_like
+            measurement for this update. z can be a scalar if dim_z is 1,
+            otherwise it must be convertible to a column vector.
+        Examples
+        --------
+        >>> cv = kinematic_kf(dim=3, order=2) # 3D const velocity filter
+        >>> # let filter converge on representative data, then save k and P
+        >>> for i in range(100):
+        >>>     cv.predict()
+        >>>     cv.update([i, i, i])
+        >>> saved_k = np.copy(cv.K)
+        >>> saved_P = np.copy(cv.P)
+        later on:
+        >>> cv = kinematic_kf(dim=3, order=2) # 3D const velocity filter
+        >>> cv.K = np.copy(saved_K)
+        >>> cv.P = np.copy(saved_P)
+        >>> for i in range(100):
+        >>>     cv.predict_steadystate()
+        >>>     cv.update_steadystate([i, i, i])
+        """
+
+        # set to None to force recompute
+        self._log_likelihood = None
+        self._likelihood = None
+        self._mahalanobis = None
+
+        if z is None:
+            self.z = np.array([[None]*self.dim_z]).T
+            self.x_post = self.x.copy()
+            self.P_post = self.P.copy()
+            self.y = zeros((self.dim_z, 1))
+            return
+
+        z = reshape_z(z, self.dim_z, self.x.ndim)
+
+        # y = z - Hx
+        # error (residual) between measurement and prediction
+        self.y = z - dot(self.H, self.x)
+
+        # x = x + Ky
+        # predict new x with residual scaled by the kalman gain
+        self.x = self.x + dot(self.K, self.y)
+
+        self.z = deepcopy(z)
+        self.x_post = self.x.copy()
+        self.P_post = self.P.copy()
+
+        # set to None to force recompute
+        self._log_likelihood = None
+        self._likelihood = None
+        self._mahalanobis = None
+
+    def update_correlated(self, z, R=None, H=None):
+        """ Add a new measurement (z) to the Kalman filter assuming that
+        process noise and measurement noise are correlated as defined in
+        the `self.M` matrix.
+        A partial derivation can be found in [1]
+        If z is None, nothing is changed.
+        Parameters
+        ----------
+        z : (dim_z, 1): array_like
+            measurement for this update. z can be a scalar if dim_z is 1,
+            otherwise it must be convertible to a column vector.
+        R : np.array, scalar, or None
+            Optionally provide R to override the measurement noise for this
+            one call, otherwise  self.R will be used.
+        H : np.array,  or None
+            Optionally provide H to override the measurement function for this
+            one call, otherwise  self.H will be used.
+        References
+        ----------
+        .. [1] Bulut, Y. (2011). Applied Kalman filter theory (Doctoral dissertation, Northeastern University).
+               http://people.duke.edu/~hpgavin/SystemID/References/Balut-KalmanFilter-PhD-NEU-2011.pdf
+        """
+
+        # set to None to force recompute
+        self._log_likelihood = None
+        self._likelihood = None
+        self._mahalanobis = None
+
+        if z is None:
+            self.z = np.array([[None]*self.dim_z]).T
+            self.x_post = self.x.copy()
+            self.P_post = self.P.copy()
+            self.y = zeros((self.dim_z, 1))
+            return
+
+        if R is None:
+            R = self.R
+        elif isscalar(R):
+            R = eye(self.dim_z) * R
+
+        # rename for readability and a tiny extra bit of speed
+        if H is None:
+            z = reshape_z(z, self.dim_z, self.x.ndim)
+            H = self.H
+
+        # handle special case: if z is in form [[z]] but x is not a column
+        # vector dimensions will not match
+        if self.x.ndim == 1 and shape(z) == (1, 1):
+            z = z[0]
+
+        if shape(z) == (): # is it scalar, e.g. z=3 or z=np.array(3)
+            z = np.asarray([z])
+
+        # y = z - Hx
+        # error (residual) between measurement and prediction
+        self.y = z - dot(H, self.x)
+
+        # common subexpression for speed
+        PHT = dot(self.P, H.T)
+
+        # project system uncertainty into measurement space
+        self.S = dot(H, PHT) + dot(H, self.M) + dot(self.M.T, H.T) + R
+        self.SI = self.inv(self.S)
+
+        # K = PH'inv(S)
+        # map system uncertainty into kalman gain
+        self.K = dot(PHT + self.M, self.SI)
+
+        # x = x + Ky
+        # predict new x with residual scaled by the kalman gain
+        self.x = self.x + dot(self.K, self.y)
+        self.P = self.P - dot(self.K, dot(H, self.P) + self.M.T)
+
+        self.z = deepcopy(z)
+        self.x_post = self.x.copy()
+        self.P_post = self.P.copy()
+
+    def batch_filter(self, zs, Fs=None, Qs=None, Hs=None,
+                     Rs=None, Bs=None, us=None, update_first=False,
+                     saver=None):
+        """ Batch processes a sequences of measurements.
+        Parameters
+        ----------
+        zs : list-like
+            list of measurements at each time step `self.dt`. Missing
+            measurements must be represented by `None`.
+        Fs : None, list-like, default=None
+            optional value or list of values to use for the state transition
+            matrix F.
+            If Fs is None then self.F is used for all epochs.
+            Otherwise it must contain a list-like list of F's, one for
+            each epoch.  This allows you to have varying F per epoch.
+        Qs : None, np.array or list-like, default=None
+            optional value or list of values to use for the process error
+            covariance Q.
+            If Qs is None then self.Q is used for all epochs.
+            Otherwise it must contain a list-like list of Q's, one for
+            each epoch.  This allows you to have varying Q per epoch.
+        Hs : None, np.array or list-like, default=None
+            optional list of values to use for the measurement matrix H.
+            If Hs is None then self.H is used for all epochs.
+            If Hs contains a single matrix, then it is used as H for all
+            epochs.
+            Otherwise it must contain a list-like list of H's, one for
+            each epoch.  This allows you to have varying H per epoch.
+        Rs : None, np.array or list-like, default=None
+            optional list of values to use for the measurement error
+            covariance R.
+            If Rs is None then self.R is used for all epochs.
+            Otherwise it must contain a list-like list of R's, one for
+            each epoch.  This allows you to have varying R per epoch.
+        Bs : None, np.array or list-like, default=None
+            optional list of values to use for the control transition matrix B.
+            If Bs is None then self.B is used for all epochs.
+            Otherwise it must contain a list-like list of B's, one for
+            each epoch.  This allows you to have varying B per epoch.
+        us : None, np.array or list-like, default=None
+            optional list of values to use for the control input vector;
+            If us is None then None is used for all epochs (equivalent to 0,
+            or no control input).
+            Otherwise it must contain a list-like list of u's, one for
+            each epoch.
+       update_first : bool, optional, default=False
+            controls whether the order of operations is update followed by
+            predict, or predict followed by update. Default is predict->update.
+        saver : filterpy.common.Saver, optional
+            filterpy.common.Saver object. If provided, saver.save() will be
+            called after every epoch
+        Returns
+        -------
+        means : np.array((n,dim_x,1))
+            array of the state for each time step after the update. Each entry
+            is an np.array. In other words `means[k,:]` is the state at step
+            `k`.
+        covariance : np.array((n,dim_x,dim_x))
+            array of the covariances for each time step after the update.
+            In other words `covariance[k,:,:]` is the covariance at step `k`.
+        means_predictions : np.array((n,dim_x,1))
+            array of the state for each time step after the predictions. Each
+            entry is an np.array. In other words `means[k,:]` is the state at
+            step `k`.
+        covariance_predictions : np.array((n,dim_x,dim_x))
+            array of the covariances for each time step after the prediction.
+            In other words `covariance[k,:,:]` is the covariance at step `k`.
+        Examples
+        --------
+        .. code-block:: Python
+            # this example demonstrates tracking a measurement where the time
+            # between measurement varies, as stored in dts. This requires
+            # that F be recomputed for each epoch. The output is then smoothed
+            # with an RTS smoother.
+            zs = [t + random.randn()*4 for t in range (40)]
+            Fs = [np.array([[1., dt], [0, 1]] for dt in dts]
+            (mu, cov, _, _) = kf.batch_filter(zs, Fs=Fs)
+            (xs, Ps, Ks, Pps) = kf.rts_smoother(mu, cov, Fs=Fs)
+        """
+
+        #pylint: disable=too-many-statements
+        n = np.size(zs, 0)
+        if Fs is None:
+            Fs = [self.F] * n
+        if Qs is None:
+            Qs = [self.Q] * n
+        if Hs is None:
+            Hs = [self.H] * n
+        if Rs is None:
+            Rs = [self.R] * n
+        if Bs is None:
+            Bs = [self.B] * n
+        if us is None:
+            us = [0] * n
+
+        # mean estimates from Kalman Filter
+        if self.x.ndim == 1:
+            means = zeros((n, self.dim_x))
+            means_p = zeros((n, self.dim_x))
+        else:
+            means = zeros((n, self.dim_x, 1))
+            means_p = zeros((n, self.dim_x, 1))
+
+        # state covariances from Kalman Filter
+        covariances = zeros((n, self.dim_x, self.dim_x))
+        covariances_p = zeros((n, self.dim_x, self.dim_x))
+
+        if update_first:
+            for i, (z, F, Q, H, R, B, u) in enumerate(zip(zs, Fs, Qs, Hs, Rs, Bs, us)):
+
+                self.update(z, R=R, H=H)
+                means[i, :] = self.x
+                covariances[i, :, :] = self.P
+
+                self.predict(u=u, B=B, F=F, Q=Q)
+                means_p[i, :] = self.x
+                covariances_p[i, :, :] = self.P
+
+                if saver is not None:
+                    saver.save()
+        else:
+            for i, (z, F, Q, H, R, B, u) in enumerate(zip(zs, Fs, Qs, Hs, Rs, Bs, us)):
+
+                self.predict(u=u, B=B, F=F, Q=Q)
+                means_p[i, :] = self.x
+                covariances_p[i, :, :] = self.P
+
+                self.update(z, R=R, H=H)
+                means[i, :] = self.x
+                covariances[i, :, :] = self.P
+
+                if saver is not None:
+                    saver.save()
+
+        return (means, covariances, means_p, covariances_p)
+
+    def rts_smoother(self, Xs, Ps, Fs=None, Qs=None, inv=np.linalg.inv):
+        """
+        Runs the Rauch-Tung-Striebel Kalman smoother on a set of
+        means and covariances computed by a Kalman filter. The usual input
+        would come from the output of `KalmanFilter.batch_filter()`.
+        Parameters
+        ----------
+        Xs : numpy.array
+           array of the means (state variable x) of the output of a Kalman
+           filter.
+        Ps : numpy.array
+            array of the covariances of the output of a kalman filter.
+        Fs : list-like collection of numpy.array, optional
+            State transition matrix of the Kalman filter at each time step.
+            Optional, if not provided the filter's self.F will be used
+        Qs : list-like collection of numpy.array, optional
+            Process noise of the Kalman filter at each time step. Optional,
+            if not provided the filter's self.Q will be used
+        inv : function, default numpy.linalg.inv
+            If you prefer another inverse function, such as the Moore-Penrose
+            pseudo inverse, set it to that instead: kf.inv = np.linalg.pinv
+        Returns
+        -------
+        x : numpy.ndarray
+           smoothed means
+        P : numpy.ndarray
+           smoothed state covariances
+        K : numpy.ndarray
+            smoother gain at each step
+        Pp : numpy.ndarray
+           Predicted state covariances
+        Examples
+        --------
+        .. code-block:: Python
+            zs = [t + random.randn()*4 for t in range (40)]
+            (mu, cov, _, _) = kalman.batch_filter(zs)
+            (x, P, K, Pp) = rts_smoother(mu, cov, kf.F, kf.Q)
+        """
+
+        if len(Xs) != len(Ps):
+            raise ValueError('length of Xs and Ps must be the same')
+
+        n = Xs.shape[0]
+        dim_x = Xs.shape[1]
+
+        if Fs is None:
+            Fs = [self.F] * n
+        if Qs is None:
+            Qs = [self.Q] * n
+
+        # smoother gain
+        K = zeros((n, dim_x, dim_x))
+
+        x, P, Pp = Xs.copy(), Ps.copy(), Ps.copy()
+        for k in range(n-2, -1, -1):
+            Pp[k] = dot(dot(Fs[k+1], P[k]), Fs[k+1].T) + Qs[k+1]
+
+            #pylint: disable=bad-whitespace
+            K[k]  = dot(dot(P[k], Fs[k+1].T), inv(Pp[k]))
+            x[k] += dot(K[k], x[k+1] - dot(Fs[k+1], x[k]))
+            P[k] += dot(dot(K[k], P[k+1] - Pp[k]), K[k].T)
+
+        return (x, P, K, Pp)
+
+    def get_prediction(self, u=None, B=None, F=None, Q=None):
+        """
+        Predict next state (prior) using the Kalman filter state propagation
+        equations and returns it without modifying the object.
+        Parameters
+        ----------
+        u : np.array, default 0
+            Optional control vector.
+        B : np.array(dim_x, dim_u), or None
+            Optional control transition matrix; a value of None
+            will cause the filter to use `self.B`.
+        F : np.array(dim_x, dim_x), or None
+            Optional state transition matrix; a value of None
+            will cause the filter to use `self.F`.
+        Q : np.array(dim_x, dim_x), scalar, or None
+            Optional process noise matrix; a value of None will cause the
+            filter to use `self.Q`.
+        Returns
+        -------
+        (x, P) : tuple
+            State vector and covariance array of the prediction.
+        """
+
+        if B is None:
+            B = self.B
+        if F is None:
+            F = self.F
+        if Q is None:
+            Q = self.Q
+        elif isscalar(Q):
+            Q = eye(self.dim_x) * Q
+
+        # x = Fx + Bu
+        if B is not None and u is not None:
+            x = dot(F, self.x) + dot(B, u)
+        else:
+            x = dot(F, self.x)
+
+        # P = FPF' + Q
+        P = self._alpha_sq * dot(dot(F, self.P), F.T) + Q
+
+        return x, P
+
+    def get_update(self, z=None):
+        """
+        Computes the new estimate based on measurement `z` and returns it
+        without altering the state of the filter.
+        Parameters
+        ----------
+        z : (dim_z, 1): array_like
+            measurement for this update. z can be a scalar if dim_z is 1,
+            otherwise it must be convertible to a column vector.
+        Returns
+        -------
+        (x, P) : tuple
+            State vector and covariance array of the update.
+       """
+
+        if z is None:
+            return self.x, self.P
+        z = reshape_z(z, self.dim_z, self.x.ndim)
+
+        R = self.R
+        H = self.H
+        P = self.P
+        x = self.x
+
+        # error (residual) between measurement and prediction
+        y = z - dot(H, x)
+
+        # common subexpression for speed
+        PHT = dot(P, H.T)
+
+        # project system uncertainty into measurement space
+        S = dot(H, PHT) + R
+
+        # map system uncertainty into kalman gain
+        K = dot(PHT, self.inv(S))
+
+        # predict new x with residual scaled by the kalman gain
+        x = x + dot(K, y)
+
+        # P = (I-KH)P(I-KH)' + KRK'
+        I_KH = self._I - dot(K, H)
+        P = dot(dot(I_KH, P), I_KH.T) + dot(dot(K, R), K.T)
+
+        return x, P
+
+    def residual_of(self, z):
+        """
+        Returns the residual for the given measurement (z). Does not alter
+        the state of the filter.
+        """
+        z = reshape_z(z, self.dim_z, self.x.ndim)
+        return z - dot(self.H, self.x_prior)
+
+    def measurement_of_state(self, x):
+        """
+        Helper function that converts a state into a measurement.
+        Parameters
+        ----------
+        x : np.array
+            kalman state vector
+        Returns
+        -------
+        z : (dim_z, 1): array_like
+            measurement for this update. z can be a scalar if dim_z is 1,
+            otherwise it must be convertible to a column vector.
+        """
+
+        return dot(self.H, x)
+
+    @property
+    def log_likelihood(self):
+        """
+        log-likelihood of the last measurement.
+        """
+        if self._log_likelihood is None:
+            self._log_likelihood = logpdf(x=self.y, cov=self.S)
+        return self._log_likelihood
+
+    @property
+    def likelihood(self):
+        """
+        Computed from the log-likelihood. The log-likelihood can be very
+        small,  meaning a large negative value such as -28000. Taking the
+        exp() of that results in 0.0, which can break typical algorithms
+        which multiply by this value, so by default we always return a
+        number >= sys.float_info.min.
+        """
+        if self._likelihood is None:
+            self._likelihood = exp(self.log_likelihood)
+            if self._likelihood == 0:
+                self._likelihood = sys.float_info.min
+        return self._likelihood
+
+    @property
+    def mahalanobis(self):
+        """"
+        Mahalanobis distance of measurement. E.g. 3 means measurement
+        was 3 standard deviations away from the predicted value.
+        Returns
+        -------
+        mahalanobis : float
+        """
+        if self._mahalanobis is None:
+            self._mahalanobis = sqrt(float(dot(dot(self.y.T, self.SI), self.y)))
+        return self._mahalanobis
+
+    @property
+    def alpha(self):
+        """
+        Fading memory setting. 1.0 gives the normal Kalman filter, and
+        values slightly larger than 1.0 (such as 1.02) give a fading
+        memory effect - previous measurements have less influence on the
+        filter's estimates. This formulation of the Fading memory filter
+        (there are many) is due to Dan Simon [1]_.
+        """
+        return self._alpha_sq**.5
+
+    def log_likelihood_of(self, z):
+        """
+        log likelihood of the measurement `z`. This should only be called
+        after a call to update(). Calling after predict() will yield an
+        incorrect result."""
+
+        if z is None:
+            return log(sys.float_info.min)
+        return logpdf(z, dot(self.H, self.x), self.S)
+
+    @alpha.setter
+    def alpha(self, value):
+        if not np.isscalar(value) or value < 1:
+            raise ValueError('alpha must be a float greater than 1')
+
+        self._alpha_sq = value**2
+
+    def __repr__(self):
+        return '\n'.join([
+            'KalmanFilter object',
+            pretty_str('dim_x', self.dim_x),
+            pretty_str('dim_z', self.dim_z),
+            pretty_str('dim_u', self.dim_u),
+            pretty_str('x', self.x),
+            pretty_str('P', self.P),
+            pretty_str('x_prior', self.x_prior),
+            pretty_str('P_prior', self.P_prior),
+            pretty_str('x_post', self.x_post),
+            pretty_str('P_post', self.P_post),
+            pretty_str('F', self.F),
+            pretty_str('Q', self.Q),
+            pretty_str('R', self.R),
+            pretty_str('H', self.H),
+            pretty_str('K', self.K),
+            pretty_str('y', self.y),
+            pretty_str('S', self.S),
+            pretty_str('SI', self.SI),
+            pretty_str('M', self.M),
+            pretty_str('B', self.B),
+            pretty_str('z', self.z),
+            pretty_str('log-likelihood', self.log_likelihood),
+            pretty_str('likelihood', self.likelihood),
+            pretty_str('mahalanobis', self.mahalanobis),
+            pretty_str('alpha', self.alpha),
+            pretty_str('inv', self.inv)
+            ])
+
+    def test_matrix_dimensions(self, z=None, H=None, R=None, F=None, Q=None):
+        """
+        Performs a series of asserts to check that the size of everything
+        is what it should be. This can help you debug problems in your design.
+        If you pass in H, R, F, Q those will be used instead of this object's
+        value for those matrices.
+        Testing `z` (the measurement) is problamatic. x is a vector, and can be
+        implemented as either a 1D array or as a nx1 column vector. Thus Hx
+        can be of different shapes. Then, if Hx is a single value, it can
+        be either a 1D array or 2D vector. If either is true, z can reasonably
+        be a scalar (either '3' or np.array('3') are scalars under this
+        definition), a 1D, 1 element array, or a 2D, 1 element array. You are
+        allowed to pass in any combination that works.
+        """
+
+        if H is None:
+            H = self.H
+        if R is None:
+            R = self.R
+        if F is None:
+            F = self.F
+        if Q is None:
+            Q = self.Q
+        x = self.x
+        P = self.P
+
+        assert x.ndim == 1 or x.ndim == 2, \
+                "x must have one or two dimensions, but has {}".format(x.ndim)
+
+        if x.ndim == 1:
+            assert x.shape[0] == self.dim_x, \
+                   "Shape of x must be ({},{}), but is {}".format(
+                       self.dim_x, 1, x.shape)
+        else:
+            assert x.shape == (self.dim_x, 1), \
+                   "Shape of x must be ({},{}), but is {}".format(
+                       self.dim_x, 1, x.shape)
+
+        assert P.shape == (self.dim_x, self.dim_x), \
+               "Shape of P must be ({},{}), but is {}".format(
+                   self.dim_x, self.dim_x, P.shape)
+
+        assert Q.shape == (self.dim_x, self.dim_x), \
+               "Shape of Q must be ({},{}), but is {}".format(
+                   self.dim_x, self.dim_x, P.shape)
+
+        assert F.shape == (self.dim_x, self.dim_x), \
+               "Shape of F must be ({},{}), but is {}".format(
+                   self.dim_x, self.dim_x, F.shape)
+
+        assert np.ndim(H) == 2, \
+               "Shape of H must be (dim_z, {}), but is {}".format(
+                   P.shape[0], shape(H))
+
+        assert H.shape[1] == P.shape[0], \
+               "Shape of H must be (dim_z, {}), but is {}".format(
+                   P.shape[0], H.shape)
+
+        # shape of R must be the same as HPH'
+        hph_shape = (H.shape[0], H.shape[0])
+        r_shape = shape(R)
+
+        if H.shape[0] == 1:
+            # r can be scalar, 1D, or 2D in this case
+            assert r_shape in [(), (1,), (1, 1)], \
+            "R must be scalar or one element array, but is shaped {}".format(
+                r_shape)
+        else:
+            assert r_shape == hph_shape, \
+            "shape of R should be {} but it is {}".format(hph_shape, r_shape)
+
+
+        if z is not None:
+            z_shape = shape(z)
+        else:
+            z_shape = (self.dim_z, 1)
+
+        # H@x must have shape of z
+        Hx = dot(H, x)
+
+        if z_shape == (): # scalar or np.array(scalar)
+            assert Hx.ndim == 1 or shape(Hx) == (1, 1), \
+            "shape of z should be {}, not {} for the given H".format(
+                shape(Hx), z_shape)
+
+        elif shape(Hx) == (1,):
+            assert z_shape[0] == 1, 'Shape of z must be {} for the given H'.format(shape(Hx))
+
+        else:
+            assert (z_shape == shape(Hx) or
+                    (len(z_shape) == 1 and shape(Hx) == (z_shape[0], 1))), \
+                    "shape of z should be {}, not {} for the given H".format(
+                        shape(Hx), z_shape)
+
+        if np.ndim(Hx) > 1 and shape(Hx) != (1, 1):
+            assert shape(Hx) == z_shape, \
+               'shape of z should be {} for the given H, but it is {}'.format(
+                   shape(Hx), z_shape)
+
+
+def update(x, P, z, R, H=None, return_all=False):
+    """
+    Add a new measurement (z) to the Kalman filter. If z is None, nothing
+    is changed.
+    This can handle either the multidimensional or unidimensional case. If
+    all parameters are floats instead of arrays the filter will still work,
+    and return floats for x, P as the result.
+    update(1, 2, 1, 1, 1)  # univariate
+    update(x, P, 1
+    Parameters
+    ----------
+    x : numpy.array(dim_x, 1), or float
+        State estimate vector
+    P : numpy.array(dim_x, dim_x), or float
+        Covariance matrix
+    z : (dim_z, 1): array_like
+        measurement for this update. z can be a scalar if dim_z is 1,
+        otherwise it must be convertible to a column vector.
+    R : numpy.array(dim_z, dim_z), or float
+        Measurement noise matrix
+    H : numpy.array(dim_x, dim_x), or float, optional
+        Measurement function. If not provided, a value of 1 is assumed.
+    return_all : bool, default False
+        If true, y, K, S, and log_likelihood are returned, otherwise
+        only x and P are returned.
+    Returns
+    -------
+    x : numpy.array
+        Posterior state estimate vector
+    P : numpy.array
+        Posterior covariance matrix
+    y : numpy.array or scalar
+        Residua. Difference between measurement and state in measurement space
+    K : numpy.array
+        Kalman gain
+    S : numpy.array
+        System uncertainty in measurement space
+    log_likelihood : float
+        log likelihood of the measurement
+    """
+
+    #pylint: disable=bare-except
+
+    if z is None:
+        if return_all:
+            return x, P, None, None, None, None
+        return x, P
+
+    if H is None:
+        H = np.array([1])
+
+    if np.isscalar(H):
+        H = np.array([H])
+
+    Hx = np.atleast_1d(dot(H, x))
+    z = reshape_z(z, Hx.shape[0], x.ndim)
+
+    # error (residual) between measurement and prediction
+    y = z - Hx
+
+    # project system uncertainty into measurement space
+    S = dot(dot(H, P), H.T) + R
+
+
+    # map system uncertainty into kalman gain
+    try:
+        K = dot(dot(P, H.T), linalg.inv(S))
+    except:
+        # can't invert a 1D array, annoyingly
+        K = dot(dot(P, H.T), 1./S)
+
+
+    # predict new x with residual scaled by the kalman gain
+    x = x + dot(K, y)
+
+    # P = (I-KH)P(I-KH)' + KRK'
+    KH = dot(K, H)
+
+    try:
+        I_KH = np.eye(KH.shape[0]) - KH
+    except:
+        I_KH = np.array([1 - KH])
+    P = dot(dot(I_KH, P), I_KH.T) + dot(dot(K, R), K.T)
+
+
+    if return_all:
+        # compute log likelihood
+        log_likelihood = logpdf(z, dot(H, x), S)
+        return x, P, y, K, S, log_likelihood
+    return x, P
+
+
+def update_steadystate(x, z, K, H=None):
+    """
+    Add a new measurement (z) to the Kalman filter. If z is None, nothing
+    is changed.
+    Parameters
+    ----------
+    x : numpy.array(dim_x, 1), or float
+        State estimate vector
+    z : (dim_z, 1): array_like
+        measurement for this update. z can be a scalar if dim_z is 1,
+        otherwise it must be convertible to a column vector.
+    K : numpy.array, or float
+        Kalman gain matrix
+    H : numpy.array(dim_x, dim_x), or float, optional
+        Measurement function. If not provided, a value of 1 is assumed.
+    Returns
+    -------
+    x : numpy.array
+        Posterior state estimate vector
+    Examples
+    --------
+    This can handle either the multidimensional or unidimensional case. If
+    all parameters are floats instead of arrays the filter will still work,
+    and return floats for x, P as the result.
+    >>> update_steadystate(1, 2, 1)  # univariate
+    >>> update_steadystate(x, P, z, H)
+    """
+
+
+    if z is None:
+        return x
+
+    if H is None:
+        H = np.array([1])
+
+    if np.isscalar(H):
+        H = np.array([H])
+
+    Hx = np.atleast_1d(dot(H, x))
+    z = reshape_z(z, Hx.shape[0], x.ndim)
+
+    # error (residual) between measurement and prediction
+    y = z - Hx
+
+    # estimate new x with residual scaled by the kalman gain
+    return x + dot(K, y)
+
+
+def predict(x, P, F=1, Q=0, u=0, B=1, alpha=1.):
+    """
+    Predict next state (prior) using the Kalman filter state propagation
+    equations.
+    Parameters
+    ----------
+    x : numpy.array
+        State estimate vector
+    P : numpy.array
+        Covariance matrix
+    F : numpy.array()
+        State Transition matrix
+    Q : numpy.array, Optional
+        Process noise matrix
+    u : numpy.array, Optional, default 0.
+        Control vector. If non-zero, it is multiplied by B
+        to create the control input into the system.
+    B : numpy.array, optional, default 0.
+        Control transition matrix.
+    alpha : float, Optional, default=1.0
+        Fading memory setting. 1.0 gives the normal Kalman filter, and
+        values slightly larger than 1.0 (such as 1.02) give a fading
+        memory effect - previous measurements have less influence on the
+        filter's estimates. This formulation of the Fading memory filter
+        (there are many) is due to Dan Simon
+    Returns
+    -------
+    x : numpy.array
+        Prior state estimate vector
+    P : numpy.array
+        Prior covariance matrix
+    """
+
+    if np.isscalar(F):
+        F = np.array(F)
+    x = dot(F, x) + dot(B, u)
+    P = (alpha * alpha) * dot(dot(F, P), F.T) + Q
+
+    return x, P
+
+
+def predict_steadystate(x, F=1, u=0, B=1):
+    """
+    Predict next state (prior) using the Kalman filter state propagation
+    equations. This steady state form only computes x, assuming that the
+    covariance is constant.
+    Parameters
+    ----------
+    x : numpy.array
+        State estimate vector
+    P : numpy.array
+        Covariance matrix
+    F : numpy.array()
+        State Transition matrix
+    u : numpy.array, Optional, default 0.
+        Control vector. If non-zero, it is multiplied by B
+        to create the control input into the system.
+    B : numpy.array, optional, default 0.
+        Control transition matrix.
+    Returns
+    -------
+    x : numpy.array
+        Prior state estimate vector
+    """
+
+    if np.isscalar(F):
+        F = np.array(F)
+    x = dot(F, x) + dot(B, u)
+
+    return x
+
+
+
+def batch_filter(x, P, zs, Fs, Qs, Hs, Rs, Bs=None, us=None,
+                 update_first=False, saver=None):
+    """
+    Batch processes a sequences of measurements.
+    Parameters
+    ----------
+    zs : list-like
+        list of measurements at each time step. Missing measurements must be
+        represented by None.
+    Fs : list-like
+        list of values to use for the state transition matrix matrix.
+    Qs : list-like
+        list of values to use for the process error
+        covariance.
+    Hs : list-like
+        list of values to use for the measurement matrix.
+    Rs : list-like
+        list of values to use for the measurement error
+        covariance.
+    Bs : list-like, optional
+        list of values to use for the control transition matrix;
+        a value of None in any position will cause the filter
+        to use `self.B` for that time step.
+    us : list-like, optional
+        list of values to use for the control input vector;
+        a value of None in any position will cause the filter to use
+        0 for that time step.
+    update_first : bool, optional
+        controls whether the order of operations is update followed by
+        predict, or predict followed by update. Default is predict->update.
+        saver : filterpy.common.Saver, optional
+            filterpy.common.Saver object. If provided, saver.save() will be
+            called after every epoch
+    Returns
+    -------
+    means : np.array((n,dim_x,1))
+        array of the state for each time step after the update. Each entry
+        is an np.array. In other words `means[k,:]` is the state at step
+        `k`.
+    covariance : np.array((n,dim_x,dim_x))
+        array of the covariances for each time step after the update.
+        In other words `covariance[k,:,:]` is the covariance at step `k`.
+    means_predictions : np.array((n,dim_x,1))
+        array of the state for each time step after the predictions. Each
+        entry is an np.array. In other words `means[k,:]` is the state at
+        step `k`.
+    covariance_predictions : np.array((n,dim_x,dim_x))
+        array of the covariances for each time step after the prediction.
+        In other words `covariance[k,:,:]` is the covariance at step `k`.
+    Examples
+    --------
+    .. code-block:: Python
+        zs = [t + random.randn()*4 for t in range (40)]
+        Fs = [kf.F for t in range (40)]
+        Hs = [kf.H for t in range (40)]
+        (mu, cov, _, _) = kf.batch_filter(zs, Rs=R_list, Fs=Fs, Hs=Hs, Qs=None,
+                                          Bs=None, us=None, update_first=False)
+        (xs, Ps, Ks, Pps) = kf.rts_smoother(mu, cov, Fs=Fs, Qs=None)
+    """
+
+    n = np.size(zs, 0)
+    dim_x = x.shape[0]
+
+    # mean estimates from Kalman Filter
+    if x.ndim == 1:
+        means = zeros((n, dim_x))
+        means_p = zeros((n, dim_x))
+    else:
+        means = zeros((n, dim_x, 1))
+        means_p = zeros((n, dim_x, 1))
+
+    # state covariances from Kalman Filter
+    covariances = zeros((n, dim_x, dim_x))
+    covariances_p = zeros((n, dim_x, dim_x))
+
+    if us is None:
+        us = [0.] * n
+        Bs = [0.] * n
+
+    if update_first:
+        for i, (z, F, Q, H, R, B, u) in enumerate(zip(zs, Fs, Qs, Hs, Rs, Bs, us)):
+
+            x, P = update(x, P, z, R=R, H=H)
+            means[i, :] = x
+            covariances[i, :, :] = P
+
+            x, P = predict(x, P, u=u, B=B, F=F, Q=Q)
+            means_p[i, :] = x
+            covariances_p[i, :, :] = P
+            if saver is not None:
+                saver.save()
+    else:
+        for i, (z, F, Q, H, R, B, u) in enumerate(zip(zs, Fs, Qs, Hs, Rs, Bs, us)):
+
+            x, P = predict(x, P, u=u, B=B, F=F, Q=Q)
+            means_p[i, :] = x
+            covariances_p[i, :, :] = P
+
+            x, P = update(x, P, z, R=R, H=H)
+            means[i, :] = x
+            covariances[i, :, :] = P
+            if saver is not None:
+                saver.save()
+
+    return (means, covariances, means_p, covariances_p)
+
+
+
+def rts_smoother(Xs, Ps, Fs, Qs):
+    """
+    Runs the Rauch-Tung-Striebel Kalman smoother on a set of
+    means and covariances computed by a Kalman filter. The usual input
+    would come from the output of `KalmanFilter.batch_filter()`.
+    Parameters
+    ----------
+    Xs : numpy.array
+       array of the means (state variable x) of the output of a Kalman
+       filter.
+    Ps : numpy.array
+        array of the covariances of the output of a kalman filter.
+    Fs : list-like collection of numpy.array
+        State transition matrix of the Kalman filter at each time step.
+    Qs : list-like collection of numpy.array, optional
+        Process noise of the Kalman filter at each time step.
+    Returns
+    -------
+    x : numpy.ndarray
+       smoothed means
+    P : numpy.ndarray
+       smoothed state covariances
+    K : numpy.ndarray
+        smoother gain at each step
+    pP : numpy.ndarray
+       predicted state covariances
+    Examples
+    --------
+    .. code-block:: Python
+        zs = [t + random.randn()*4 for t in range (40)]
+        (mu, cov, _, _) = kalman.batch_filter(zs)
+        (x, P, K, pP) = rts_smoother(mu, cov, kf.F, kf.Q)
+    """
+
+    if len(Xs) != len(Ps):
+        raise ValueError('length of Xs and Ps must be the same')
+
+    n = Xs.shape[0]
+    dim_x = Xs.shape[1]
+
+    # smoother gain
+    K = zeros((n, dim_x, dim_x))
+    x, P, pP = Xs.copy(), Ps.copy(), Ps.copy()
+
+    for k in range(n-2, -1, -1):
+        pP[k] = dot(dot(Fs[k], P[k]), Fs[k].T) + Qs[k]
+
+        #pylint: disable=bad-whitespace
+        K[k]  = dot(dot(P[k], Fs[k].T), linalg.inv(pP[k]))
+        x[k] += dot(K[k], x[k+1] - dot(Fs[k], x[k]))
+        P[k] += dot(dot(K[k], P[k+1] - pP[k]), K[k].T)
+
+    return (x, P, K, pP)

trackers/ocsort/ocsort.py

@@ -0,0 +1,328 @@
+"""
+    This script is adopted from the SORT script by Alex Bewley alex@bewley.ai
+"""
+from __future__ import print_function
+
+import numpy as np
+from .association import *
+from ultralytics.yolo.utils.ops import xywh2xyxy
+
+
+def k_previous_obs(observations, cur_age, k):
+    if len(observations) == 0:
+        return [-1, -1, -1, -1, -1]
+    for i in range(k):
+        dt = k - i
+        if cur_age - dt in observations:
+            return observations[cur_age-dt]
+    max_age = max(observations.keys())
+    return observations[max_age]
+
+
+def convert_bbox_to_z(bbox):
+    """
+    Takes a bounding box in the form [x1,y1,x2,y2] and returns z in the form
+      [x,y,s,r] where x,y is the centre of the box and s is the scale/area and r is
+      the aspect ratio
+    """
+    w = bbox[2] - bbox[0]
+    h = bbox[3] - bbox[1]
+    x = bbox[0] + w/2.
+    y = bbox[1] + h/2.
+    s = w * h  # scale is just area
+    r = w / float(h+1e-6)
+    return np.array([x, y, s, r]).reshape((4, 1))
+
+
+def convert_x_to_bbox(x, score=None):
+    """
+    Takes a bounding box in the centre form [x,y,s,r] and returns it in the form
+      [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right
+    """
+    w = np.sqrt(x[2] * x[3])
+    h = x[2] / w
+    if(score == None):
+      return np.array([x[0]-w/2., x[1]-h/2., x[0]+w/2., x[1]+h/2.]).reshape((1, 4))
+    else:
+      return np.array([x[0]-w/2., x[1]-h/2., x[0]+w/2., x[1]+h/2., score]).reshape((1, 5))
+
+
+def speed_direction(bbox1, bbox2):
+    cx1, cy1 = (bbox1[0]+bbox1[2]) / 2.0, (bbox1[1]+bbox1[3])/2.0
+    cx2, cy2 = (bbox2[0]+bbox2[2]) / 2.0, (bbox2[1]+bbox2[3])/2.0
+    speed = np.array([cy2-cy1, cx2-cx1])
+    norm = np.sqrt((cy2-cy1)**2 + (cx2-cx1)**2) + 1e-6
+    return speed / norm
+
+
+class KalmanBoxTracker(object):
+    """
+    This class represents the internal state of individual tracked objects observed as bbox.
+    """
+    count = 0
+
+    def __init__(self, bbox, cls, delta_t=3, orig=False):
+        """
+        Initialises a tracker using initial bounding box.
+
+        """
+        # define constant velocity model
+        if not orig:
+          from .kalmanfilter import KalmanFilterNew as KalmanFilter
+          self.kf = KalmanFilter(dim_x=7, dim_z=4)
+        else:
+          from filterpy.kalman import KalmanFilter
+          self.kf = KalmanFilter(dim_x=7, dim_z=4)
+        self.kf.F = np.array([[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 1, 0], [0, 0, 1, 0, 0, 0, 1], [
+                            0, 0, 0, 1, 0, 0, 0],  [0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 0, 1]])
+        self.kf.H = np.array([[1, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0],
+                            [0, 0, 1, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 0]])
+
+        self.kf.R[2:, 2:] *= 10.
+        self.kf.P[4:, 4:] *= 1000.  # give high uncertainty to the unobservable initial velocities
+        self.kf.P *= 10.
+        self.kf.Q[-1, -1] *= 0.01
+        self.kf.Q[4:, 4:] *= 0.01
+
+        self.kf.x[:4] = convert_bbox_to_z(bbox)
+        self.time_since_update = 0
+        self.id = KalmanBoxTracker.count
+        KalmanBoxTracker.count += 1
+        self.history = []
+        self.hits = 0
+        self.hit_streak = 0
+        self.age = 0
+        self.conf = bbox[-1]
+        self.cls = cls
+        """
+        NOTE: [-1,-1,-1,-1,-1] is a compromising placeholder for non-observation status, the same for the return of 
+        function k_previous_obs. It is ugly and I do not like it. But to support generate observation array in a 
+        fast and unified way, which you would see below k_observations = np.array([k_previous_obs(...]]), let's bear it for now.
+        """
+        self.last_observation = np.array([-1, -1, -1, -1, -1])  # placeholder
+        self.observations = dict()
+        self.history_observations = []
+        self.velocity = None
+        self.delta_t = delta_t
+
+    def update(self, bbox, cls):
+        """
+        Updates the state vector with observed bbox.
+        """
+        
+        if bbox is not None:
+            self.conf = bbox[-1]
+            self.cls = cls
+            if self.last_observation.sum() >= 0:  # no previous observation
+                previous_box = None
+                for i in range(self.delta_t):
+                    dt = self.delta_t - i
+                    if self.age - dt in self.observations:
+                        previous_box = self.observations[self.age-dt]
+                        break
+                if previous_box is None:
+                    previous_box = self.last_observation
+                """
+                  Estimate the track speed direction with observations \Delta t steps away
+                """
+                self.velocity = speed_direction(previous_box, bbox)
+            
+            """
+              Insert new observations. This is a ugly way to maintain both self.observations
+              and self.history_observations. Bear it for the moment.
+            """
+            self.last_observation = bbox
+            self.observations[self.age] = bbox
+            self.history_observations.append(bbox)
+
+            self.time_since_update = 0
+            self.history = []
+            self.hits += 1
+            self.hit_streak += 1
+            self.kf.update(convert_bbox_to_z(bbox))
+        else:
+            self.kf.update(bbox)
+
+    def predict(self):
+        """
+        Advances the state vector and returns the predicted bounding box estimate.
+        """
+        if((self.kf.x[6]+self.kf.x[2]) <= 0):
+            self.kf.x[6] *= 0.0
+
+        self.kf.predict()
+        self.age += 1
+        if(self.time_since_update > 0):
+            self.hit_streak = 0
+        self.time_since_update += 1
+        self.history.append(convert_x_to_bbox(self.kf.x))
+        return self.history[-1]
+
+    def get_state(self):
+        """
+        Returns the current bounding box estimate.
+        """
+        return convert_x_to_bbox(self.kf.x)
+
+
+"""
+    We support multiple ways for association cost calculation, by default
+    we use IoU. GIoU may have better performance in some situations. We note 
+    that we hardly normalize the cost by all methods to (0,1) which may not be 
+    the best practice.
+"""
+ASSO_FUNCS = {  "iou": iou_batch,
+                "giou": giou_batch,
+                "ciou": ciou_batch,
+                "diou": diou_batch,
+                "ct_dist": ct_dist}
+
+
+class OCSort(object):
+    def __init__(self, det_thresh, max_age=30, min_hits=3, 
+        iou_threshold=0.3, delta_t=3, asso_func="iou", inertia=0.2, use_byte=False):
+        """
+        Sets key parameters for SORT
+        """
+        self.max_age = max_age
+        self.min_hits = min_hits
+        self.iou_threshold = iou_threshold
+        self.trackers = []
+        self.frame_count = 0
+        self.det_thresh = det_thresh
+        self.delta_t = delta_t
+        self.asso_func = ASSO_FUNCS[asso_func]
+        self.inertia = inertia
+        self.use_byte = use_byte
+        KalmanBoxTracker.count = 0
+
+    def update(self, dets, _):
+        """
+        Params:
+          dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
+        Requires: this method must be called once for each frame even with empty detections (use np.empty((0, 5)) for frames without detections).
+        Returns the a similar array, where the last column is the object ID.
+        NOTE: The number of objects returned may differ from the number of detections provided.
+        """
+
+        self.frame_count += 1
+        
+        xyxys = dets[:, 0:4]
+        confs = dets[:, 4]
+        clss = dets[:, 5]
+        
+        classes = clss.numpy()
+        xyxys = xyxys.numpy()
+        confs = confs.numpy()
+
+        output_results = np.column_stack((xyxys, confs, classes))
+        
+        inds_low = confs > 0.1
+        inds_high = confs < self.det_thresh
+        inds_second = np.logical_and(inds_low, inds_high)  # self.det_thresh > score > 0.1, for second matching
+        dets_second = output_results[inds_second]  # detections for second matching
+        remain_inds = confs > self.det_thresh
+        dets = output_results[remain_inds]
+
+        # get predicted locations from existing trackers.
+        trks = np.zeros((len(self.trackers), 5))
+        to_del = []
+        ret = []
+        for t, trk in enumerate(trks):
+            pos = self.trackers[t].predict()[0]
+            trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
+            if np.any(np.isnan(pos)):
+                to_del.append(t)
+        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
+        for t in reversed(to_del):
+            self.trackers.pop(t)
+
+        velocities = np.array(
+            [trk.velocity if trk.velocity is not None else np.array((0, 0)) for trk in self.trackers])
+        last_boxes = np.array([trk.last_observation for trk in self.trackers])
+        k_observations = np.array(
+            [k_previous_obs(trk.observations, trk.age, self.delta_t) for trk in self.trackers])
+
+        """
+            First round of association
+        """
+        matched, unmatched_dets, unmatched_trks = associate(
+            dets, trks, self.iou_threshold, velocities, k_observations, self.inertia)
+        for m in matched:
+            self.trackers[m[1]].update(dets[m[0], :5], dets[m[0], 5])
+
+        """
+            Second round of associaton by OCR
+        """
+        # BYTE association
+        if self.use_byte and len(dets_second) > 0 and unmatched_trks.shape[0] > 0:
+            u_trks = trks[unmatched_trks]
+            iou_left = self.asso_func(dets_second, u_trks)          # iou between low score detections and unmatched tracks
+            iou_left = np.array(iou_left)
+            if iou_left.max() > self.iou_threshold:
+                """
+                    NOTE: by using a lower threshold, e.g., self.iou_threshold - 0.1, you may
+                    get a higher performance especially on MOT17/MOT20 datasets. But we keep it
+                    uniform here for simplicity
+                """
+                matched_indices = linear_assignment(-iou_left)
+                to_remove_trk_indices = []
+                for m in matched_indices:
+                    det_ind, trk_ind = m[0], unmatched_trks[m[1]]
+                    if iou_left[m[0], m[1]] < self.iou_threshold:
+                        continue
+                    self.trackers[trk_ind].update(dets_second[det_ind, :5], dets_second[det_ind, 5])
+                    to_remove_trk_indices.append(trk_ind)
+                unmatched_trks = np.setdiff1d(unmatched_trks, np.array(to_remove_trk_indices))
+
+        if unmatched_dets.shape[0] > 0 and unmatched_trks.shape[0] > 0:
+            left_dets = dets[unmatched_dets]
+            left_trks = last_boxes[unmatched_trks]
+            iou_left = self.asso_func(left_dets, left_trks)
+            iou_left = np.array(iou_left)
+            if iou_left.max() > self.iou_threshold:
+                """
+                    NOTE: by using a lower threshold, e.g., self.iou_threshold - 0.1, you may
+                    get a higher performance especially on MOT17/MOT20 datasets. But we keep it
+                    uniform here for simplicity
+                """
+                rematched_indices = linear_assignment(-iou_left)
+                to_remove_det_indices = []
+                to_remove_trk_indices = []
+                for m in rematched_indices:
+                    det_ind, trk_ind = unmatched_dets[m[0]], unmatched_trks[m[1]]
+                    if iou_left[m[0], m[1]] < self.iou_threshold:
+                        continue
+                    self.trackers[trk_ind].update(dets[det_ind, :5], dets[det_ind, 5])
+                    to_remove_det_indices.append(det_ind)
+                    to_remove_trk_indices.append(trk_ind)
+                unmatched_dets = np.setdiff1d(unmatched_dets, np.array(to_remove_det_indices))
+                unmatched_trks = np.setdiff1d(unmatched_trks, np.array(to_remove_trk_indices))
+
+        for m in unmatched_trks:
+            self.trackers[m].update(None, None)
+
+        # create and initialise new trackers for unmatched detections
+        for i in unmatched_dets:
+            trk = KalmanBoxTracker(dets[i, :5], dets[i, 5], delta_t=self.delta_t)
+            self.trackers.append(trk)
+        i = len(self.trackers)
+        for trk in reversed(self.trackers):
+            if trk.last_observation.sum() < 0:
+                d = trk.get_state()[0]
+            else:
+                """
+                    this is optional to use the recent observation or the kalman filter prediction,
+                    we didn't notice significant difference here
+                """
+                d = trk.last_observation[:4]
+            if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits):
+                # +1 as MOT benchmark requires positive
+                ret.append(np.concatenate((d, [trk.id+1], [trk.cls], [trk.conf])).reshape(1, -1))
+            i -= 1
+            # remove dead tracklet
+            if(trk.time_since_update > self.max_age):
+                self.trackers.pop(i)
+        if(len(ret) > 0):
+            return np.concatenate(ret)
+        return np.empty((0, 5))

trackers/reid_export.py

@@ -0,0 +1,313 @@
+import argparse
+
+import os
+# limit the number of cpus used by high performance libraries
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+
+import sys
+import numpy as np
+from pathlib import Path
+import torch
+import time
+import platform
+import pandas as pd
+import subprocess
+import torch.backends.cudnn as cudnn
+from torch.utils.mobile_optimizer import optimize_for_mobile
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0].parents[0]  # yolov5 strongsort root directory
+WEIGHTS = ROOT / 'weights'
+
+
+if str(ROOT) not in sys.path:
+    sys.path.append(str(ROOT))  # add ROOT to PATH
+if str(ROOT / 'yolov5') not in sys.path:
+    sys.path.append(str(ROOT / 'yolov5'))  # add yolov5 ROOT to PATH
+
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+
+import logging
+from ultralytics.yolo.utils.torch_utils import select_device
+from ultralytics.yolo.utils import LOGGER, colorstr, ops
+from ultralytics.yolo.utils.checks import check_requirements, check_version
+from trackers.strongsort.deep.models import build_model
+from trackers.strongsort.deep.reid_model_factory import get_model_name, load_pretrained_weights
+
+
+def file_size(path):
+    # Return file/dir size (MB)
+    path = Path(path)
+    if path.is_file():
+        return path.stat().st_size / 1E6
+    elif path.is_dir():
+        return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6
+    else:
+        return 0.0
+
+
+def export_formats():
+    # YOLOv5 export formats
+    x = [
+        ['PyTorch', '-', '.pt', True, True],
+        ['TorchScript', 'torchscript', '.torchscript', True, True],
+        ['ONNX', 'onnx', '.onnx', True, True],
+        ['OpenVINO', 'openvino', '_openvino_model', True, False],
+        ['TensorRT', 'engine', '.engine', False, True],
+        ['TensorFlow Lite', 'tflite', '.tflite', True, False],
+    ]
+    return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])
+
+
+def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
+    # YOLOv5 TorchScript model export
+    try:
+        LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
+        f = file.with_suffix('.torchscript')
+
+        ts = torch.jit.trace(model, im, strict=False)
+        if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
+            optimize_for_mobile(ts)._save_for_lite_interpreter(str(f))
+        else:
+            ts.save(str(f))
+
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
+    except Exception as e:
+        LOGGER.info(f'{prefix} export failure: {e}')
+
+
+def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')):
+    # ONNX export
+    try:
+        check_requirements(('onnx',))
+        import onnx
+
+        f = file.with_suffix('.onnx')
+        LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
+        
+        if dynamic:
+            dynamic = {'images': {0: 'batch'}}  # shape(1,3,640,640)
+            dynamic['output'] = {0: 'batch'}  # shape(1,25200,85)
+
+        torch.onnx.export(
+            model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
+            im.cpu() if dynamic else im,
+            f,
+            verbose=False,
+            opset_version=opset,
+            do_constant_folding=True,
+            input_names=['images'],
+            output_names=['output'],
+            dynamic_axes=dynamic or None
+        )
+        # Checks
+        model_onnx = onnx.load(f)  # load onnx model
+        onnx.checker.check_model(model_onnx)  # check onnx model
+        onnx.save(model_onnx, f)
+
+        # Simplify
+        if simplify:
+            try:
+                cuda = torch.cuda.is_available()
+                check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
+                import onnxsim
+
+                LOGGER.info(f'simplifying with onnx-simplifier {onnxsim.__version__}...')
+                model_onnx, check = onnxsim.simplify(model_onnx)
+                assert check, 'assert check failed'
+                onnx.save(model_onnx, f)
+            except Exception as e:
+                LOGGER.info(f'simplifier failure: {e}')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
+    except Exception as e:
+        LOGGER.info(f'export failure: {e}')
+    
+        
+        
+def export_openvino(file, half, prefix=colorstr('OpenVINO:')):
+    # YOLOv5 OpenVINO export
+    check_requirements(('openvino-dev',))  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+    import openvino.inference_engine as ie
+    try:
+        LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
+        f = str(file).replace('.pt', f'_openvino_model{os.sep}')
+
+        cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}"
+        subprocess.check_output(cmd.split())  # export
+    except Exception as e:
+        LOGGER.info(f'export failure: {e}')
+    LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+    return f
+        
+
+def export_tflite(file, half, prefix=colorstr('TFLite:')):
+    # YOLOv5 OpenVINO export
+    try:
+        check_requirements(('openvino2tensorflow', 'tensorflow', 'tensorflow_datasets'))  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+        import openvino.inference_engine as ie
+        LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...')
+        output = Path(str(file).replace(f'_openvino_model{os.sep}', f'_tflite_model{os.sep}'))
+        modelxml = list(Path(file).glob('*.xml'))[0]
+        cmd = f"openvino2tensorflow \
+            --model_path {modelxml} \
+            --model_output_path {output} \
+            --output_pb \
+            --output_saved_model \
+            --output_no_quant_float32_tflite \
+            --output_dynamic_range_quant_tflite"
+        subprocess.check_output(cmd.split())  # export
+
+        LOGGER.info(f'{prefix} export success, results saved in {output} ({file_size(f):.1f} MB)')
+        return f
+    except Exception as e:
+        LOGGER.info(f'\n{prefix} export failure: {e}')
+        
+        
+def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
+    # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
+    try:
+        assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
+        try:
+            import tensorrt as trt
+        except Exception:
+            if platform.system() == 'Linux':
+                check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
+            import tensorrt as trt
+
+        if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
+            grid = model.model[-1].anchor_grid
+            model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
+            export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
+            model.model[-1].anchor_grid = grid
+        else:  # TensorRT >= 8
+            check_version(trt.__version__, '8.0.0', hard=True)  # require tensorrt>=8.0.0
+            export_onnx(model, im, file, 12, dynamic, simplify)  # opset 13
+        onnx = file.with_suffix('.onnx')
+
+        LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
+        assert onnx.exists(), f'failed to export ONNX file: {onnx}'
+        f = file.with_suffix('.engine')  # TensorRT engine file
+        logger = trt.Logger(trt.Logger.INFO)
+        if verbose:
+            logger.min_severity = trt.Logger.Severity.VERBOSE
+
+        builder = trt.Builder(logger)
+        config = builder.create_builder_config()
+        config.max_workspace_size = workspace * 1 << 30
+        # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)  # fix TRT 8.4 deprecation notice
+
+        flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+        network = builder.create_network(flag)
+        parser = trt.OnnxParser(network, logger)
+        if not parser.parse_from_file(str(onnx)):
+            raise RuntimeError(f'failed to load ONNX file: {onnx}')
+
+        inputs = [network.get_input(i) for i in range(network.num_inputs)]
+        outputs = [network.get_output(i) for i in range(network.num_outputs)]
+        LOGGER.info(f'{prefix} Network Description:')
+        for inp in inputs:
+            LOGGER.info(f'{prefix}\tinput "{inp.name}" with shape {inp.shape} and dtype {inp.dtype}')
+        for out in outputs:
+            LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
+
+        if dynamic:
+            if im.shape[0] <= 1:
+                LOGGER.warning(f"{prefix}WARNING: --dynamic model requires maximum --batch-size argument")
+            profile = builder.create_optimization_profile()
+            for inp in inputs:
+                profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
+            config.add_optimization_profile(profile)
+
+        LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine in {f}')
+        if builder.platform_has_fast_fp16 and half:
+            config.set_flag(trt.BuilderFlag.FP16)
+        with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
+            t.write(engine.serialize())
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        return f
+    except Exception as e:
+        LOGGER.info(f'\n{prefix} export failure: {e}')
+        
+        
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="ReID export")
+    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
+    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[256, 128], help='image (h, w)')
+    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
+    parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
+    parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes')
+    parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
+    parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
+    parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
+    parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
+    parser.add_argument('--weights', nargs='+', type=str, default=WEIGHTS / 'osnet_x0_25_msmt17.pt', help='model.pt path(s)')
+    parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
+    parser.add_argument('--include',
+                        nargs='+',
+                        default=['torchscript'],
+                        help='torchscript, onnx, openvino, engine')
+    args = parser.parse_args()
+
+    t = time.time()
+
+    include = [x.lower() for x in args.include]  # to lowercase
+    fmts = tuple(export_formats()['Argument'][1:])  # --include arguments
+    flags = [x in include for x in fmts]
+    assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}'
+    jit, onnx, openvino, engine, tflite = flags  # export booleans
+
+    args.device = select_device(args.device)
+    if args.half:
+        assert args.device.type != 'cpu', '--half only compatible with GPU export, i.e. use --device 0'
+        assert not args.dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both'
+    
+    if type(args.weights) is list:
+        args.weights = Path(args.weights[0])
+
+    model = build_model(
+        get_model_name(args.weights),
+        num_classes=1,
+        pretrained=not (args.weights and args.weights.is_file() and args.weights.suffix == '.pt'),
+        use_gpu=args.device
+    ).to(args.device)
+    load_pretrained_weights(model, args.weights)
+    model.eval()
+
+    if args.optimize:
+        assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
+    
+    im = torch.zeros(args.batch_size, 3, args.imgsz[0], args.imgsz[1]).to(args.device)  # image size(1,3,640,480) BCHW iDetection
+    for _ in range(2):
+        y = model(im)  # dry runs
+    if args.half:
+        im, model = im.half(), model.half()  # to FP16
+    shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
+    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {args.weights} with output shape {shape} ({file_size(args.weights):.1f} MB)")
+    
+    # Exports
+    f = [''] * len(fmts)  # exported filenames
+    if jit:
+        f[0] = export_torchscript(model, im, args.weights, args.optimize)  # opset 12
+    if engine:  # TensorRT required before ONNX
+        f[1] = export_engine(model, im, args.weights, args.half, args.dynamic, args.simplify, args.workspace, args.verbose)
+    if onnx:  # OpenVINO requires ONNX
+        f[2] = export_onnx(model, im, args.weights, args.opset, args.dynamic, args.simplify)  # opset 12
+    if openvino:
+        f[3] = export_openvino(args.weights, args.half)
+    if tflite:
+        export_tflite(f, False)
+
+    # Finish
+    f = [str(x) for x in f if x]  # filter out '' and None
+    if any(f):
+        LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)'
+                    f"\nResults saved to {colorstr('bold', args.weights.parent.resolve())}"
+                    f"\nVisualize:       https://netron.app")
+

trackers/strongsort/.gitignore

@@ -0,0 +1,13 @@
+# Folders
+__pycache__/
+build/
+*.egg-info
+
+
+# Files
+*.weights
+*.t7
+*.mp4
+*.avi
+*.so
+*.txt

trackers/strongsort/configs/strongsort.yaml

@@ -0,0 +1,11 @@
+strongsort:
+  ecc: true
+  ema_alpha: 0.8962157769329083
+  max_age: 40
+  max_dist: 0.1594374041012136
+  max_iou_dist: 0.5431835667667874
+  max_unmatched_preds: 0
+  mc_lambda: 0.995
+  n_init: 3
+  nn_budget: 100
+  conf_thres: 0.5122620708221085

trackers/strongsort/deep/checkpoint/osnet_x0_25_market1501.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0ff09177a21417a19bc73bbf5af4eb5d2b097c2074ed35e67819ee6cd93612c
+size 2462783

trackers/strongsort/deep/checkpoint/osnet_x0_25_msmt17.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f57607fed9f502b9efed546108132ee715df5a5b6e6932c6269bacb47f59f99
+size 3057863

trackers/strongsort/deep/checkpoint/osnet_x1_0_msmt17.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7d73dc67c016fd044e4027ff856019496392a7aca8fa0ed56d862a1632c1cf2
+size 10994685

trackers/strongsort/deep/models/__init__.py

@@ -0,0 +1,122 @@
+from __future__ import absolute_import
+import torch
+
+from .pcb import *
+from .mlfn import *
+from .hacnn import *
+from .osnet import *
+from .senet import *
+from .mudeep import *
+from .nasnet import *
+from .resnet import *
+from .densenet import *
+from .xception import *
+from .osnet_ain import *
+from .resnetmid import *
+from .shufflenet import *
+from .squeezenet import *
+from .inceptionv4 import *
+from .mobilenetv2 import *
+from .resnet_ibn_a import *
+from .resnet_ibn_b import *
+from .shufflenetv2 import *
+from .inceptionresnetv2 import *
+
+__model_factory = {
+    # image classification models
+    'resnet18': resnet18,
+    'resnet34': resnet34,
+    'resnet50': resnet50,
+    'resnet101': resnet101,
+    'resnet152': resnet152,
+    'resnext50_32x4d': resnext50_32x4d,
+    'resnext101_32x8d': resnext101_32x8d,
+    'resnet50_fc512': resnet50_fc512,
+    'se_resnet50': se_resnet50,
+    'se_resnet50_fc512': se_resnet50_fc512,
+    'se_resnet101': se_resnet101,
+    'se_resnext50_32x4d': se_resnext50_32x4d,
+    'se_resnext101_32x4d': se_resnext101_32x4d,
+    'densenet121': densenet121,
+    'densenet169': densenet169,
+    'densenet201': densenet201,
+    'densenet161': densenet161,
+    'densenet121_fc512': densenet121_fc512,
+    'inceptionresnetv2': inceptionresnetv2,
+    'inceptionv4': inceptionv4,
+    'xception': xception,
+    'resnet50_ibn_a': resnet50_ibn_a,
+    'resnet50_ibn_b': resnet50_ibn_b,
+    # lightweight models
+    'nasnsetmobile': nasnetamobile,
+    'mobilenetv2_x1_0': mobilenetv2_x1_0,
+    'mobilenetv2_x1_4': mobilenetv2_x1_4,
+    'shufflenet': shufflenet,
+    'squeezenet1_0': squeezenet1_0,
+    'squeezenet1_0_fc512': squeezenet1_0_fc512,
+    'squeezenet1_1': squeezenet1_1,
+    'shufflenet_v2_x0_5': shufflenet_v2_x0_5,
+    'shufflenet_v2_x1_0': shufflenet_v2_x1_0,
+    'shufflenet_v2_x1_5': shufflenet_v2_x1_5,
+    'shufflenet_v2_x2_0': shufflenet_v2_x2_0,
+    # reid-specific models
+    'mudeep': MuDeep,
+    'resnet50mid': resnet50mid,
+    'hacnn': HACNN,
+    'pcb_p6': pcb_p6,
+    'pcb_p4': pcb_p4,
+    'mlfn': mlfn,
+    'osnet_x1_0': osnet_x1_0,
+    'osnet_x0_75': osnet_x0_75,
+    'osnet_x0_5': osnet_x0_5,
+    'osnet_x0_25': osnet_x0_25,
+    'osnet_ibn_x1_0': osnet_ibn_x1_0,
+    'osnet_ain_x1_0': osnet_ain_x1_0,
+    'osnet_ain_x0_75': osnet_ain_x0_75,
+    'osnet_ain_x0_5': osnet_ain_x0_5,
+    'osnet_ain_x0_25': osnet_ain_x0_25
+}
+
+
+def show_avai_models():
+    """Displays available models.
+
+    Examples::
+        >>> from torchreid import models
+        >>> models.show_avai_models()
+    """
+    print(list(__model_factory.keys()))
+
+
+def build_model(
+    name, num_classes, loss='softmax', pretrained=True, use_gpu=True
+):
+    """A function wrapper for building a model.
+
+    Args:
+        name (str): model name.
+        num_classes (int): number of training identities.
+        loss (str, optional): loss function to optimize the model. Currently
+            supports "softmax" and "triplet". Default is "softmax".
+        pretrained (bool, optional): whether to load ImageNet-pretrained weights.
+            Default is True.
+        use_gpu (bool, optional): whether to use gpu. Default is True.
+
+    Returns:
+        nn.Module
+
+    Examples::
+        >>> from torchreid import models
+        >>> model = models.build_model('resnet50', 751, loss='softmax')
+    """
+    avai_models = list(__model_factory.keys())
+    if name not in avai_models:
+        raise KeyError(
+            'Unknown model: {}. Must be one of {}'.format(name, avai_models)
+        )
+    return __model_factory[name](
+        num_classes=num_classes,
+        loss=loss,
+        pretrained=pretrained,
+        use_gpu=use_gpu
+    )

trackers/strongsort/deep/models/densenet.py

@@ -0,0 +1,380 @@
+"""
+Code source: https://github.com/pytorch/vision
+"""
+from __future__ import division, absolute_import
+import re
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils import model_zoo
+
+__all__ = [
+    'densenet121', 'densenet169', 'densenet201', 'densenet161',
+    'densenet121_fc512'
+]
+
+model_urls = {
+    'densenet121':
+    'https://download.pytorch.org/models/densenet121-a639ec97.pth',
+    'densenet169':
+    'https://download.pytorch.org/models/densenet169-b2777c0a.pth',
+    'densenet201':
+    'https://download.pytorch.org/models/densenet201-c1103571.pth',
+    'densenet161':
+    'https://download.pytorch.org/models/densenet161-8d451a50.pth',
+}
+
+
+class _DenseLayer(nn.Sequential):
+
+    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):
+        super(_DenseLayer, self).__init__()
+        self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
+        self.add_module('relu1', nn.ReLU(inplace=True)),
+        self.add_module(
+            'conv1',
+            nn.Conv2d(
+                num_input_features,
+                bn_size * growth_rate,
+                kernel_size=1,
+                stride=1,
+                bias=False
+            )
+        ),
+        self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
+        self.add_module('relu2', nn.ReLU(inplace=True)),
+        self.add_module(
+            'conv2',
+            nn.Conv2d(
+                bn_size * growth_rate,
+                growth_rate,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False
+            )
+        ),
+        self.drop_rate = drop_rate
+
+    def forward(self, x):
+        new_features = super(_DenseLayer, self).forward(x)
+        if self.drop_rate > 0:
+            new_features = F.dropout(
+                new_features, p=self.drop_rate, training=self.training
+            )
+        return torch.cat([x, new_features], 1)
+
+
+class _DenseBlock(nn.Sequential):
+
+    def __init__(
+        self, num_layers, num_input_features, bn_size, growth_rate, drop_rate
+    ):
+        super(_DenseBlock, self).__init__()
+        for i in range(num_layers):
+            layer = _DenseLayer(
+                num_input_features + i*growth_rate, growth_rate, bn_size,
+                drop_rate
+            )
+            self.add_module('denselayer%d' % (i+1), layer)
+
+
+class _Transition(nn.Sequential):
+
+    def __init__(self, num_input_features, num_output_features):
+        super(_Transition, self).__init__()
+        self.add_module('norm', nn.BatchNorm2d(num_input_features))
+        self.add_module('relu', nn.ReLU(inplace=True))
+        self.add_module(
+            'conv',
+            nn.Conv2d(
+                num_input_features,
+                num_output_features,
+                kernel_size=1,
+                stride=1,
+                bias=False
+            )
+        )
+        self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
+
+
+class DenseNet(nn.Module):
+    """Densely connected network.
+    
+    Reference:
+        Huang et al. Densely Connected Convolutional Networks. CVPR 2017.
+
+    Public keys:
+        - ``densenet121``: DenseNet121.
+        - ``densenet169``: DenseNet169.
+        - ``densenet201``: DenseNet201.
+        - ``densenet161``: DenseNet161.
+        - ``densenet121_fc512``: DenseNet121 + FC.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        growth_rate=32,
+        block_config=(6, 12, 24, 16),
+        num_init_features=64,
+        bn_size=4,
+        drop_rate=0,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    ):
+
+        super(DenseNet, self).__init__()
+        self.loss = loss
+
+        # First convolution
+        self.features = nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        'conv0',
+                        nn.Conv2d(
+                            3,
+                            num_init_features,
+                            kernel_size=7,
+                            stride=2,
+                            padding=3,
+                            bias=False
+                        )
+                    ),
+                    ('norm0', nn.BatchNorm2d(num_init_features)),
+                    ('relu0', nn.ReLU(inplace=True)),
+                    (
+                        'pool0',
+                        nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+                    ),
+                ]
+            )
+        )
+
+        # Each denseblock
+        num_features = num_init_features
+        for i, num_layers in enumerate(block_config):
+            block = _DenseBlock(
+                num_layers=num_layers,
+                num_input_features=num_features,
+                bn_size=bn_size,
+                growth_rate=growth_rate,
+                drop_rate=drop_rate
+            )
+            self.features.add_module('denseblock%d' % (i+1), block)
+            num_features = num_features + num_layers*growth_rate
+            if i != len(block_config) - 1:
+                trans = _Transition(
+                    num_input_features=num_features,
+                    num_output_features=num_features // 2
+                )
+                self.features.add_module('transition%d' % (i+1), trans)
+                num_features = num_features // 2
+
+        # Final batch norm
+        self.features.add_module('norm5', nn.BatchNorm2d(num_features))
+
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.feature_dim = num_features
+        self.fc = self._construct_fc_layer(fc_dims, num_features, dropout_p)
+
+        # Linear layer
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer.
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        f = self.features(x)
+        f = F.relu(f, inplace=True)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if self.fc is not None:
+            v = self.fc(v)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+
+    # '.'s are no longer allowed in module names, but pervious _DenseLayer
+    # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
+    # They are also in the checkpoints in model_urls. This pattern is used
+    # to find such keys.
+    pattern = re.compile(
+        r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$'
+    )
+    for key in list(pretrain_dict.keys()):
+        res = pattern.match(key)
+        if res:
+            new_key = res.group(1) + res.group(2)
+            pretrain_dict[new_key] = pretrain_dict[key]
+            del pretrain_dict[key]
+
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+"""
+Dense network configurations:
+--
+densenet121: num_init_features=64, growth_rate=32, block_config=(6, 12, 24, 16)
+densenet169: num_init_features=64, growth_rate=32, block_config=(6, 12, 32, 32)
+densenet201: num_init_features=64, growth_rate=32, block_config=(6, 12, 48, 32)
+densenet161: num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24)
+"""
+
+
+def densenet121(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = DenseNet(
+        num_classes=num_classes,
+        loss=loss,
+        num_init_features=64,
+        growth_rate=32,
+        block_config=(6, 12, 24, 16),
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['densenet121'])
+    return model
+
+
+def densenet169(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = DenseNet(
+        num_classes=num_classes,
+        loss=loss,
+        num_init_features=64,
+        growth_rate=32,
+        block_config=(6, 12, 32, 32),
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['densenet169'])
+    return model
+
+
+def densenet201(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = DenseNet(
+        num_classes=num_classes,
+        loss=loss,
+        num_init_features=64,
+        growth_rate=32,
+        block_config=(6, 12, 48, 32),
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['densenet201'])
+    return model
+
+
+def densenet161(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = DenseNet(
+        num_classes=num_classes,
+        loss=loss,
+        num_init_features=96,
+        growth_rate=48,
+        block_config=(6, 12, 36, 24),
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['densenet161'])
+    return model
+
+
+def densenet121_fc512(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = DenseNet(
+        num_classes=num_classes,
+        loss=loss,
+        num_init_features=64,
+        growth_rate=32,
+        block_config=(6, 12, 24, 16),
+        fc_dims=[512],
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['densenet121'])
+    return model

trackers/strongsort/deep/models/hacnn.py

@@ -0,0 +1,414 @@
+from __future__ import division, absolute_import
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ['HACNN']
+
+
+class ConvBlock(nn.Module):
+    """Basic convolutional block.
+    
+    convolution + batch normalization + relu.
+
+    Args:
+        in_c (int): number of input channels.
+        out_c (int): number of output channels.
+        k (int or tuple): kernel size.
+        s (int or tuple): stride.
+        p (int or tuple): padding.
+    """
+
+    def __init__(self, in_c, out_c, k, s=1, p=0):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p)
+        self.bn = nn.BatchNorm2d(out_c)
+
+    def forward(self, x):
+        return F.relu(self.bn(self.conv(x)))
+
+
+class InceptionA(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super(InceptionA, self).__init__()
+        mid_channels = out_channels // 4
+
+        self.stream1 = nn.Sequential(
+            ConvBlock(in_channels, mid_channels, 1),
+            ConvBlock(mid_channels, mid_channels, 3, p=1),
+        )
+        self.stream2 = nn.Sequential(
+            ConvBlock(in_channels, mid_channels, 1),
+            ConvBlock(mid_channels, mid_channels, 3, p=1),
+        )
+        self.stream3 = nn.Sequential(
+            ConvBlock(in_channels, mid_channels, 1),
+            ConvBlock(mid_channels, mid_channels, 3, p=1),
+        )
+        self.stream4 = nn.Sequential(
+            nn.AvgPool2d(3, stride=1, padding=1),
+            ConvBlock(in_channels, mid_channels, 1),
+        )
+
+    def forward(self, x):
+        s1 = self.stream1(x)
+        s2 = self.stream2(x)
+        s3 = self.stream3(x)
+        s4 = self.stream4(x)
+        y = torch.cat([s1, s2, s3, s4], dim=1)
+        return y
+
+
+class InceptionB(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super(InceptionB, self).__init__()
+        mid_channels = out_channels // 4
+
+        self.stream1 = nn.Sequential(
+            ConvBlock(in_channels, mid_channels, 1),
+            ConvBlock(mid_channels, mid_channels, 3, s=2, p=1),
+        )
+        self.stream2 = nn.Sequential(
+            ConvBlock(in_channels, mid_channels, 1),
+            ConvBlock(mid_channels, mid_channels, 3, p=1),
+            ConvBlock(mid_channels, mid_channels, 3, s=2, p=1),
+        )
+        self.stream3 = nn.Sequential(
+            nn.MaxPool2d(3, stride=2, padding=1),
+            ConvBlock(in_channels, mid_channels * 2, 1),
+        )
+
+    def forward(self, x):
+        s1 = self.stream1(x)
+        s2 = self.stream2(x)
+        s3 = self.stream3(x)
+        y = torch.cat([s1, s2, s3], dim=1)
+        return y
+
+
+class SpatialAttn(nn.Module):
+    """Spatial Attention (Sec. 3.1.I.1)"""
+
+    def __init__(self):
+        super(SpatialAttn, self).__init__()
+        self.conv1 = ConvBlock(1, 1, 3, s=2, p=1)
+        self.conv2 = ConvBlock(1, 1, 1)
+
+    def forward(self, x):
+        # global cross-channel averaging
+        x = x.mean(1, keepdim=True)
+        # 3-by-3 conv
+        x = self.conv1(x)
+        # bilinear resizing
+        x = F.upsample(
+            x, (x.size(2) * 2, x.size(3) * 2),
+            mode='bilinear',
+            align_corners=True
+        )
+        # scaling conv
+        x = self.conv2(x)
+        return x
+
+
+class ChannelAttn(nn.Module):
+    """Channel Attention (Sec. 3.1.I.2)"""
+
+    def __init__(self, in_channels, reduction_rate=16):
+        super(ChannelAttn, self).__init__()
+        assert in_channels % reduction_rate == 0
+        self.conv1 = ConvBlock(in_channels, in_channels // reduction_rate, 1)
+        self.conv2 = ConvBlock(in_channels // reduction_rate, in_channels, 1)
+
+    def forward(self, x):
+        # squeeze operation (global average pooling)
+        x = F.avg_pool2d(x, x.size()[2:])
+        # excitation operation (2 conv layers)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+
+class SoftAttn(nn.Module):
+    """Soft Attention (Sec. 3.1.I)
+    
+    Aim: Spatial Attention + Channel Attention
+    
+    Output: attention maps with shape identical to input.
+    """
+
+    def __init__(self, in_channels):
+        super(SoftAttn, self).__init__()
+        self.spatial_attn = SpatialAttn()
+        self.channel_attn = ChannelAttn(in_channels)
+        self.conv = ConvBlock(in_channels, in_channels, 1)
+
+    def forward(self, x):
+        y_spatial = self.spatial_attn(x)
+        y_channel = self.channel_attn(x)
+        y = y_spatial * y_channel
+        y = torch.sigmoid(self.conv(y))
+        return y
+
+
+class HardAttn(nn.Module):
+    """Hard Attention (Sec. 3.1.II)"""
+
+    def __init__(self, in_channels):
+        super(HardAttn, self).__init__()
+        self.fc = nn.Linear(in_channels, 4 * 2)
+        self.init_params()
+
+    def init_params(self):
+        self.fc.weight.data.zero_()
+        self.fc.bias.data.copy_(
+            torch.tensor(
+                [0, -0.75, 0, -0.25, 0, 0.25, 0, 0.75], dtype=torch.float
+            )
+        )
+
+    def forward(self, x):
+        # squeeze operation (global average pooling)
+        x = F.avg_pool2d(x, x.size()[2:]).view(x.size(0), x.size(1))
+        # predict transformation parameters
+        theta = torch.tanh(self.fc(x))
+        theta = theta.view(-1, 4, 2)
+        return theta
+
+
+class HarmAttn(nn.Module):
+    """Harmonious Attention (Sec. 3.1)"""
+
+    def __init__(self, in_channels):
+        super(HarmAttn, self).__init__()
+        self.soft_attn = SoftAttn(in_channels)
+        self.hard_attn = HardAttn(in_channels)
+
+    def forward(self, x):
+        y_soft_attn = self.soft_attn(x)
+        theta = self.hard_attn(x)
+        return y_soft_attn, theta
+
+
+class HACNN(nn.Module):
+    """Harmonious Attention Convolutional Neural Network.
+
+    Reference:
+        Li et al. Harmonious Attention Network for Person Re-identification. CVPR 2018.
+
+    Public keys:
+        - ``hacnn``: HACNN.
+    """
+
+    # Args:
+    #    num_classes (int): number of classes to predict
+    #    nchannels (list): number of channels AFTER concatenation
+    #    feat_dim (int): feature dimension for a single stream
+    #    learn_region (bool): whether to learn region features (i.e. local branch)
+
+    def __init__(
+        self,
+        num_classes,
+        loss='softmax',
+        nchannels=[128, 256, 384],
+        feat_dim=512,
+        learn_region=True,
+        use_gpu=True,
+        **kwargs
+    ):
+        super(HACNN, self).__init__()
+        self.loss = loss
+        self.learn_region = learn_region
+        self.use_gpu = use_gpu
+
+        self.conv = ConvBlock(3, 32, 3, s=2, p=1)
+
+        # Construct Inception + HarmAttn blocks
+        # ============== Block 1 ==============
+        self.inception1 = nn.Sequential(
+            InceptionA(32, nchannels[0]),
+            InceptionB(nchannels[0], nchannels[0]),
+        )
+        self.ha1 = HarmAttn(nchannels[0])
+
+        # ============== Block 2 ==============
+        self.inception2 = nn.Sequential(
+            InceptionA(nchannels[0], nchannels[1]),
+            InceptionB(nchannels[1], nchannels[1]),
+        )
+        self.ha2 = HarmAttn(nchannels[1])
+
+        # ============== Block 3 ==============
+        self.inception3 = nn.Sequential(
+            InceptionA(nchannels[1], nchannels[2]),
+            InceptionB(nchannels[2], nchannels[2]),
+        )
+        self.ha3 = HarmAttn(nchannels[2])
+
+        self.fc_global = nn.Sequential(
+            nn.Linear(nchannels[2], feat_dim),
+            nn.BatchNorm1d(feat_dim),
+            nn.ReLU(),
+        )
+        self.classifier_global = nn.Linear(feat_dim, num_classes)
+
+        if self.learn_region:
+            self.init_scale_factors()
+            self.local_conv1 = InceptionB(32, nchannels[0])
+            self.local_conv2 = InceptionB(nchannels[0], nchannels[1])
+            self.local_conv3 = InceptionB(nchannels[1], nchannels[2])
+            self.fc_local = nn.Sequential(
+                nn.Linear(nchannels[2] * 4, feat_dim),
+                nn.BatchNorm1d(feat_dim),
+                nn.ReLU(),
+            )
+            self.classifier_local = nn.Linear(feat_dim, num_classes)
+            self.feat_dim = feat_dim * 2
+        else:
+            self.feat_dim = feat_dim
+
+    def init_scale_factors(self):
+        # initialize scale factors (s_w, s_h) for four regions
+        self.scale_factors = []
+        self.scale_factors.append(
+            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)
+        )
+        self.scale_factors.append(
+            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)
+        )
+        self.scale_factors.append(
+            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)
+        )
+        self.scale_factors.append(
+            torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)
+        )
+
+    def stn(self, x, theta):
+        """Performs spatial transform
+        
+        x: (batch, channel, height, width)
+        theta: (batch, 2, 3)
+        """
+        grid = F.affine_grid(theta, x.size())
+        x = F.grid_sample(x, grid)
+        return x
+
+    def transform_theta(self, theta_i, region_idx):
+        """Transforms theta to include (s_w, s_h), resulting in (batch, 2, 3)"""
+        scale_factors = self.scale_factors[region_idx]
+        theta = torch.zeros(theta_i.size(0), 2, 3)
+        theta[:, :, :2] = scale_factors
+        theta[:, :, -1] = theta_i
+        if self.use_gpu:
+            theta = theta.cuda()
+        return theta
+
+    def forward(self, x):
+        assert x.size(2) == 160 and x.size(3) == 64, \
+            'Input size does not match, expected (160, 64) but got ({}, {})'.format(x.size(2), x.size(3))
+        x = self.conv(x)
+
+        # ============== Block 1 ==============
+        # global branch
+        x1 = self.inception1(x)
+        x1_attn, x1_theta = self.ha1(x1)
+        x1_out = x1 * x1_attn
+        # local branch
+        if self.learn_region:
+            x1_local_list = []
+            for region_idx in range(4):
+                x1_theta_i = x1_theta[:, region_idx, :]
+                x1_theta_i = self.transform_theta(x1_theta_i, region_idx)
+                x1_trans_i = self.stn(x, x1_theta_i)
+                x1_trans_i = F.upsample(
+                    x1_trans_i, (24, 28), mode='bilinear', align_corners=True
+                )
+                x1_local_i = self.local_conv1(x1_trans_i)
+                x1_local_list.append(x1_local_i)
+
+        # ============== Block 2 ==============
+        # Block 2
+        # global branch
+        x2 = self.inception2(x1_out)
+        x2_attn, x2_theta = self.ha2(x2)
+        x2_out = x2 * x2_attn
+        # local branch
+        if self.learn_region:
+            x2_local_list = []
+            for region_idx in range(4):
+                x2_theta_i = x2_theta[:, region_idx, :]
+                x2_theta_i = self.transform_theta(x2_theta_i, region_idx)
+                x2_trans_i = self.stn(x1_out, x2_theta_i)
+                x2_trans_i = F.upsample(
+                    x2_trans_i, (12, 14), mode='bilinear', align_corners=True
+                )
+                x2_local_i = x2_trans_i + x1_local_list[region_idx]
+                x2_local_i = self.local_conv2(x2_local_i)
+                x2_local_list.append(x2_local_i)
+
+        # ============== Block 3 ==============
+        # Block 3
+        # global branch
+        x3 = self.inception3(x2_out)
+        x3_attn, x3_theta = self.ha3(x3)
+        x3_out = x3 * x3_attn
+        # local branch
+        if self.learn_region:
+            x3_local_list = []
+            for region_idx in range(4):
+                x3_theta_i = x3_theta[:, region_idx, :]
+                x3_theta_i = self.transform_theta(x3_theta_i, region_idx)
+                x3_trans_i = self.stn(x2_out, x3_theta_i)
+                x3_trans_i = F.upsample(
+                    x3_trans_i, (6, 7), mode='bilinear', align_corners=True
+                )
+                x3_local_i = x3_trans_i + x2_local_list[region_idx]
+                x3_local_i = self.local_conv3(x3_local_i)
+                x3_local_list.append(x3_local_i)
+
+        # ============== Feature generation ==============
+        # global branch
+        x_global = F.avg_pool2d(x3_out,
+                                x3_out.size()[2:]
+                                ).view(x3_out.size(0), x3_out.size(1))
+        x_global = self.fc_global(x_global)
+        # local branch
+        if self.learn_region:
+            x_local_list = []
+            for region_idx in range(4):
+                x_local_i = x3_local_list[region_idx]
+                x_local_i = F.avg_pool2d(x_local_i,
+                                         x_local_i.size()[2:]
+                                         ).view(x_local_i.size(0), -1)
+                x_local_list.append(x_local_i)
+            x_local = torch.cat(x_local_list, 1)
+            x_local = self.fc_local(x_local)
+
+        if not self.training:
+            # l2 normalization before concatenation
+            if self.learn_region:
+                x_global = x_global / x_global.norm(p=2, dim=1, keepdim=True)
+                x_local = x_local / x_local.norm(p=2, dim=1, keepdim=True)
+                return torch.cat([x_global, x_local], 1)
+            else:
+                return x_global
+
+        prelogits_global = self.classifier_global(x_global)
+        if self.learn_region:
+            prelogits_local = self.classifier_local(x_local)
+
+        if self.loss == 'softmax':
+            if self.learn_region:
+                return (prelogits_global, prelogits_local)
+            else:
+                return prelogits_global
+
+        elif self.loss == 'triplet':
+            if self.learn_region:
+                return (prelogits_global, prelogits_local), (x_global, x_local)
+            else:
+                return prelogits_global, x_global
+
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))

trackers/strongsort/deep/models/inceptionresnetv2.py

@@ -0,0 +1,361 @@
+"""
+Code imported from https://github.com/Cadene/pretrained-models.pytorch
+"""
+from __future__ import division, absolute_import
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['inceptionresnetv2']
+
+pretrained_settings = {
+    'inceptionresnetv2': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/inceptionresnetv2-520b38e4.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 299, 299],
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1000
+        },
+        'imagenet+background': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/inceptionresnetv2-520b38e4.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 299, 299],
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1001
+        }
+    }
+}
+
+
+class BasicConv2d(nn.Module):
+
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False
+        ) # verify bias false
+        self.bn = nn.BatchNorm2d(
+            out_planes,
+            eps=0.001, # value found in tensorflow
+            momentum=0.1, # default pytorch value
+            affine=True
+        )
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class Mixed_5b(nn.Module):
+
+    def __init__(self):
+        super(Mixed_5b, self).__init__()
+
+        self.branch0 = BasicConv2d(192, 96, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(192, 48, kernel_size=1, stride=1),
+            BasicConv2d(48, 64, kernel_size=5, stride=1, padding=2)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(192, 64, kernel_size=1, stride=1),
+            BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.branch3 = nn.Sequential(
+            nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),
+            BasicConv2d(192, 64, kernel_size=1, stride=1)
+        )
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class Block35(nn.Module):
+
+    def __init__(self, scale=1.0):
+        super(Block35, self).__init__()
+
+        self.scale = scale
+
+        self.branch0 = BasicConv2d(320, 32, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(320, 32, kernel_size=1, stride=1),
+            BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(320, 32, kernel_size=1, stride=1),
+            BasicConv2d(32, 48, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(48, 64, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.conv2d = nn.Conv2d(128, 320, kernel_size=1, stride=1)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        out = torch.cat((x0, x1, x2), 1)
+        out = self.conv2d(out)
+        out = out * self.scale + x
+        out = self.relu(out)
+        return out
+
+
+class Mixed_6a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_6a, self).__init__()
+
+        self.branch0 = BasicConv2d(320, 384, kernel_size=3, stride=2)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(320, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(256, 384, kernel_size=3, stride=2)
+        )
+
+        self.branch2 = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        out = torch.cat((x0, x1, x2), 1)
+        return out
+
+
+class Block17(nn.Module):
+
+    def __init__(self, scale=1.0):
+        super(Block17, self).__init__()
+
+        self.scale = scale
+
+        self.branch0 = BasicConv2d(1088, 192, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(1088, 128, kernel_size=1, stride=1),
+            BasicConv2d(
+                128, 160, kernel_size=(1, 7), stride=1, padding=(0, 3)
+            ),
+            BasicConv2d(
+                160, 192, kernel_size=(7, 1), stride=1, padding=(3, 0)
+            )
+        )
+
+        self.conv2d = nn.Conv2d(384, 1088, kernel_size=1, stride=1)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        out = torch.cat((x0, x1), 1)
+        out = self.conv2d(out)
+        out = out * self.scale + x
+        out = self.relu(out)
+        return out
+
+
+class Mixed_7a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_7a, self).__init__()
+
+        self.branch0 = nn.Sequential(
+            BasicConv2d(1088, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 384, kernel_size=3, stride=2)
+        )
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(1088, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 288, kernel_size=3, stride=2)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(1088, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 288, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(288, 320, kernel_size=3, stride=2)
+        )
+
+        self.branch3 = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class Block8(nn.Module):
+
+    def __init__(self, scale=1.0, noReLU=False):
+        super(Block8, self).__init__()
+
+        self.scale = scale
+        self.noReLU = noReLU
+
+        self.branch0 = BasicConv2d(2080, 192, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(2080, 192, kernel_size=1, stride=1),
+            BasicConv2d(
+                192, 224, kernel_size=(1, 3), stride=1, padding=(0, 1)
+            ),
+            BasicConv2d(
+                224, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)
+            )
+        )
+
+        self.conv2d = nn.Conv2d(448, 2080, kernel_size=1, stride=1)
+        if not self.noReLU:
+            self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        out = torch.cat((x0, x1), 1)
+        out = self.conv2d(out)
+        out = out * self.scale + x
+        if not self.noReLU:
+            out = self.relu(out)
+        return out
+
+
+# ----------------
+# Model Definition
+# ----------------
+class InceptionResNetV2(nn.Module):
+    """Inception-ResNet-V2.
+
+    Reference:
+        Szegedy et al. Inception-v4, Inception-ResNet and the Impact of Residual
+        Connections on Learning. AAAI 2017.
+
+    Public keys:
+        - ``inceptionresnetv2``: Inception-ResNet-V2.
+    """
+
+    def __init__(self, num_classes, loss='softmax', **kwargs):
+        super(InceptionResNetV2, self).__init__()
+        self.loss = loss
+
+        # Modules
+        self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2)
+        self.conv2d_2a = BasicConv2d(32, 32, kernel_size=3, stride=1)
+        self.conv2d_2b = BasicConv2d(
+            32, 64, kernel_size=3, stride=1, padding=1
+        )
+        self.maxpool_3a = nn.MaxPool2d(3, stride=2)
+        self.conv2d_3b = BasicConv2d(64, 80, kernel_size=1, stride=1)
+        self.conv2d_4a = BasicConv2d(80, 192, kernel_size=3, stride=1)
+        self.maxpool_5a = nn.MaxPool2d(3, stride=2)
+        self.mixed_5b = Mixed_5b()
+        self.repeat = nn.Sequential(
+            Block35(scale=0.17), Block35(scale=0.17), Block35(scale=0.17),
+            Block35(scale=0.17), Block35(scale=0.17), Block35(scale=0.17),
+            Block35(scale=0.17), Block35(scale=0.17), Block35(scale=0.17),
+            Block35(scale=0.17)
+        )
+        self.mixed_6a = Mixed_6a()
+        self.repeat_1 = nn.Sequential(
+            Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),
+            Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),
+            Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),
+            Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),
+            Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),
+            Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),
+            Block17(scale=0.10), Block17(scale=0.10)
+        )
+        self.mixed_7a = Mixed_7a()
+        self.repeat_2 = nn.Sequential(
+            Block8(scale=0.20), Block8(scale=0.20), Block8(scale=0.20),
+            Block8(scale=0.20), Block8(scale=0.20), Block8(scale=0.20),
+            Block8(scale=0.20), Block8(scale=0.20), Block8(scale=0.20)
+        )
+
+        self.block8 = Block8(noReLU=True)
+        self.conv2d_7b = BasicConv2d(2080, 1536, kernel_size=1, stride=1)
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.classifier = nn.Linear(1536, num_classes)
+
+    def load_imagenet_weights(self):
+        settings = pretrained_settings['inceptionresnetv2']['imagenet']
+        pretrain_dict = model_zoo.load_url(settings['url'])
+        model_dict = self.state_dict()
+        pretrain_dict = {
+            k: v
+            for k, v in pretrain_dict.items()
+            if k in model_dict and model_dict[k].size() == v.size()
+        }
+        model_dict.update(pretrain_dict)
+        self.load_state_dict(model_dict)
+
+    def featuremaps(self, x):
+        x = self.conv2d_1a(x)
+        x = self.conv2d_2a(x)
+        x = self.conv2d_2b(x)
+        x = self.maxpool_3a(x)
+        x = self.conv2d_3b(x)
+        x = self.conv2d_4a(x)
+        x = self.maxpool_5a(x)
+        x = self.mixed_5b(x)
+        x = self.repeat(x)
+        x = self.mixed_6a(x)
+        x = self.repeat_1(x)
+        x = self.mixed_7a(x)
+        x = self.repeat_2(x)
+        x = self.block8(x)
+        x = self.conv2d_7b(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def inceptionresnetv2(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = InceptionResNetV2(num_classes=num_classes, loss=loss, **kwargs)
+    if pretrained:
+        model.load_imagenet_weights()
+    return model

trackers/strongsort/deep/models/inceptionv4.py

@@ -0,0 +1,381 @@
+from __future__ import division, absolute_import
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['inceptionv4']
+"""
+Code imported from https://github.com/Cadene/pretrained-models.pytorch
+"""
+
+pretrained_settings = {
+    'inceptionv4': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/inceptionv4-8e4777a0.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 299, 299],
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1000
+        },
+        'imagenet+background': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/inceptionv4-8e4777a0.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 299, 299],
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1001
+        }
+    }
+}
+
+
+class BasicConv2d(nn.Module):
+
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False
+        ) # verify bias false
+        self.bn = nn.BatchNorm2d(
+            out_planes,
+            eps=0.001, # value found in tensorflow
+            momentum=0.1, # default pytorch value
+            affine=True
+        )
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class Mixed_3a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_3a, self).__init__()
+        self.maxpool = nn.MaxPool2d(3, stride=2)
+        self.conv = BasicConv2d(64, 96, kernel_size=3, stride=2)
+
+    def forward(self, x):
+        x0 = self.maxpool(x)
+        x1 = self.conv(x)
+        out = torch.cat((x0, x1), 1)
+        return out
+
+
+class Mixed_4a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_4a, self).__init__()
+
+        self.branch0 = nn.Sequential(
+            BasicConv2d(160, 64, kernel_size=1, stride=1),
+            BasicConv2d(64, 96, kernel_size=3, stride=1)
+        )
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(160, 64, kernel_size=1, stride=1),
+            BasicConv2d(64, 64, kernel_size=(1, 7), stride=1, padding=(0, 3)),
+            BasicConv2d(64, 64, kernel_size=(7, 1), stride=1, padding=(3, 0)),
+            BasicConv2d(64, 96, kernel_size=(3, 3), stride=1)
+        )
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        out = torch.cat((x0, x1), 1)
+        return out
+
+
+class Mixed_5a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_5a, self).__init__()
+        self.conv = BasicConv2d(192, 192, kernel_size=3, stride=2)
+        self.maxpool = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.conv(x)
+        x1 = self.maxpool(x)
+        out = torch.cat((x0, x1), 1)
+        return out
+
+
+class Inception_A(nn.Module):
+
+    def __init__(self):
+        super(Inception_A, self).__init__()
+        self.branch0 = BasicConv2d(384, 96, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(384, 64, kernel_size=1, stride=1),
+            BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(384, 64, kernel_size=1, stride=1),
+            BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.branch3 = nn.Sequential(
+            nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),
+            BasicConv2d(384, 96, kernel_size=1, stride=1)
+        )
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class Reduction_A(nn.Module):
+
+    def __init__(self):
+        super(Reduction_A, self).__init__()
+        self.branch0 = BasicConv2d(384, 384, kernel_size=3, stride=2)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(384, 192, kernel_size=1, stride=1),
+            BasicConv2d(192, 224, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(224, 256, kernel_size=3, stride=2)
+        )
+
+        self.branch2 = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        out = torch.cat((x0, x1, x2), 1)
+        return out
+
+
+class Inception_B(nn.Module):
+
+    def __init__(self):
+        super(Inception_B, self).__init__()
+        self.branch0 = BasicConv2d(1024, 384, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(1024, 192, kernel_size=1, stride=1),
+            BasicConv2d(
+                192, 224, kernel_size=(1, 7), stride=1, padding=(0, 3)
+            ),
+            BasicConv2d(
+                224, 256, kernel_size=(7, 1), stride=1, padding=(3, 0)
+            )
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(1024, 192, kernel_size=1, stride=1),
+            BasicConv2d(
+                192, 192, kernel_size=(7, 1), stride=1, padding=(3, 0)
+            ),
+            BasicConv2d(
+                192, 224, kernel_size=(1, 7), stride=1, padding=(0, 3)
+            ),
+            BasicConv2d(
+                224, 224, kernel_size=(7, 1), stride=1, padding=(3, 0)
+            ),
+            BasicConv2d(
+                224, 256, kernel_size=(1, 7), stride=1, padding=(0, 3)
+            )
+        )
+
+        self.branch3 = nn.Sequential(
+            nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),
+            BasicConv2d(1024, 128, kernel_size=1, stride=1)
+        )
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class Reduction_B(nn.Module):
+
+    def __init__(self):
+        super(Reduction_B, self).__init__()
+
+        self.branch0 = nn.Sequential(
+            BasicConv2d(1024, 192, kernel_size=1, stride=1),
+            BasicConv2d(192, 192, kernel_size=3, stride=2)
+        )
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(1024, 256, kernel_size=1, stride=1),
+            BasicConv2d(
+                256, 256, kernel_size=(1, 7), stride=1, padding=(0, 3)
+            ),
+            BasicConv2d(
+                256, 320, kernel_size=(7, 1), stride=1, padding=(3, 0)
+            ), BasicConv2d(320, 320, kernel_size=3, stride=2)
+        )
+
+        self.branch2 = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        out = torch.cat((x0, x1, x2), 1)
+        return out
+
+
+class Inception_C(nn.Module):
+
+    def __init__(self):
+        super(Inception_C, self).__init__()
+
+        self.branch0 = BasicConv2d(1536, 256, kernel_size=1, stride=1)
+
+        self.branch1_0 = BasicConv2d(1536, 384, kernel_size=1, stride=1)
+        self.branch1_1a = BasicConv2d(
+            384, 256, kernel_size=(1, 3), stride=1, padding=(0, 1)
+        )
+        self.branch1_1b = BasicConv2d(
+            384, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)
+        )
+
+        self.branch2_0 = BasicConv2d(1536, 384, kernel_size=1, stride=1)
+        self.branch2_1 = BasicConv2d(
+            384, 448, kernel_size=(3, 1), stride=1, padding=(1, 0)
+        )
+        self.branch2_2 = BasicConv2d(
+            448, 512, kernel_size=(1, 3), stride=1, padding=(0, 1)
+        )
+        self.branch2_3a = BasicConv2d(
+            512, 256, kernel_size=(1, 3), stride=1, padding=(0, 1)
+        )
+        self.branch2_3b = BasicConv2d(
+            512, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)
+        )
+
+        self.branch3 = nn.Sequential(
+            nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),
+            BasicConv2d(1536, 256, kernel_size=1, stride=1)
+        )
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+
+        x1_0 = self.branch1_0(x)
+        x1_1a = self.branch1_1a(x1_0)
+        x1_1b = self.branch1_1b(x1_0)
+        x1 = torch.cat((x1_1a, x1_1b), 1)
+
+        x2_0 = self.branch2_0(x)
+        x2_1 = self.branch2_1(x2_0)
+        x2_2 = self.branch2_2(x2_1)
+        x2_3a = self.branch2_3a(x2_2)
+        x2_3b = self.branch2_3b(x2_2)
+        x2 = torch.cat((x2_3a, x2_3b), 1)
+
+        x3 = self.branch3(x)
+
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class InceptionV4(nn.Module):
+    """Inception-v4.
+
+    Reference:
+        Szegedy et al. Inception-v4, Inception-ResNet and the Impact of Residual
+        Connections on Learning. AAAI 2017.
+
+    Public keys:
+        - ``inceptionv4``: InceptionV4.
+    """
+
+    def __init__(self, num_classes, loss, **kwargs):
+        super(InceptionV4, self).__init__()
+        self.loss = loss
+
+        self.features = nn.Sequential(
+            BasicConv2d(3, 32, kernel_size=3, stride=2),
+            BasicConv2d(32, 32, kernel_size=3, stride=1),
+            BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1),
+            Mixed_3a(),
+            Mixed_4a(),
+            Mixed_5a(),
+            Inception_A(),
+            Inception_A(),
+            Inception_A(),
+            Inception_A(),
+            Reduction_A(), # Mixed_6a
+            Inception_B(),
+            Inception_B(),
+            Inception_B(),
+            Inception_B(),
+            Inception_B(),
+            Inception_B(),
+            Inception_B(),
+            Reduction_B(), # Mixed_7a
+            Inception_C(),
+            Inception_C(),
+            Inception_C()
+        )
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.classifier = nn.Linear(1536, num_classes)
+
+    def forward(self, x):
+        f = self.features(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def inceptionv4(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = InceptionV4(num_classes, loss, **kwargs)
+    if pretrained:
+        model_url = pretrained_settings['inceptionv4']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model

trackers/strongsort/deep/models/mlfn.py

@@ -0,0 +1,269 @@
+from __future__ import division, absolute_import
+import torch
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ['mlfn']
+
+model_urls = {
+    # training epoch = 5, top1 = 51.6
+    'imagenet':
+    'https://mega.nz/#!YHxAhaxC!yu9E6zWl0x5zscSouTdbZu8gdFFytDdl-RAdD2DEfpk',
+}
+
+
+class MLFNBlock(nn.Module):
+
+    def __init__(
+        self, in_channels, out_channels, stride, fsm_channels, groups=32
+    ):
+        super(MLFNBlock, self).__init__()
+        self.groups = groups
+        mid_channels = out_channels // 2
+
+        # Factor Modules
+        self.fm_conv1 = nn.Conv2d(in_channels, mid_channels, 1, bias=False)
+        self.fm_bn1 = nn.BatchNorm2d(mid_channels)
+        self.fm_conv2 = nn.Conv2d(
+            mid_channels,
+            mid_channels,
+            3,
+            stride=stride,
+            padding=1,
+            bias=False,
+            groups=self.groups
+        )
+        self.fm_bn2 = nn.BatchNorm2d(mid_channels)
+        self.fm_conv3 = nn.Conv2d(mid_channels, out_channels, 1, bias=False)
+        self.fm_bn3 = nn.BatchNorm2d(out_channels)
+
+        # Factor Selection Module
+        self.fsm = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, fsm_channels[0], 1),
+            nn.BatchNorm2d(fsm_channels[0]),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(fsm_channels[0], fsm_channels[1], 1),
+            nn.BatchNorm2d(fsm_channels[1]),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(fsm_channels[1], self.groups, 1),
+            nn.BatchNorm2d(self.groups),
+            nn.Sigmoid(),
+        )
+
+        self.downsample = None
+        if in_channels != out_channels or stride > 1:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(
+                    in_channels, out_channels, 1, stride=stride, bias=False
+                ),
+                nn.BatchNorm2d(out_channels),
+            )
+
+    def forward(self, x):
+        residual = x
+        s = self.fsm(x)
+
+        # reduce dimension
+        x = self.fm_conv1(x)
+        x = self.fm_bn1(x)
+        x = F.relu(x, inplace=True)
+
+        # group convolution
+        x = self.fm_conv2(x)
+        x = self.fm_bn2(x)
+        x = F.relu(x, inplace=True)
+
+        # factor selection
+        b, c = x.size(0), x.size(1)
+        n = c // self.groups
+        ss = s.repeat(1, n, 1, 1) # from (b, g, 1, 1) to (b, g*n=c, 1, 1)
+        ss = ss.view(b, n, self.groups, 1, 1)
+        ss = ss.permute(0, 2, 1, 3, 4).contiguous()
+        ss = ss.view(b, c, 1, 1)
+        x = ss * x
+
+        # recover dimension
+        x = self.fm_conv3(x)
+        x = self.fm_bn3(x)
+        x = F.relu(x, inplace=True)
+
+        if self.downsample is not None:
+            residual = self.downsample(residual)
+
+        return F.relu(residual + x, inplace=True), s
+
+
+class MLFN(nn.Module):
+    """Multi-Level Factorisation Net.
+
+    Reference:
+        Chang et al. Multi-Level Factorisation Net for
+        Person Re-Identification. CVPR 2018.
+
+    Public keys:
+        - ``mlfn``: MLFN (Multi-Level Factorisation Net).
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss='softmax',
+        groups=32,
+        channels=[64, 256, 512, 1024, 2048],
+        embed_dim=1024,
+        **kwargs
+    ):
+        super(MLFN, self).__init__()
+        self.loss = loss
+        self.groups = groups
+
+        # first convolutional layer
+        self.conv1 = nn.Conv2d(3, channels[0], 7, stride=2, padding=3)
+        self.bn1 = nn.BatchNorm2d(channels[0])
+        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
+
+        # main body
+        self.feature = nn.ModuleList(
+            [
+                # layer 1-3
+                MLFNBlock(channels[0], channels[1], 1, [128, 64], self.groups),
+                MLFNBlock(channels[1], channels[1], 1, [128, 64], self.groups),
+                MLFNBlock(channels[1], channels[1], 1, [128, 64], self.groups),
+                # layer 4-7
+                MLFNBlock(
+                    channels[1], channels[2], 2, [256, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[2], channels[2], 1, [256, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[2], channels[2], 1, [256, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[2], channels[2], 1, [256, 128], self.groups
+                ),
+                # layer 8-13
+                MLFNBlock(
+                    channels[2], channels[3], 2, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[3], channels[3], 1, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[3], channels[3], 1, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[3], channels[3], 1, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[3], channels[3], 1, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[3], channels[3], 1, [512, 128], self.groups
+                ),
+                # layer 14-16
+                MLFNBlock(
+                    channels[3], channels[4], 2, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[4], channels[4], 1, [512, 128], self.groups
+                ),
+                MLFNBlock(
+                    channels[4], channels[4], 1, [512, 128], self.groups
+                ),
+            ]
+        )
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+
+        # projection functions
+        self.fc_x = nn.Sequential(
+            nn.Conv2d(channels[4], embed_dim, 1, bias=False),
+            nn.BatchNorm2d(embed_dim),
+            nn.ReLU(inplace=True),
+        )
+        self.fc_s = nn.Sequential(
+            nn.Conv2d(self.groups * 16, embed_dim, 1, bias=False),
+            nn.BatchNorm2d(embed_dim),
+            nn.ReLU(inplace=True),
+        )
+
+        self.classifier = nn.Linear(embed_dim, num_classes)
+
+        self.init_params()
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = F.relu(x, inplace=True)
+        x = self.maxpool(x)
+
+        s_hat = []
+        for block in self.feature:
+            x, s = block(x)
+            s_hat.append(s)
+        s_hat = torch.cat(s_hat, 1)
+
+        x = self.global_avgpool(x)
+        x = self.fc_x(x)
+        s_hat = self.fc_s(s_hat)
+
+        v = (x+s_hat) * 0.5
+        v = v.view(v.size(0), -1)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def mlfn(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = MLFN(num_classes, loss, **kwargs)
+    if pretrained:
+        # init_pretrained_weights(model, model_urls['imagenet'])
+        import warnings
+        warnings.warn(
+            'The imagenet pretrained weights need to be manually downloaded from {}'
+            .format(model_urls['imagenet'])
+        )
+    return model

trackers/strongsort/deep/models/mobilenetv2.py

@@ -0,0 +1,274 @@
+from __future__ import division, absolute_import
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ['mobilenetv2_x1_0', 'mobilenetv2_x1_4']
+
+model_urls = {
+    # 1.0: top-1 71.3
+    'mobilenetv2_x1_0':
+    'https://mega.nz/#!NKp2wAIA!1NH1pbNzY_M2hVk_hdsxNM1NUOWvvGPHhaNr-fASF6c',
+    # 1.4: top-1 73.9
+    'mobilenetv2_x1_4':
+    'https://mega.nz/#!RGhgEIwS!xN2s2ZdyqI6vQ3EwgmRXLEW3khr9tpXg96G9SUJugGk',
+}
+
+
+class ConvBlock(nn.Module):
+    """Basic convolutional block.
+    
+    convolution (bias discarded) + batch normalization + relu6.
+
+    Args:
+        in_c (int): number of input channels.
+        out_c (int): number of output channels.
+        k (int or tuple): kernel size.
+        s (int or tuple): stride.
+        p (int or tuple): padding.
+        g (int): number of blocked connections from input channels
+            to output channels (default: 1).
+    """
+
+    def __init__(self, in_c, out_c, k, s=1, p=0, g=1):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Conv2d(
+            in_c, out_c, k, stride=s, padding=p, bias=False, groups=g
+        )
+        self.bn = nn.BatchNorm2d(out_c)
+
+    def forward(self, x):
+        return F.relu6(self.bn(self.conv(x)))
+
+
+class Bottleneck(nn.Module):
+
+    def __init__(self, in_channels, out_channels, expansion_factor, stride=1):
+        super(Bottleneck, self).__init__()
+        mid_channels = in_channels * expansion_factor
+        self.use_residual = stride == 1 and in_channels == out_channels
+        self.conv1 = ConvBlock(in_channels, mid_channels, 1)
+        self.dwconv2 = ConvBlock(
+            mid_channels, mid_channels, 3, stride, 1, g=mid_channels
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(mid_channels, out_channels, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+        )
+
+    def forward(self, x):
+        m = self.conv1(x)
+        m = self.dwconv2(m)
+        m = self.conv3(m)
+        if self.use_residual:
+            return x + m
+        else:
+            return m
+
+
+class MobileNetV2(nn.Module):
+    """MobileNetV2.
+
+    Reference:
+        Sandler et al. MobileNetV2: Inverted Residuals and
+        Linear Bottlenecks. CVPR 2018.
+
+    Public keys:
+        - ``mobilenetv2_x1_0``: MobileNetV2 x1.0.
+        - ``mobilenetv2_x1_4``: MobileNetV2 x1.4.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        width_mult=1,
+        loss='softmax',
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    ):
+        super(MobileNetV2, self).__init__()
+        self.loss = loss
+        self.in_channels = int(32 * width_mult)
+        self.feature_dim = int(1280 * width_mult) if width_mult > 1 else 1280
+
+        # construct layers
+        self.conv1 = ConvBlock(3, self.in_channels, 3, s=2, p=1)
+        self.conv2 = self._make_layer(
+            Bottleneck, 1, int(16 * width_mult), 1, 1
+        )
+        self.conv3 = self._make_layer(
+            Bottleneck, 6, int(24 * width_mult), 2, 2
+        )
+        self.conv4 = self._make_layer(
+            Bottleneck, 6, int(32 * width_mult), 3, 2
+        )
+        self.conv5 = self._make_layer(
+            Bottleneck, 6, int(64 * width_mult), 4, 2
+        )
+        self.conv6 = self._make_layer(
+            Bottleneck, 6, int(96 * width_mult), 3, 1
+        )
+        self.conv7 = self._make_layer(
+            Bottleneck, 6, int(160 * width_mult), 3, 2
+        )
+        self.conv8 = self._make_layer(
+            Bottleneck, 6, int(320 * width_mult), 1, 1
+        )
+        self.conv9 = ConvBlock(self.in_channels, self.feature_dim, 1)
+
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = self._construct_fc_layer(
+            fc_dims, self.feature_dim, dropout_p
+        )
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _make_layer(self, block, t, c, n, s):
+        # t: expansion factor
+        # c: output channels
+        # n: number of blocks
+        # s: stride for first layer
+        layers = []
+        layers.append(block(self.in_channels, c, t, s))
+        self.in_channels = c
+        for i in range(1, n):
+            layers.append(block(self.in_channels, c, t))
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer.
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = self.conv5(x)
+        x = self.conv6(x)
+        x = self.conv7(x)
+        x = self.conv8(x)
+        x = self.conv9(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if self.fc is not None:
+            v = self.fc(v)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def mobilenetv2_x1_0(num_classes, loss, pretrained=True, **kwargs):
+    model = MobileNetV2(
+        num_classes,
+        loss=loss,
+        width_mult=1,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        # init_pretrained_weights(model, model_urls['mobilenetv2_x1_0'])
+        import warnings
+        warnings.warn(
+            'The imagenet pretrained weights need to be manually downloaded from {}'
+            .format(model_urls['mobilenetv2_x1_0'])
+        )
+    return model
+
+
+def mobilenetv2_x1_4(num_classes, loss, pretrained=True, **kwargs):
+    model = MobileNetV2(
+        num_classes,
+        loss=loss,
+        width_mult=1.4,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        # init_pretrained_weights(model, model_urls['mobilenetv2_x1_4'])
+        import warnings
+        warnings.warn(
+            'The imagenet pretrained weights need to be manually downloaded from {}'
+            .format(model_urls['mobilenetv2_x1_4'])
+        )
+    return model

trackers/strongsort/deep/models/mudeep.py

@@ -0,0 +1,206 @@
+from __future__ import division, absolute_import
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ['MuDeep']
+
+
+class ConvBlock(nn.Module):
+    """Basic convolutional block.
+    
+    convolution + batch normalization + relu.
+
+    Args:
+        in_c (int): number of input channels.
+        out_c (int): number of output channels.
+        k (int or tuple): kernel size.
+        s (int or tuple): stride.
+        p (int or tuple): padding.
+    """
+
+    def __init__(self, in_c, out_c, k, s, p):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p)
+        self.bn = nn.BatchNorm2d(out_c)
+
+    def forward(self, x):
+        return F.relu(self.bn(self.conv(x)))
+
+
+class ConvLayers(nn.Module):
+    """Preprocessing layers."""
+
+    def __init__(self):
+        super(ConvLayers, self).__init__()
+        self.conv1 = ConvBlock(3, 48, k=3, s=1, p=1)
+        self.conv2 = ConvBlock(48, 96, k=3, s=1, p=1)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.maxpool(x)
+        return x
+
+
+class MultiScaleA(nn.Module):
+    """Multi-scale stream layer A (Sec.3.1)"""
+
+    def __init__(self):
+        super(MultiScaleA, self).__init__()
+        self.stream1 = nn.Sequential(
+            ConvBlock(96, 96, k=1, s=1, p=0),
+            ConvBlock(96, 24, k=3, s=1, p=1),
+        )
+        self.stream2 = nn.Sequential(
+            nn.AvgPool2d(kernel_size=3, stride=1, padding=1),
+            ConvBlock(96, 24, k=1, s=1, p=0),
+        )
+        self.stream3 = ConvBlock(96, 24, k=1, s=1, p=0)
+        self.stream4 = nn.Sequential(
+            ConvBlock(96, 16, k=1, s=1, p=0),
+            ConvBlock(16, 24, k=3, s=1, p=1),
+            ConvBlock(24, 24, k=3, s=1, p=1),
+        )
+
+    def forward(self, x):
+        s1 = self.stream1(x)
+        s2 = self.stream2(x)
+        s3 = self.stream3(x)
+        s4 = self.stream4(x)
+        y = torch.cat([s1, s2, s3, s4], dim=1)
+        return y
+
+
+class Reduction(nn.Module):
+    """Reduction layer (Sec.3.1)"""
+
+    def __init__(self):
+        super(Reduction, self).__init__()
+        self.stream1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.stream2 = ConvBlock(96, 96, k=3, s=2, p=1)
+        self.stream3 = nn.Sequential(
+            ConvBlock(96, 48, k=1, s=1, p=0),
+            ConvBlock(48, 56, k=3, s=1, p=1),
+            ConvBlock(56, 64, k=3, s=2, p=1),
+        )
+
+    def forward(self, x):
+        s1 = self.stream1(x)
+        s2 = self.stream2(x)
+        s3 = self.stream3(x)
+        y = torch.cat([s1, s2, s3], dim=1)
+        return y
+
+
+class MultiScaleB(nn.Module):
+    """Multi-scale stream layer B (Sec.3.1)"""
+
+    def __init__(self):
+        super(MultiScaleB, self).__init__()
+        self.stream1 = nn.Sequential(
+            nn.AvgPool2d(kernel_size=3, stride=1, padding=1),
+            ConvBlock(256, 256, k=1, s=1, p=0),
+        )
+        self.stream2 = nn.Sequential(
+            ConvBlock(256, 64, k=1, s=1, p=0),
+            ConvBlock(64, 128, k=(1, 3), s=1, p=(0, 1)),
+            ConvBlock(128, 256, k=(3, 1), s=1, p=(1, 0)),
+        )
+        self.stream3 = ConvBlock(256, 256, k=1, s=1, p=0)
+        self.stream4 = nn.Sequential(
+            ConvBlock(256, 64, k=1, s=1, p=0),
+            ConvBlock(64, 64, k=(1, 3), s=1, p=(0, 1)),
+            ConvBlock(64, 128, k=(3, 1), s=1, p=(1, 0)),
+            ConvBlock(128, 128, k=(1, 3), s=1, p=(0, 1)),
+            ConvBlock(128, 256, k=(3, 1), s=1, p=(1, 0)),
+        )
+
+    def forward(self, x):
+        s1 = self.stream1(x)
+        s2 = self.stream2(x)
+        s3 = self.stream3(x)
+        s4 = self.stream4(x)
+        return s1, s2, s3, s4
+
+
+class Fusion(nn.Module):
+    """Saliency-based learning fusion layer (Sec.3.2)"""
+
+    def __init__(self):
+        super(Fusion, self).__init__()
+        self.a1 = nn.Parameter(torch.rand(1, 256, 1, 1))
+        self.a2 = nn.Parameter(torch.rand(1, 256, 1, 1))
+        self.a3 = nn.Parameter(torch.rand(1, 256, 1, 1))
+        self.a4 = nn.Parameter(torch.rand(1, 256, 1, 1))
+
+        # We add an average pooling layer to reduce the spatial dimension
+        # of feature maps, which differs from the original paper.
+        self.avgpool = nn.AvgPool2d(kernel_size=4, stride=4, padding=0)
+
+    def forward(self, x1, x2, x3, x4):
+        s1 = self.a1.expand_as(x1) * x1
+        s2 = self.a2.expand_as(x2) * x2
+        s3 = self.a3.expand_as(x3) * x3
+        s4 = self.a4.expand_as(x4) * x4
+        y = self.avgpool(s1 + s2 + s3 + s4)
+        return y
+
+
+class MuDeep(nn.Module):
+    """Multiscale deep neural network.
+
+    Reference:
+        Qian et al. Multi-scale Deep Learning Architectures
+        for Person Re-identification. ICCV 2017.
+
+    Public keys:
+        - ``mudeep``: Multiscale deep neural network.
+    """
+
+    def __init__(self, num_classes, loss='softmax', **kwargs):
+        super(MuDeep, self).__init__()
+        self.loss = loss
+
+        self.block1 = ConvLayers()
+        self.block2 = MultiScaleA()
+        self.block3 = Reduction()
+        self.block4 = MultiScaleB()
+        self.block5 = Fusion()
+
+        # Due to this fully connected layer, input image has to be fixed
+        # in shape, i.e. (3, 256, 128), such that the last convolutional feature
+        # maps are of shape (256, 16, 8). If input shape is changed,
+        # the input dimension of this layer has to be changed accordingly.
+        self.fc = nn.Sequential(
+            nn.Linear(256 * 16 * 8, 4096),
+            nn.BatchNorm1d(4096),
+            nn.ReLU(),
+        )
+        self.classifier = nn.Linear(4096, num_classes)
+        self.feat_dim = 4096
+
+    def featuremaps(self, x):
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        x = self.block4(x)
+        x = self.block5(*x)
+        return x
+
+    def forward(self, x):
+        x = self.featuremaps(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        y = self.classifier(x)
+
+        if not self.training:
+            return x
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, x
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))

trackers/strongsort/deep/models/nasnet.py

@@ -0,0 +1,1131 @@
+from __future__ import division, absolute_import
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['nasnetamobile']
+"""
+NASNet Mobile
+Thanks to Anastasiia (https://github.com/DagnyT) for the great help, support and motivation!
+
+
+------------------------------------------------------------------------------------
+      Architecture       | Top-1 Acc | Top-5 Acc |  Multiply-Adds |  Params (M)
+------------------------------------------------------------------------------------
+|   NASNet-A (4 @ 1056)  |   74.08%  |   91.74%  |       564 M    |     5.3        |
+------------------------------------------------------------------------------------
+# References:
+ - [Learning Transferable Architectures for Scalable Image Recognition]
+    (https://arxiv.org/abs/1707.07012)
+"""
+"""
+Code imported from https://github.com/Cadene/pretrained-models.pytorch
+"""
+
+pretrained_settings = {
+    'nasnetamobile': {
+        'imagenet': {
+            # 'url': 'https://github.com/veronikayurchuk/pretrained-models.pytorch/releases/download/v1.0/nasnetmobile-7e03cead.pth.tar',
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/nasnetamobile-7e03cead.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224], # resize 256
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1000
+        },
+        # 'imagenet+background': {
+        #     # 'url': 'http://data.lip6.fr/cadene/pretrainedmodels/nasnetalarge-a1897284.pth',
+        #     'input_space': 'RGB',
+        #     'input_size': [3, 224, 224], # resize 256
+        #     'input_range': [0, 1],
+        #     'mean': [0.5, 0.5, 0.5],
+        #     'std': [0.5, 0.5, 0.5],
+        #     'num_classes': 1001
+        # }
+    }
+}
+
+
+class MaxPoolPad(nn.Module):
+
+    def __init__(self):
+        super(MaxPoolPad, self).__init__()
+        self.pad = nn.ZeroPad2d((1, 0, 1, 0))
+        self.pool = nn.MaxPool2d(3, stride=2, padding=1)
+
+    def forward(self, x):
+        x = self.pad(x)
+        x = self.pool(x)
+        x = x[:, :, 1:, 1:].contiguous()
+        return x
+
+
+class AvgPoolPad(nn.Module):
+
+    def __init__(self, stride=2, padding=1):
+        super(AvgPoolPad, self).__init__()
+        self.pad = nn.ZeroPad2d((1, 0, 1, 0))
+        self.pool = nn.AvgPool2d(
+            3, stride=stride, padding=padding, count_include_pad=False
+        )
+
+    def forward(self, x):
+        x = self.pad(x)
+        x = self.pool(x)
+        x = x[:, :, 1:, 1:].contiguous()
+        return x
+
+
+class SeparableConv2d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        dw_kernel,
+        dw_stride,
+        dw_padding,
+        bias=False
+    ):
+        super(SeparableConv2d, self).__init__()
+        self.depthwise_conv2d = nn.Conv2d(
+            in_channels,
+            in_channels,
+            dw_kernel,
+            stride=dw_stride,
+            padding=dw_padding,
+            bias=bias,
+            groups=in_channels
+        )
+        self.pointwise_conv2d = nn.Conv2d(
+            in_channels, out_channels, 1, stride=1, bias=bias
+        )
+
+    def forward(self, x):
+        x = self.depthwise_conv2d(x)
+        x = self.pointwise_conv2d(x)
+        return x
+
+
+class BranchSeparables(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        name=None,
+        bias=False
+    ):
+        super(BranchSeparables, self).__init__()
+        self.relu = nn.ReLU()
+        self.separable_1 = SeparableConv2d(
+            in_channels, in_channels, kernel_size, stride, padding, bias=bias
+        )
+        self.bn_sep_1 = nn.BatchNorm2d(
+            in_channels, eps=0.001, momentum=0.1, affine=True
+        )
+        self.relu1 = nn.ReLU()
+        self.separable_2 = SeparableConv2d(
+            in_channels, out_channels, kernel_size, 1, padding, bias=bias
+        )
+        self.bn_sep_2 = nn.BatchNorm2d(
+            out_channels, eps=0.001, momentum=0.1, affine=True
+        )
+        self.name = name
+
+    def forward(self, x):
+        x = self.relu(x)
+        if self.name == 'specific':
+            x = nn.ZeroPad2d((1, 0, 1, 0))(x)
+        x = self.separable_1(x)
+        if self.name == 'specific':
+            x = x[:, :, 1:, 1:].contiguous()
+
+        x = self.bn_sep_1(x)
+        x = self.relu1(x)
+        x = self.separable_2(x)
+        x = self.bn_sep_2(x)
+        return x
+
+
+class BranchSeparablesStem(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        bias=False
+    ):
+        super(BranchSeparablesStem, self).__init__()
+        self.relu = nn.ReLU()
+        self.separable_1 = SeparableConv2d(
+            in_channels, out_channels, kernel_size, stride, padding, bias=bias
+        )
+        self.bn_sep_1 = nn.BatchNorm2d(
+            out_channels, eps=0.001, momentum=0.1, affine=True
+        )
+        self.relu1 = nn.ReLU()
+        self.separable_2 = SeparableConv2d(
+            out_channels, out_channels, kernel_size, 1, padding, bias=bias
+        )
+        self.bn_sep_2 = nn.BatchNorm2d(
+            out_channels, eps=0.001, momentum=0.1, affine=True
+        )
+
+    def forward(self, x):
+        x = self.relu(x)
+        x = self.separable_1(x)
+        x = self.bn_sep_1(x)
+        x = self.relu1(x)
+        x = self.separable_2(x)
+        x = self.bn_sep_2(x)
+        return x
+
+
+class BranchSeparablesReduction(BranchSeparables):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        z_padding=1,
+        bias=False
+    ):
+        BranchSeparables.__init__(
+            self, in_channels, out_channels, kernel_size, stride, padding, bias
+        )
+        self.padding = nn.ZeroPad2d((z_padding, 0, z_padding, 0))
+
+    def forward(self, x):
+        x = self.relu(x)
+        x = self.padding(x)
+        x = self.separable_1(x)
+        x = x[:, :, 1:, 1:].contiguous()
+        x = self.bn_sep_1(x)
+        x = self.relu1(x)
+        x = self.separable_2(x)
+        x = self.bn_sep_2(x)
+        return x
+
+
+class CellStem0(nn.Module):
+
+    def __init__(self, stem_filters, num_filters=42):
+        super(CellStem0, self).__init__()
+        self.num_filters = num_filters
+        self.stem_filters = stem_filters
+        self.conv_1x1 = nn.Sequential()
+        self.conv_1x1.add_module('relu', nn.ReLU())
+        self.conv_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                self.stem_filters, self.num_filters, 1, stride=1, bias=False
+            )
+        )
+        self.conv_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                self.num_filters, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.comb_iter_0_left = BranchSeparables(
+            self.num_filters, self.num_filters, 5, 2, 2
+        )
+        self.comb_iter_0_right = BranchSeparablesStem(
+            self.stem_filters, self.num_filters, 7, 2, 3, bias=False
+        )
+
+        self.comb_iter_1_left = nn.MaxPool2d(3, stride=2, padding=1)
+        self.comb_iter_1_right = BranchSeparablesStem(
+            self.stem_filters, self.num_filters, 7, 2, 3, bias=False
+        )
+
+        self.comb_iter_2_left = nn.AvgPool2d(
+            3, stride=2, padding=1, count_include_pad=False
+        )
+        self.comb_iter_2_right = BranchSeparablesStem(
+            self.stem_filters, self.num_filters, 5, 2, 2, bias=False
+        )
+
+        self.comb_iter_3_right = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_4_left = BranchSeparables(
+            self.num_filters, self.num_filters, 3, 1, 1, bias=False
+        )
+        self.comb_iter_4_right = nn.MaxPool2d(3, stride=2, padding=1)
+
+    def forward(self, x):
+        x1 = self.conv_1x1(x)
+
+        x_comb_iter_0_left = self.comb_iter_0_left(x1)
+        x_comb_iter_0_right = self.comb_iter_0_right(x)
+        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right
+
+        x_comb_iter_1_left = self.comb_iter_1_left(x1)
+        x_comb_iter_1_right = self.comb_iter_1_right(x)
+        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right
+
+        x_comb_iter_2_left = self.comb_iter_2_left(x1)
+        x_comb_iter_2_right = self.comb_iter_2_right(x)
+        x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right
+
+        x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)
+        x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1
+
+        x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)
+        x_comb_iter_4_right = self.comb_iter_4_right(x1)
+        x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right
+
+        x_out = torch.cat(
+            [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1
+        )
+        return x_out
+
+
+class CellStem1(nn.Module):
+
+    def __init__(self, stem_filters, num_filters):
+        super(CellStem1, self).__init__()
+        self.num_filters = num_filters
+        self.stem_filters = stem_filters
+        self.conv_1x1 = nn.Sequential()
+        self.conv_1x1.add_module('relu', nn.ReLU())
+        self.conv_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                2 * self.num_filters,
+                self.num_filters,
+                1,
+                stride=1,
+                bias=False
+            )
+        )
+        self.conv_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                self.num_filters, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.relu = nn.ReLU()
+        self.path_1 = nn.Sequential()
+        self.path_1.add_module(
+            'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)
+        )
+        self.path_1.add_module(
+            'conv',
+            nn.Conv2d(
+                self.stem_filters,
+                self.num_filters // 2,
+                1,
+                stride=1,
+                bias=False
+            )
+        )
+        self.path_2 = nn.ModuleList()
+        self.path_2.add_module('pad', nn.ZeroPad2d((0, 1, 0, 1)))
+        self.path_2.add_module(
+            'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)
+        )
+        self.path_2.add_module(
+            'conv',
+            nn.Conv2d(
+                self.stem_filters,
+                self.num_filters // 2,
+                1,
+                stride=1,
+                bias=False
+            )
+        )
+
+        self.final_path_bn = nn.BatchNorm2d(
+            self.num_filters, eps=0.001, momentum=0.1, affine=True
+        )
+
+        self.comb_iter_0_left = BranchSeparables(
+            self.num_filters,
+            self.num_filters,
+            5,
+            2,
+            2,
+            name='specific',
+            bias=False
+        )
+        self.comb_iter_0_right = BranchSeparables(
+            self.num_filters,
+            self.num_filters,
+            7,
+            2,
+            3,
+            name='specific',
+            bias=False
+        )
+
+        # self.comb_iter_1_left = nn.MaxPool2d(3, stride=2, padding=1)
+        self.comb_iter_1_left = MaxPoolPad()
+        self.comb_iter_1_right = BranchSeparables(
+            self.num_filters,
+            self.num_filters,
+            7,
+            2,
+            3,
+            name='specific',
+            bias=False
+        )
+
+        # self.comb_iter_2_left = nn.AvgPool2d(3, stride=2, padding=1, count_include_pad=False)
+        self.comb_iter_2_left = AvgPoolPad()
+        self.comb_iter_2_right = BranchSeparables(
+            self.num_filters,
+            self.num_filters,
+            5,
+            2,
+            2,
+            name='specific',
+            bias=False
+        )
+
+        self.comb_iter_3_right = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_4_left = BranchSeparables(
+            self.num_filters,
+            self.num_filters,
+            3,
+            1,
+            1,
+            name='specific',
+            bias=False
+        )
+        # self.comb_iter_4_right = nn.MaxPool2d(3, stride=2, padding=1)
+        self.comb_iter_4_right = MaxPoolPad()
+
+    def forward(self, x_conv0, x_stem_0):
+        x_left = self.conv_1x1(x_stem_0)
+
+        x_relu = self.relu(x_conv0)
+        # path 1
+        x_path1 = self.path_1(x_relu)
+        # path 2
+        x_path2 = self.path_2.pad(x_relu)
+        x_path2 = x_path2[:, :, 1:, 1:]
+        x_path2 = self.path_2.avgpool(x_path2)
+        x_path2 = self.path_2.conv(x_path2)
+        # final path
+        x_right = self.final_path_bn(torch.cat([x_path1, x_path2], 1))
+
+        x_comb_iter_0_left = self.comb_iter_0_left(x_left)
+        x_comb_iter_0_right = self.comb_iter_0_right(x_right)
+        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right
+
+        x_comb_iter_1_left = self.comb_iter_1_left(x_left)
+        x_comb_iter_1_right = self.comb_iter_1_right(x_right)
+        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right
+
+        x_comb_iter_2_left = self.comb_iter_2_left(x_left)
+        x_comb_iter_2_right = self.comb_iter_2_right(x_right)
+        x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right
+
+        x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)
+        x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1
+
+        x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)
+        x_comb_iter_4_right = self.comb_iter_4_right(x_left)
+        x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right
+
+        x_out = torch.cat(
+            [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1
+        )
+        return x_out
+
+
+class FirstCell(nn.Module):
+
+    def __init__(
+        self, in_channels_left, out_channels_left, in_channels_right,
+        out_channels_right
+    ):
+        super(FirstCell, self).__init__()
+        self.conv_1x1 = nn.Sequential()
+        self.conv_1x1.add_module('relu', nn.ReLU())
+        self.conv_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_right, out_channels_right, 1, stride=1, bias=False
+            )
+        )
+        self.conv_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_right, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.relu = nn.ReLU()
+        self.path_1 = nn.Sequential()
+        self.path_1.add_module(
+            'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)
+        )
+        self.path_1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_left, out_channels_left, 1, stride=1, bias=False
+            )
+        )
+        self.path_2 = nn.ModuleList()
+        self.path_2.add_module('pad', nn.ZeroPad2d((0, 1, 0, 1)))
+        self.path_2.add_module(
+            'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)
+        )
+        self.path_2.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_left, out_channels_left, 1, stride=1, bias=False
+            )
+        )
+
+        self.final_path_bn = nn.BatchNorm2d(
+            out_channels_left * 2, eps=0.001, momentum=0.1, affine=True
+        )
+
+        self.comb_iter_0_left = BranchSeparables(
+            out_channels_right, out_channels_right, 5, 1, 2, bias=False
+        )
+        self.comb_iter_0_right = BranchSeparables(
+            out_channels_right, out_channels_right, 3, 1, 1, bias=False
+        )
+
+        self.comb_iter_1_left = BranchSeparables(
+            out_channels_right, out_channels_right, 5, 1, 2, bias=False
+        )
+        self.comb_iter_1_right = BranchSeparables(
+            out_channels_right, out_channels_right, 3, 1, 1, bias=False
+        )
+
+        self.comb_iter_2_left = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_3_left = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+        self.comb_iter_3_right = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_4_left = BranchSeparables(
+            out_channels_right, out_channels_right, 3, 1, 1, bias=False
+        )
+
+    def forward(self, x, x_prev):
+        x_relu = self.relu(x_prev)
+        # path 1
+        x_path1 = self.path_1(x_relu)
+        # path 2
+        x_path2 = self.path_2.pad(x_relu)
+        x_path2 = x_path2[:, :, 1:, 1:]
+        x_path2 = self.path_2.avgpool(x_path2)
+        x_path2 = self.path_2.conv(x_path2)
+        # final path
+        x_left = self.final_path_bn(torch.cat([x_path1, x_path2], 1))
+
+        x_right = self.conv_1x1(x)
+
+        x_comb_iter_0_left = self.comb_iter_0_left(x_right)
+        x_comb_iter_0_right = self.comb_iter_0_right(x_left)
+        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right
+
+        x_comb_iter_1_left = self.comb_iter_1_left(x_left)
+        x_comb_iter_1_right = self.comb_iter_1_right(x_left)
+        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right
+
+        x_comb_iter_2_left = self.comb_iter_2_left(x_right)
+        x_comb_iter_2 = x_comb_iter_2_left + x_left
+
+        x_comb_iter_3_left = self.comb_iter_3_left(x_left)
+        x_comb_iter_3_right = self.comb_iter_3_right(x_left)
+        x_comb_iter_3 = x_comb_iter_3_left + x_comb_iter_3_right
+
+        x_comb_iter_4_left = self.comb_iter_4_left(x_right)
+        x_comb_iter_4 = x_comb_iter_4_left + x_right
+
+        x_out = torch.cat(
+            [
+                x_left, x_comb_iter_0, x_comb_iter_1, x_comb_iter_2,
+                x_comb_iter_3, x_comb_iter_4
+            ], 1
+        )
+        return x_out
+
+
+class NormalCell(nn.Module):
+
+    def __init__(
+        self, in_channels_left, out_channels_left, in_channels_right,
+        out_channels_right
+    ):
+        super(NormalCell, self).__init__()
+        self.conv_prev_1x1 = nn.Sequential()
+        self.conv_prev_1x1.add_module('relu', nn.ReLU())
+        self.conv_prev_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_left, out_channels_left, 1, stride=1, bias=False
+            )
+        )
+        self.conv_prev_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_left, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.conv_1x1 = nn.Sequential()
+        self.conv_1x1.add_module('relu', nn.ReLU())
+        self.conv_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_right, out_channels_right, 1, stride=1, bias=False
+            )
+        )
+        self.conv_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_right, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.comb_iter_0_left = BranchSeparables(
+            out_channels_right, out_channels_right, 5, 1, 2, bias=False
+        )
+        self.comb_iter_0_right = BranchSeparables(
+            out_channels_left, out_channels_left, 3, 1, 1, bias=False
+        )
+
+        self.comb_iter_1_left = BranchSeparables(
+            out_channels_left, out_channels_left, 5, 1, 2, bias=False
+        )
+        self.comb_iter_1_right = BranchSeparables(
+            out_channels_left, out_channels_left, 3, 1, 1, bias=False
+        )
+
+        self.comb_iter_2_left = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_3_left = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+        self.comb_iter_3_right = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_4_left = BranchSeparables(
+            out_channels_right, out_channels_right, 3, 1, 1, bias=False
+        )
+
+    def forward(self, x, x_prev):
+        x_left = self.conv_prev_1x1(x_prev)
+        x_right = self.conv_1x1(x)
+
+        x_comb_iter_0_left = self.comb_iter_0_left(x_right)
+        x_comb_iter_0_right = self.comb_iter_0_right(x_left)
+        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right
+
+        x_comb_iter_1_left = self.comb_iter_1_left(x_left)
+        x_comb_iter_1_right = self.comb_iter_1_right(x_left)
+        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right
+
+        x_comb_iter_2_left = self.comb_iter_2_left(x_right)
+        x_comb_iter_2 = x_comb_iter_2_left + x_left
+
+        x_comb_iter_3_left = self.comb_iter_3_left(x_left)
+        x_comb_iter_3_right = self.comb_iter_3_right(x_left)
+        x_comb_iter_3 = x_comb_iter_3_left + x_comb_iter_3_right
+
+        x_comb_iter_4_left = self.comb_iter_4_left(x_right)
+        x_comb_iter_4 = x_comb_iter_4_left + x_right
+
+        x_out = torch.cat(
+            [
+                x_left, x_comb_iter_0, x_comb_iter_1, x_comb_iter_2,
+                x_comb_iter_3, x_comb_iter_4
+            ], 1
+        )
+        return x_out
+
+
+class ReductionCell0(nn.Module):
+
+    def __init__(
+        self, in_channels_left, out_channels_left, in_channels_right,
+        out_channels_right
+    ):
+        super(ReductionCell0, self).__init__()
+        self.conv_prev_1x1 = nn.Sequential()
+        self.conv_prev_1x1.add_module('relu', nn.ReLU())
+        self.conv_prev_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_left, out_channels_left, 1, stride=1, bias=False
+            )
+        )
+        self.conv_prev_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_left, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.conv_1x1 = nn.Sequential()
+        self.conv_1x1.add_module('relu', nn.ReLU())
+        self.conv_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_right, out_channels_right, 1, stride=1, bias=False
+            )
+        )
+        self.conv_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_right, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.comb_iter_0_left = BranchSeparablesReduction(
+            out_channels_right, out_channels_right, 5, 2, 2, bias=False
+        )
+        self.comb_iter_0_right = BranchSeparablesReduction(
+            out_channels_right, out_channels_right, 7, 2, 3, bias=False
+        )
+
+        self.comb_iter_1_left = MaxPoolPad()
+        self.comb_iter_1_right = BranchSeparablesReduction(
+            out_channels_right, out_channels_right, 7, 2, 3, bias=False
+        )
+
+        self.comb_iter_2_left = AvgPoolPad()
+        self.comb_iter_2_right = BranchSeparablesReduction(
+            out_channels_right, out_channels_right, 5, 2, 2, bias=False
+        )
+
+        self.comb_iter_3_right = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_4_left = BranchSeparablesReduction(
+            out_channels_right, out_channels_right, 3, 1, 1, bias=False
+        )
+        self.comb_iter_4_right = MaxPoolPad()
+
+    def forward(self, x, x_prev):
+        x_left = self.conv_prev_1x1(x_prev)
+        x_right = self.conv_1x1(x)
+
+        x_comb_iter_0_left = self.comb_iter_0_left(x_right)
+        x_comb_iter_0_right = self.comb_iter_0_right(x_left)
+        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right
+
+        x_comb_iter_1_left = self.comb_iter_1_left(x_right)
+        x_comb_iter_1_right = self.comb_iter_1_right(x_left)
+        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right
+
+        x_comb_iter_2_left = self.comb_iter_2_left(x_right)
+        x_comb_iter_2_right = self.comb_iter_2_right(x_left)
+        x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right
+
+        x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)
+        x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1
+
+        x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)
+        x_comb_iter_4_right = self.comb_iter_4_right(x_right)
+        x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right
+
+        x_out = torch.cat(
+            [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1
+        )
+        return x_out
+
+
+class ReductionCell1(nn.Module):
+
+    def __init__(
+        self, in_channels_left, out_channels_left, in_channels_right,
+        out_channels_right
+    ):
+        super(ReductionCell1, self).__init__()
+        self.conv_prev_1x1 = nn.Sequential()
+        self.conv_prev_1x1.add_module('relu', nn.ReLU())
+        self.conv_prev_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_left, out_channels_left, 1, stride=1, bias=False
+            )
+        )
+        self.conv_prev_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_left, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.conv_1x1 = nn.Sequential()
+        self.conv_1x1.add_module('relu', nn.ReLU())
+        self.conv_1x1.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels_right, out_channels_right, 1, stride=1, bias=False
+            )
+        )
+        self.conv_1x1.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                out_channels_right, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.comb_iter_0_left = BranchSeparables(
+            out_channels_right,
+            out_channels_right,
+            5,
+            2,
+            2,
+            name='specific',
+            bias=False
+        )
+        self.comb_iter_0_right = BranchSeparables(
+            out_channels_right,
+            out_channels_right,
+            7,
+            2,
+            3,
+            name='specific',
+            bias=False
+        )
+
+        # self.comb_iter_1_left = nn.MaxPool2d(3, stride=2, padding=1)
+        self.comb_iter_1_left = MaxPoolPad()
+        self.comb_iter_1_right = BranchSeparables(
+            out_channels_right,
+            out_channels_right,
+            7,
+            2,
+            3,
+            name='specific',
+            bias=False
+        )
+
+        # self.comb_iter_2_left = nn.AvgPool2d(3, stride=2, padding=1, count_include_pad=False)
+        self.comb_iter_2_left = AvgPoolPad()
+        self.comb_iter_2_right = BranchSeparables(
+            out_channels_right,
+            out_channels_right,
+            5,
+            2,
+            2,
+            name='specific',
+            bias=False
+        )
+
+        self.comb_iter_3_right = nn.AvgPool2d(
+            3, stride=1, padding=1, count_include_pad=False
+        )
+
+        self.comb_iter_4_left = BranchSeparables(
+            out_channels_right,
+            out_channels_right,
+            3,
+            1,
+            1,
+            name='specific',
+            bias=False
+        )
+        # self.comb_iter_4_right = nn.MaxPool2d(3, stride=2, padding=1)
+        self.comb_iter_4_right = MaxPoolPad()
+
+    def forward(self, x, x_prev):
+        x_left = self.conv_prev_1x1(x_prev)
+        x_right = self.conv_1x1(x)
+
+        x_comb_iter_0_left = self.comb_iter_0_left(x_right)
+        x_comb_iter_0_right = self.comb_iter_0_right(x_left)
+        x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right
+
+        x_comb_iter_1_left = self.comb_iter_1_left(x_right)
+        x_comb_iter_1_right = self.comb_iter_1_right(x_left)
+        x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right
+
+        x_comb_iter_2_left = self.comb_iter_2_left(x_right)
+        x_comb_iter_2_right = self.comb_iter_2_right(x_left)
+        x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right
+
+        x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)
+        x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1
+
+        x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)
+        x_comb_iter_4_right = self.comb_iter_4_right(x_right)
+        x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right
+
+        x_out = torch.cat(
+            [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1
+        )
+        return x_out
+
+
+class NASNetAMobile(nn.Module):
+    """Neural Architecture Search (NAS).
+
+    Reference:
+        Zoph et al. Learning Transferable Architectures
+        for Scalable Image Recognition. CVPR 2018.
+
+    Public keys:
+        - ``nasnetamobile``: NASNet-A Mobile.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        stem_filters=32,
+        penultimate_filters=1056,
+        filters_multiplier=2,
+        **kwargs
+    ):
+        super(NASNetAMobile, self).__init__()
+        self.stem_filters = stem_filters
+        self.penultimate_filters = penultimate_filters
+        self.filters_multiplier = filters_multiplier
+        self.loss = loss
+
+        filters = self.penultimate_filters // 24
+        # 24 is default value for the architecture
+
+        self.conv0 = nn.Sequential()
+        self.conv0.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels=3,
+                out_channels=self.stem_filters,
+                kernel_size=3,
+                padding=0,
+                stride=2,
+                bias=False
+            )
+        )
+        self.conv0.add_module(
+            'bn',
+            nn.BatchNorm2d(
+                self.stem_filters, eps=0.001, momentum=0.1, affine=True
+            )
+        )
+
+        self.cell_stem_0 = CellStem0(
+            self.stem_filters, num_filters=filters // (filters_multiplier**2)
+        )
+        self.cell_stem_1 = CellStem1(
+            self.stem_filters, num_filters=filters // filters_multiplier
+        )
+
+        self.cell_0 = FirstCell(
+            in_channels_left=filters,
+            out_channels_left=filters // 2, # 1, 0.5
+            in_channels_right=2 * filters,
+            out_channels_right=filters
+        ) # 2, 1
+        self.cell_1 = NormalCell(
+            in_channels_left=2 * filters,
+            out_channels_left=filters, # 2, 1
+            in_channels_right=6 * filters,
+            out_channels_right=filters
+        ) # 6, 1
+        self.cell_2 = NormalCell(
+            in_channels_left=6 * filters,
+            out_channels_left=filters, # 6, 1
+            in_channels_right=6 * filters,
+            out_channels_right=filters
+        ) # 6, 1
+        self.cell_3 = NormalCell(
+            in_channels_left=6 * filters,
+            out_channels_left=filters, # 6, 1
+            in_channels_right=6 * filters,
+            out_channels_right=filters
+        ) # 6, 1
+
+        self.reduction_cell_0 = ReductionCell0(
+            in_channels_left=6 * filters,
+            out_channels_left=2 * filters, # 6, 2
+            in_channels_right=6 * filters,
+            out_channels_right=2 * filters
+        ) # 6, 2
+
+        self.cell_6 = FirstCell(
+            in_channels_left=6 * filters,
+            out_channels_left=filters, # 6, 1
+            in_channels_right=8 * filters,
+            out_channels_right=2 * filters
+        ) # 8, 2
+        self.cell_7 = NormalCell(
+            in_channels_left=8 * filters,
+            out_channels_left=2 * filters, # 8, 2
+            in_channels_right=12 * filters,
+            out_channels_right=2 * filters
+        ) # 12, 2
+        self.cell_8 = NormalCell(
+            in_channels_left=12 * filters,
+            out_channels_left=2 * filters, # 12, 2
+            in_channels_right=12 * filters,
+            out_channels_right=2 * filters
+        ) # 12, 2
+        self.cell_9 = NormalCell(
+            in_channels_left=12 * filters,
+            out_channels_left=2 * filters, # 12, 2
+            in_channels_right=12 * filters,
+            out_channels_right=2 * filters
+        ) # 12, 2
+
+        self.reduction_cell_1 = ReductionCell1(
+            in_channels_left=12 * filters,
+            out_channels_left=4 * filters, # 12, 4
+            in_channels_right=12 * filters,
+            out_channels_right=4 * filters
+        ) # 12, 4
+
+        self.cell_12 = FirstCell(
+            in_channels_left=12 * filters,
+            out_channels_left=2 * filters, # 12, 2
+            in_channels_right=16 * filters,
+            out_channels_right=4 * filters
+        ) # 16, 4
+        self.cell_13 = NormalCell(
+            in_channels_left=16 * filters,
+            out_channels_left=4 * filters, # 16, 4
+            in_channels_right=24 * filters,
+            out_channels_right=4 * filters
+        ) # 24, 4
+        self.cell_14 = NormalCell(
+            in_channels_left=24 * filters,
+            out_channels_left=4 * filters, # 24, 4
+            in_channels_right=24 * filters,
+            out_channels_right=4 * filters
+        ) # 24, 4
+        self.cell_15 = NormalCell(
+            in_channels_left=24 * filters,
+            out_channels_left=4 * filters, # 24, 4
+            in_channels_right=24 * filters,
+            out_channels_right=4 * filters
+        ) # 24, 4
+
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout()
+        self.classifier = nn.Linear(24 * filters, num_classes)
+
+        self._init_params()
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def features(self, input):
+        x_conv0 = self.conv0(input)
+        x_stem_0 = self.cell_stem_0(x_conv0)
+        x_stem_1 = self.cell_stem_1(x_conv0, x_stem_0)
+
+        x_cell_0 = self.cell_0(x_stem_1, x_stem_0)
+        x_cell_1 = self.cell_1(x_cell_0, x_stem_1)
+        x_cell_2 = self.cell_2(x_cell_1, x_cell_0)
+        x_cell_3 = self.cell_3(x_cell_2, x_cell_1)
+
+        x_reduction_cell_0 = self.reduction_cell_0(x_cell_3, x_cell_2)
+
+        x_cell_6 = self.cell_6(x_reduction_cell_0, x_cell_3)
+        x_cell_7 = self.cell_7(x_cell_6, x_reduction_cell_0)
+        x_cell_8 = self.cell_8(x_cell_7, x_cell_6)
+        x_cell_9 = self.cell_9(x_cell_8, x_cell_7)
+
+        x_reduction_cell_1 = self.reduction_cell_1(x_cell_9, x_cell_8)
+
+        x_cell_12 = self.cell_12(x_reduction_cell_1, x_cell_9)
+        x_cell_13 = self.cell_13(x_cell_12, x_reduction_cell_1)
+        x_cell_14 = self.cell_14(x_cell_13, x_cell_12)
+        x_cell_15 = self.cell_15(x_cell_14, x_cell_13)
+
+        x_cell_15 = self.relu(x_cell_15)
+        x_cell_15 = F.avg_pool2d(
+            x_cell_15,
+            x_cell_15.size()[2:]
+        ) # global average pool
+        x_cell_15 = x_cell_15.view(x_cell_15.size(0), -1)
+        x_cell_15 = self.dropout(x_cell_15)
+
+        return x_cell_15
+
+    def forward(self, input):
+        v = self.features(input)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def nasnetamobile(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = NASNetAMobile(num_classes, loss, **kwargs)
+    if pretrained:
+        model_url = pretrained_settings['nasnetamobile']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model

trackers/strongsort/deep/models/osnet.py

@@ -0,0 +1,598 @@
+from __future__ import division, absolute_import
+import warnings
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = [
+    'osnet_x1_0', 'osnet_x0_75', 'osnet_x0_5', 'osnet_x0_25', 'osnet_ibn_x1_0'
+]
+
+pretrained_urls = {
+    'osnet_x1_0':
+    'https://drive.google.com/uc?id=1LaG1EJpHrxdAxKnSCJ_i0u-nbxSAeiFY',
+    'osnet_x0_75':
+    'https://drive.google.com/uc?id=1uwA9fElHOk3ZogwbeY5GkLI6QPTX70Hq',
+    'osnet_x0_5':
+    'https://drive.google.com/uc?id=16DGLbZukvVYgINws8u8deSaOqjybZ83i',
+    'osnet_x0_25':
+    'https://drive.google.com/uc?id=1rb8UN5ZzPKRc_xvtHlyDh-cSz88YX9hs',
+    'osnet_ibn_x1_0':
+    'https://drive.google.com/uc?id=1sr90V6irlYYDd4_4ISU2iruoRG8J__6l'
+}
+
+
+##########
+# Basic layers
+##########
+class ConvLayer(nn.Module):
+    """Convolution layer (conv + bn + relu)."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        groups=1,
+        IN=False
+    ):
+        super(ConvLayer, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False,
+            groups=groups
+        )
+        if IN:
+            self.bn = nn.InstanceNorm2d(out_channels, affine=True)
+        else:
+            self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class Conv1x1(nn.Module):
+    """1x1 convolution + bn + relu."""
+
+    def __init__(self, in_channels, out_channels, stride=1, groups=1):
+        super(Conv1x1, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            1,
+            stride=stride,
+            padding=0,
+            bias=False,
+            groups=groups
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class Conv1x1Linear(nn.Module):
+    """1x1 convolution + bn (w/o non-linearity)."""
+
+    def __init__(self, in_channels, out_channels, stride=1):
+        super(Conv1x1Linear, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels, out_channels, 1, stride=stride, padding=0, bias=False
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class Conv3x3(nn.Module):
+    """3x3 convolution + bn + relu."""
+
+    def __init__(self, in_channels, out_channels, stride=1, groups=1):
+        super(Conv3x3, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            3,
+            stride=stride,
+            padding=1,
+            bias=False,
+            groups=groups
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class LightConv3x3(nn.Module):
+    """Lightweight 3x3 convolution.
+
+    1x1 (linear) + dw 3x3 (nonlinear).
+    """
+
+    def __init__(self, in_channels, out_channels):
+        super(LightConv3x3, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, 1, stride=1, padding=0, bias=False
+        )
+        self.conv2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            3,
+            stride=1,
+            padding=1,
+            bias=False,
+            groups=out_channels
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+##########
+# Building blocks for omni-scale feature learning
+##########
+class ChannelGate(nn.Module):
+    """A mini-network that generates channel-wise gates conditioned on input tensor."""
+
+    def __init__(
+        self,
+        in_channels,
+        num_gates=None,
+        return_gates=False,
+        gate_activation='sigmoid',
+        reduction=16,
+        layer_norm=False
+    ):
+        super(ChannelGate, self).__init__()
+        if num_gates is None:
+            num_gates = in_channels
+        self.return_gates = return_gates
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc1 = nn.Conv2d(
+            in_channels,
+            in_channels // reduction,
+            kernel_size=1,
+            bias=True,
+            padding=0
+        )
+        self.norm1 = None
+        if layer_norm:
+            self.norm1 = nn.LayerNorm((in_channels // reduction, 1, 1))
+        self.relu = nn.ReLU(inplace=True)
+        self.fc2 = nn.Conv2d(
+            in_channels // reduction,
+            num_gates,
+            kernel_size=1,
+            bias=True,
+            padding=0
+        )
+        if gate_activation == 'sigmoid':
+            self.gate_activation = nn.Sigmoid()
+        elif gate_activation == 'relu':
+            self.gate_activation = nn.ReLU(inplace=True)
+        elif gate_activation == 'linear':
+            self.gate_activation = None
+        else:
+            raise RuntimeError(
+                "Unknown gate activation: {}".format(gate_activation)
+            )
+
+    def forward(self, x):
+        input = x
+        x = self.global_avgpool(x)
+        x = self.fc1(x)
+        if self.norm1 is not None:
+            x = self.norm1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        if self.gate_activation is not None:
+            x = self.gate_activation(x)
+        if self.return_gates:
+            return x
+        return input * x
+
+
+class OSBlock(nn.Module):
+    """Omni-scale feature learning block."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        IN=False,
+        bottleneck_reduction=4,
+        **kwargs
+    ):
+        super(OSBlock, self).__init__()
+        mid_channels = out_channels // bottleneck_reduction
+        self.conv1 = Conv1x1(in_channels, mid_channels)
+        self.conv2a = LightConv3x3(mid_channels, mid_channels)
+        self.conv2b = nn.Sequential(
+            LightConv3x3(mid_channels, mid_channels),
+            LightConv3x3(mid_channels, mid_channels),
+        )
+        self.conv2c = nn.Sequential(
+            LightConv3x3(mid_channels, mid_channels),
+            LightConv3x3(mid_channels, mid_channels),
+            LightConv3x3(mid_channels, mid_channels),
+        )
+        self.conv2d = nn.Sequential(
+            LightConv3x3(mid_channels, mid_channels),
+            LightConv3x3(mid_channels, mid_channels),
+            LightConv3x3(mid_channels, mid_channels),
+            LightConv3x3(mid_channels, mid_channels),
+        )
+        self.gate = ChannelGate(mid_channels)
+        self.conv3 = Conv1x1Linear(mid_channels, out_channels)
+        self.downsample = None
+        if in_channels != out_channels:
+            self.downsample = Conv1x1Linear(in_channels, out_channels)
+        self.IN = None
+        if IN:
+            self.IN = nn.InstanceNorm2d(out_channels, affine=True)
+
+    def forward(self, x):
+        identity = x
+        x1 = self.conv1(x)
+        x2a = self.conv2a(x1)
+        x2b = self.conv2b(x1)
+        x2c = self.conv2c(x1)
+        x2d = self.conv2d(x1)
+        x2 = self.gate(x2a) + self.gate(x2b) + self.gate(x2c) + self.gate(x2d)
+        x3 = self.conv3(x2)
+        if self.downsample is not None:
+            identity = self.downsample(identity)
+        out = x3 + identity
+        if self.IN is not None:
+            out = self.IN(out)
+        return F.relu(out)
+
+
+##########
+# Network architecture
+##########
+class OSNet(nn.Module):
+    """Omni-Scale Network.
+    
+    Reference:
+        - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.
+        - Zhou et al. Learning Generalisable Omni-Scale Representations
+          for Person Re-Identification. TPAMI, 2021.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        blocks,
+        layers,
+        channels,
+        feature_dim=512,
+        loss='softmax',
+        IN=False,
+        **kwargs
+    ):
+        super(OSNet, self).__init__()
+        num_blocks = len(blocks)
+        assert num_blocks == len(layers)
+        assert num_blocks == len(channels) - 1
+        self.loss = loss
+        self.feature_dim = feature_dim
+
+        # convolutional backbone
+        self.conv1 = ConvLayer(3, channels[0], 7, stride=2, padding=3, IN=IN)
+        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
+        self.conv2 = self._make_layer(
+            blocks[0],
+            layers[0],
+            channels[0],
+            channels[1],
+            reduce_spatial_size=True,
+            IN=IN
+        )
+        self.conv3 = self._make_layer(
+            blocks[1],
+            layers[1],
+            channels[1],
+            channels[2],
+            reduce_spatial_size=True
+        )
+        self.conv4 = self._make_layer(
+            blocks[2],
+            layers[2],
+            channels[2],
+            channels[3],
+            reduce_spatial_size=False
+        )
+        self.conv5 = Conv1x1(channels[3], channels[3])
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        # fully connected layer
+        self.fc = self._construct_fc_layer(
+            self.feature_dim, channels[3], dropout_p=None
+        )
+        # identity classification layer
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _make_layer(
+        self,
+        block,
+        layer,
+        in_channels,
+        out_channels,
+        reduce_spatial_size,
+        IN=False
+    ):
+        layers = []
+
+        layers.append(block(in_channels, out_channels, IN=IN))
+        for i in range(1, layer):
+            layers.append(block(out_channels, out_channels, IN=IN))
+
+        if reduce_spatial_size:
+            layers.append(
+                nn.Sequential(
+                    Conv1x1(out_channels, out_channels),
+                    nn.AvgPool2d(2, stride=2)
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        if fc_dims is None or fc_dims < 0:
+            self.feature_dim = input_dim
+            return None
+
+        if isinstance(fc_dims, int):
+            fc_dims = [fc_dims]
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = self.conv5(x)
+        return x
+
+    def forward(self, x, return_featuremaps=False):
+        x = self.featuremaps(x)
+        if return_featuremaps:
+            return x
+        v = self.global_avgpool(x)
+        v = v.view(v.size(0), -1)
+        if self.fc is not None:
+            v = self.fc(v)
+        if not self.training:
+            return v
+        y = self.classifier(v)
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, key=''):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    import os
+    import errno
+    import gdown
+    from collections import OrderedDict
+
+    def _get_torch_home():
+        ENV_TORCH_HOME = 'TORCH_HOME'
+        ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
+        DEFAULT_CACHE_DIR = '~/.cache'
+        torch_home = os.path.expanduser(
+            os.getenv(
+                ENV_TORCH_HOME,
+                os.path.join(
+                    os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
+                )
+            )
+        )
+        return torch_home
+
+    torch_home = _get_torch_home()
+    model_dir = os.path.join(torch_home, 'checkpoints')
+    try:
+        os.makedirs(model_dir)
+    except OSError as e:
+        if e.errno == errno.EEXIST:
+            # Directory already exists, ignore.
+            pass
+        else:
+            # Unexpected OSError, re-raise.
+            raise
+    filename = key + '_imagenet.pth'
+    cached_file = os.path.join(model_dir, filename)
+
+    if not os.path.exists(cached_file):
+        gdown.download(pretrained_urls[key], cached_file, quiet=False)
+
+    state_dict = torch.load(cached_file)
+    model_dict = model.state_dict()
+    new_state_dict = OrderedDict()
+    matched_layers, discarded_layers = [], []
+
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            k = k[7:] # discard module.
+
+        if k in model_dict and model_dict[k].size() == v.size():
+            new_state_dict[k] = v
+            matched_layers.append(k)
+        else:
+            discarded_layers.append(k)
+
+    model_dict.update(new_state_dict)
+    model.load_state_dict(model_dict)
+
+    if len(matched_layers) == 0:
+        warnings.warn(
+            'The pretrained weights from "{}" cannot be loaded, '
+            'please check the key names manually '
+            '(** ignored and continue **)'.format(cached_file)
+        )
+    else:
+        print(
+            'Successfully loaded imagenet pretrained weights from "{}"'.
+            format(cached_file)
+        )
+        if len(discarded_layers) > 0:
+            print(
+                '** The following layers are discarded '
+                'due to unmatched keys or layer size: {}'.
+                format(discarded_layers)
+            )
+
+
+##########
+# Instantiation
+##########
+def osnet_x1_0(num_classes=1000, pretrained=True, loss='softmax', **kwargs):
+    # standard size (width x1.0)
+    model = OSNet(
+        num_classes,
+        blocks=[OSBlock, OSBlock, OSBlock],
+        layers=[2, 2, 2],
+        channels=[64, 256, 384, 512],
+        loss=loss,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_x1_0')
+    return model
+
+
+def osnet_x0_75(num_classes=1000, pretrained=True, loss='softmax', **kwargs):
+    # medium size (width x0.75)
+    model = OSNet(
+        num_classes,
+        blocks=[OSBlock, OSBlock, OSBlock],
+        layers=[2, 2, 2],
+        channels=[48, 192, 288, 384],
+        loss=loss,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_x0_75')
+    return model
+
+
+def osnet_x0_5(num_classes=1000, pretrained=True, loss='softmax', **kwargs):
+    # tiny size (width x0.5)
+    model = OSNet(
+        num_classes,
+        blocks=[OSBlock, OSBlock, OSBlock],
+        layers=[2, 2, 2],
+        channels=[32, 128, 192, 256],
+        loss=loss,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_x0_5')
+    return model
+
+
+def osnet_x0_25(num_classes=1000, pretrained=True, loss='softmax', **kwargs):
+    # very tiny size (width x0.25)
+    model = OSNet(
+        num_classes,
+        blocks=[OSBlock, OSBlock, OSBlock],
+        layers=[2, 2, 2],
+        channels=[16, 64, 96, 128],
+        loss=loss,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_x0_25')
+    return model
+
+
+def osnet_ibn_x1_0(
+    num_classes=1000, pretrained=True, loss='softmax', **kwargs
+):
+    # standard size (width x1.0) + IBN layer
+    # Ref: Pan et al. Two at Once: Enhancing Learning and Generalization Capacities via IBN-Net. ECCV, 2018.
+    model = OSNet(
+        num_classes,
+        blocks=[OSBlock, OSBlock, OSBlock],
+        layers=[2, 2, 2],
+        channels=[64, 256, 384, 512],
+        loss=loss,
+        IN=True,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_ibn_x1_0')
+    return model

trackers/strongsort/deep/models/osnet_ain.py

@@ -0,0 +1,609 @@
+from __future__ import division, absolute_import
+import warnings
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = [
+    'osnet_ain_x1_0', 'osnet_ain_x0_75', 'osnet_ain_x0_5', 'osnet_ain_x0_25'
+]
+
+pretrained_urls = {
+    'osnet_ain_x1_0':
+    'https://drive.google.com/uc?id=1-CaioD9NaqbHK_kzSMW8VE4_3KcsRjEo',
+    'osnet_ain_x0_75':
+    'https://drive.google.com/uc?id=1apy0hpsMypqstfencdH-jKIUEFOW4xoM',
+    'osnet_ain_x0_5':
+    'https://drive.google.com/uc?id=1KusKvEYyKGDTUBVRxRiz55G31wkihB6l',
+    'osnet_ain_x0_25':
+    'https://drive.google.com/uc?id=1SxQt2AvmEcgWNhaRb2xC4rP6ZwVDP0Wt'
+}
+
+
+##########
+# Basic layers
+##########
+class ConvLayer(nn.Module):
+    """Convolution layer (conv + bn + relu)."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        groups=1,
+        IN=False
+    ):
+        super(ConvLayer, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False,
+            groups=groups
+        )
+        if IN:
+            self.bn = nn.InstanceNorm2d(out_channels, affine=True)
+        else:
+            self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return self.relu(x)
+
+
+class Conv1x1(nn.Module):
+    """1x1 convolution + bn + relu."""
+
+    def __init__(self, in_channels, out_channels, stride=1, groups=1):
+        super(Conv1x1, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            1,
+            stride=stride,
+            padding=0,
+            bias=False,
+            groups=groups
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return self.relu(x)
+
+
+class Conv1x1Linear(nn.Module):
+    """1x1 convolution + bn (w/o non-linearity)."""
+
+    def __init__(self, in_channels, out_channels, stride=1, bn=True):
+        super(Conv1x1Linear, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels, out_channels, 1, stride=stride, padding=0, bias=False
+        )
+        self.bn = None
+        if bn:
+            self.bn = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        return x
+
+
+class Conv3x3(nn.Module):
+    """3x3 convolution + bn + relu."""
+
+    def __init__(self, in_channels, out_channels, stride=1, groups=1):
+        super(Conv3x3, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            3,
+            stride=stride,
+            padding=1,
+            bias=False,
+            groups=groups
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return self.relu(x)
+
+
+class LightConv3x3(nn.Module):
+    """Lightweight 3x3 convolution.
+
+    1x1 (linear) + dw 3x3 (nonlinear).
+    """
+
+    def __init__(self, in_channels, out_channels):
+        super(LightConv3x3, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, 1, stride=1, padding=0, bias=False
+        )
+        self.conv2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            3,
+            stride=1,
+            padding=1,
+            bias=False,
+            groups=out_channels
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.bn(x)
+        return self.relu(x)
+
+
+class LightConvStream(nn.Module):
+    """Lightweight convolution stream."""
+
+    def __init__(self, in_channels, out_channels, depth):
+        super(LightConvStream, self).__init__()
+        assert depth >= 1, 'depth must be equal to or larger than 1, but got {}'.format(
+            depth
+        )
+        layers = []
+        layers += [LightConv3x3(in_channels, out_channels)]
+        for i in range(depth - 1):
+            layers += [LightConv3x3(out_channels, out_channels)]
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+##########
+# Building blocks for omni-scale feature learning
+##########
+class ChannelGate(nn.Module):
+    """A mini-network that generates channel-wise gates conditioned on input tensor."""
+
+    def __init__(
+        self,
+        in_channels,
+        num_gates=None,
+        return_gates=False,
+        gate_activation='sigmoid',
+        reduction=16,
+        layer_norm=False
+    ):
+        super(ChannelGate, self).__init__()
+        if num_gates is None:
+            num_gates = in_channels
+        self.return_gates = return_gates
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc1 = nn.Conv2d(
+            in_channels,
+            in_channels // reduction,
+            kernel_size=1,
+            bias=True,
+            padding=0
+        )
+        self.norm1 = None
+        if layer_norm:
+            self.norm1 = nn.LayerNorm((in_channels // reduction, 1, 1))
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Conv2d(
+            in_channels // reduction,
+            num_gates,
+            kernel_size=1,
+            bias=True,
+            padding=0
+        )
+        if gate_activation == 'sigmoid':
+            self.gate_activation = nn.Sigmoid()
+        elif gate_activation == 'relu':
+            self.gate_activation = nn.ReLU()
+        elif gate_activation == 'linear':
+            self.gate_activation = None
+        else:
+            raise RuntimeError(
+                "Unknown gate activation: {}".format(gate_activation)
+            )
+
+    def forward(self, x):
+        input = x
+        x = self.global_avgpool(x)
+        x = self.fc1(x)
+        if self.norm1 is not None:
+            x = self.norm1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        if self.gate_activation is not None:
+            x = self.gate_activation(x)
+        if self.return_gates:
+            return x
+        return input * x
+
+
+class OSBlock(nn.Module):
+    """Omni-scale feature learning block."""
+
+    def __init__(self, in_channels, out_channels, reduction=4, T=4, **kwargs):
+        super(OSBlock, self).__init__()
+        assert T >= 1
+        assert out_channels >= reduction and out_channels % reduction == 0
+        mid_channels = out_channels // reduction
+
+        self.conv1 = Conv1x1(in_channels, mid_channels)
+        self.conv2 = nn.ModuleList()
+        for t in range(1, T + 1):
+            self.conv2 += [LightConvStream(mid_channels, mid_channels, t)]
+        self.gate = ChannelGate(mid_channels)
+        self.conv3 = Conv1x1Linear(mid_channels, out_channels)
+        self.downsample = None
+        if in_channels != out_channels:
+            self.downsample = Conv1x1Linear(in_channels, out_channels)
+
+    def forward(self, x):
+        identity = x
+        x1 = self.conv1(x)
+        x2 = 0
+        for conv2_t in self.conv2:
+            x2_t = conv2_t(x1)
+            x2 = x2 + self.gate(x2_t)
+        x3 = self.conv3(x2)
+        if self.downsample is not None:
+            identity = self.downsample(identity)
+        out = x3 + identity
+        return F.relu(out)
+
+
+class OSBlockINin(nn.Module):
+    """Omni-scale feature learning block with instance normalization."""
+
+    def __init__(self, in_channels, out_channels, reduction=4, T=4, **kwargs):
+        super(OSBlockINin, self).__init__()
+        assert T >= 1
+        assert out_channels >= reduction and out_channels % reduction == 0
+        mid_channels = out_channels // reduction
+
+        self.conv1 = Conv1x1(in_channels, mid_channels)
+        self.conv2 = nn.ModuleList()
+        for t in range(1, T + 1):
+            self.conv2 += [LightConvStream(mid_channels, mid_channels, t)]
+        self.gate = ChannelGate(mid_channels)
+        self.conv3 = Conv1x1Linear(mid_channels, out_channels, bn=False)
+        self.downsample = None
+        if in_channels != out_channels:
+            self.downsample = Conv1x1Linear(in_channels, out_channels)
+        self.IN = nn.InstanceNorm2d(out_channels, affine=True)
+
+    def forward(self, x):
+        identity = x
+        x1 = self.conv1(x)
+        x2 = 0
+        for conv2_t in self.conv2:
+            x2_t = conv2_t(x1)
+            x2 = x2 + self.gate(x2_t)
+        x3 = self.conv3(x2)
+        x3 = self.IN(x3) # IN inside residual
+        if self.downsample is not None:
+            identity = self.downsample(identity)
+        out = x3 + identity
+        return F.relu(out)
+
+
+##########
+# Network architecture
+##########
+class OSNet(nn.Module):
+    """Omni-Scale Network.
+    
+    Reference:
+        - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.
+        - Zhou et al. Learning Generalisable Omni-Scale Representations
+          for Person Re-Identification. TPAMI, 2021.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        blocks,
+        layers,
+        channels,
+        feature_dim=512,
+        loss='softmax',
+        conv1_IN=False,
+        **kwargs
+    ):
+        super(OSNet, self).__init__()
+        num_blocks = len(blocks)
+        assert num_blocks == len(layers)
+        assert num_blocks == len(channels) - 1
+        self.loss = loss
+        self.feature_dim = feature_dim
+
+        # convolutional backbone
+        self.conv1 = ConvLayer(
+            3, channels[0], 7, stride=2, padding=3, IN=conv1_IN
+        )
+        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
+        self.conv2 = self._make_layer(
+            blocks[0], layers[0], channels[0], channels[1]
+        )
+        self.pool2 = nn.Sequential(
+            Conv1x1(channels[1], channels[1]), nn.AvgPool2d(2, stride=2)
+        )
+        self.conv3 = self._make_layer(
+            blocks[1], layers[1], channels[1], channels[2]
+        )
+        self.pool3 = nn.Sequential(
+            Conv1x1(channels[2], channels[2]), nn.AvgPool2d(2, stride=2)
+        )
+        self.conv4 = self._make_layer(
+            blocks[2], layers[2], channels[2], channels[3]
+        )
+        self.conv5 = Conv1x1(channels[3], channels[3])
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        # fully connected layer
+        self.fc = self._construct_fc_layer(
+            self.feature_dim, channels[3], dropout_p=None
+        )
+        # identity classification layer
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _make_layer(self, blocks, layer, in_channels, out_channels):
+        layers = []
+        layers += [blocks[0](in_channels, out_channels)]
+        for i in range(1, len(blocks)):
+            layers += [blocks[i](out_channels, out_channels)]
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        if fc_dims is None or fc_dims < 0:
+            self.feature_dim = input_dim
+            return None
+
+        if isinstance(fc_dims, int):
+            fc_dims = [fc_dims]
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU())
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.InstanceNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        x = self.conv2(x)
+        x = self.pool2(x)
+        x = self.conv3(x)
+        x = self.pool3(x)
+        x = self.conv4(x)
+        x = self.conv5(x)
+        return x
+
+    def forward(self, x, return_featuremaps=False):
+        x = self.featuremaps(x)
+        if return_featuremaps:
+            return x
+        v = self.global_avgpool(x)
+        v = v.view(v.size(0), -1)
+        if self.fc is not None:
+            v = self.fc(v)
+        if not self.training:
+            return v
+        y = self.classifier(v)
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, key=''):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    import os
+    import errno
+    import gdown
+    from collections import OrderedDict
+
+    def _get_torch_home():
+        ENV_TORCH_HOME = 'TORCH_HOME'
+        ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'
+        DEFAULT_CACHE_DIR = '~/.cache'
+        torch_home = os.path.expanduser(
+            os.getenv(
+                ENV_TORCH_HOME,
+                os.path.join(
+                    os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'
+                )
+            )
+        )
+        return torch_home
+
+    torch_home = _get_torch_home()
+    model_dir = os.path.join(torch_home, 'checkpoints')
+    try:
+        os.makedirs(model_dir)
+    except OSError as e:
+        if e.errno == errno.EEXIST:
+            # Directory already exists, ignore.
+            pass
+        else:
+            # Unexpected OSError, re-raise.
+            raise
+    filename = key + '_imagenet.pth'
+    cached_file = os.path.join(model_dir, filename)
+
+    if not os.path.exists(cached_file):
+        gdown.download(pretrained_urls[key], cached_file, quiet=False)
+
+    state_dict = torch.load(cached_file)
+    model_dict = model.state_dict()
+    new_state_dict = OrderedDict()
+    matched_layers, discarded_layers = [], []
+
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            k = k[7:] # discard module.
+
+        if k in model_dict and model_dict[k].size() == v.size():
+            new_state_dict[k] = v
+            matched_layers.append(k)
+        else:
+            discarded_layers.append(k)
+
+    model_dict.update(new_state_dict)
+    model.load_state_dict(model_dict)
+
+    if len(matched_layers) == 0:
+        warnings.warn(
+            'The pretrained weights from "{}" cannot be loaded, '
+            'please check the key names manually '
+            '(** ignored and continue **)'.format(cached_file)
+        )
+    else:
+        print(
+            'Successfully loaded imagenet pretrained weights from "{}"'.
+            format(cached_file)
+        )
+        if len(discarded_layers) > 0:
+            print(
+                '** The following layers are discarded '
+                'due to unmatched keys or layer size: {}'.
+                format(discarded_layers)
+            )
+
+
+##########
+# Instantiation
+##########
+def osnet_ain_x1_0(
+    num_classes=1000, pretrained=True, loss='softmax', **kwargs
+):
+    model = OSNet(
+        num_classes,
+        blocks=[
+            [OSBlockINin, OSBlockINin], [OSBlock, OSBlockINin],
+            [OSBlockINin, OSBlock]
+        ],
+        layers=[2, 2, 2],
+        channels=[64, 256, 384, 512],
+        loss=loss,
+        conv1_IN=True,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_ain_x1_0')
+    return model
+
+
+def osnet_ain_x0_75(
+    num_classes=1000, pretrained=True, loss='softmax', **kwargs
+):
+    model = OSNet(
+        num_classes,
+        blocks=[
+            [OSBlockINin, OSBlockINin], [OSBlock, OSBlockINin],
+            [OSBlockINin, OSBlock]
+        ],
+        layers=[2, 2, 2],
+        channels=[48, 192, 288, 384],
+        loss=loss,
+        conv1_IN=True,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_ain_x0_75')
+    return model
+
+
+def osnet_ain_x0_5(
+    num_classes=1000, pretrained=True, loss='softmax', **kwargs
+):
+    model = OSNet(
+        num_classes,
+        blocks=[
+            [OSBlockINin, OSBlockINin], [OSBlock, OSBlockINin],
+            [OSBlockINin, OSBlock]
+        ],
+        layers=[2, 2, 2],
+        channels=[32, 128, 192, 256],
+        loss=loss,
+        conv1_IN=True,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_ain_x0_5')
+    return model
+
+
+def osnet_ain_x0_25(
+    num_classes=1000, pretrained=True, loss='softmax', **kwargs
+):
+    model = OSNet(
+        num_classes,
+        blocks=[
+            [OSBlockINin, OSBlockINin], [OSBlock, OSBlockINin],
+            [OSBlockINin, OSBlock]
+        ],
+        layers=[2, 2, 2],
+        channels=[16, 64, 96, 128],
+        loss=loss,
+        conv1_IN=True,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, key='osnet_ain_x0_25')
+    return model

trackers/strongsort/deep/models/pcb.py

@@ -0,0 +1,314 @@
+from __future__ import division, absolute_import
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ['pcb_p6', 'pcb_p4']
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, planes * self.expansion, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class DimReduceLayer(nn.Module):
+
+    def __init__(self, in_channels, out_channels, nonlinear):
+        super(DimReduceLayer, self).__init__()
+        layers = []
+        layers.append(
+            nn.Conv2d(
+                in_channels, out_channels, 1, stride=1, padding=0, bias=False
+            )
+        )
+        layers.append(nn.BatchNorm2d(out_channels))
+
+        if nonlinear == 'relu':
+            layers.append(nn.ReLU(inplace=True))
+        elif nonlinear == 'leakyrelu':
+            layers.append(nn.LeakyReLU(0.1))
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class PCB(nn.Module):
+    """Part-based Convolutional Baseline.
+
+    Reference:
+        Sun et al. Beyond Part Models: Person Retrieval with Refined
+        Part Pooling (and A Strong Convolutional Baseline). ECCV 2018.
+
+    Public keys:
+        - ``pcb_p4``: PCB with 4-part strips.
+        - ``pcb_p6``: PCB with 6-part strips.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        block,
+        layers,
+        parts=6,
+        reduced_dim=256,
+        nonlinear='relu',
+        **kwargs
+    ):
+        self.inplanes = 64
+        super(PCB, self).__init__()
+        self.loss = loss
+        self.parts = parts
+        self.feature_dim = 512 * block.expansion
+
+        # backbone network
+        self.conv1 = nn.Conv2d(
+            3, 64, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1)
+
+        # pcb layers
+        self.parts_avgpool = nn.AdaptiveAvgPool2d((self.parts, 1))
+        self.dropout = nn.Dropout(p=0.5)
+        self.conv5 = DimReduceLayer(
+            512 * block.expansion, reduced_dim, nonlinear=nonlinear
+        )
+        self.feature_dim = reduced_dim
+        self.classifier = nn.ModuleList(
+            [
+                nn.Linear(self.feature_dim, num_classes)
+                for _ in range(self.parts)
+            ]
+        )
+
+        self._init_params()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v_g = self.parts_avgpool(f)
+
+        if not self.training:
+            v_g = F.normalize(v_g, p=2, dim=1)
+            return v_g.view(v_g.size(0), -1)
+
+        v_g = self.dropout(v_g)
+        v_h = self.conv5(v_g)
+
+        y = []
+        for i in range(self.parts):
+            v_h_i = v_h[:, :, i, :]
+            v_h_i = v_h_i.view(v_h_i.size(0), -1)
+            y_i = self.classifier[i](v_h_i)
+            y.append(y_i)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            v_g = F.normalize(v_g, p=2, dim=1)
+            return y, v_g.view(v_g.size(0), -1)
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def pcb_p6(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = PCB(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        last_stride=1,
+        parts=6,
+        reduced_dim=256,
+        nonlinear='relu',
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model
+
+
+def pcb_p4(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = PCB(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        last_stride=1,
+        parts=4,
+        reduced_dim=256,
+        nonlinear='relu',
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model

trackers/strongsort/deep/models/resnet.py

@@ -0,0 +1,530 @@
+"""
+Code source: https://github.com/pytorch/vision
+"""
+from __future__ import division, absolute_import
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+
+__all__ = [
+    'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152',
+    'resnext50_32x4d', 'resnext101_32x8d', 'resnet50_fc512'
+]
+
+model_urls = {
+    'resnet18':
+    'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34':
+    'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50':
+    'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101':
+    'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152':
+    'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d':
+    'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d':
+    'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation
+    )
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        stride=1,
+        downsample=None,
+        groups=1,
+        base_width=64,
+        dilation=1,
+        norm_layer=None
+    ):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError(
+                'BasicBlock only supports groups=1 and base_width=64'
+            )
+        if dilation > 1:
+            raise NotImplementedError(
+                "Dilation > 1 not supported in BasicBlock"
+            )
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        stride=1,
+        downsample=None,
+        groups=1,
+        base_width=64,
+        dilation=1,
+        norm_layer=None
+    ):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width/64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    """Residual network.
+    
+    Reference:
+        - He et al. Deep Residual Learning for Image Recognition. CVPR 2016.
+        - Xie et al. Aggregated Residual Transformations for Deep Neural Networks. CVPR 2017.
+
+    Public keys:
+        - ``resnet18``: ResNet18.
+        - ``resnet34``: ResNet34.
+        - ``resnet50``: ResNet50.
+        - ``resnet101``: ResNet101.
+        - ``resnet152``: ResNet152.
+        - ``resnext50_32x4d``: ResNeXt50.
+        - ``resnext101_32x8d``: ResNeXt101.
+        - ``resnet50_fc512``: ResNet50 + FC.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        block,
+        layers,
+        zero_init_residual=False,
+        groups=1,
+        width_per_group=64,
+        replace_stride_with_dilation=None,
+        norm_layer=None,
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    ):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+        self.loss = loss
+        self.feature_dim = 512 * block.expansion
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError(
+                "replace_stride_with_dilation should be None "
+                "or a 3-element tuple, got {}".
+                format(replace_stride_with_dilation)
+            )
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(
+            3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(
+            block,
+            128,
+            layers[1],
+            stride=2,
+            dilate=replace_stride_with_dilation[0]
+        )
+        self.layer3 = self._make_layer(
+            block,
+            256,
+            layers[2],
+            stride=2,
+            dilate=replace_stride_with_dilation[1]
+        )
+        self.layer4 = self._make_layer(
+            block,
+            512,
+            layers[3],
+            stride=last_stride,
+            dilate=replace_stride_with_dilation[2]
+        )
+        self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = self._construct_fc_layer(
+            fc_dims, 512 * block.expansion, dropout_p
+        )
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes, planes, stride, downsample, self.groups,
+                self.base_width, previous_dilation, norm_layer
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if self.fc is not None:
+            v = self.fc(v)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+"""ResNet"""
+
+
+def resnet18(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=BasicBlock,
+        layers=[2, 2, 2, 2],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet18'])
+    return model
+
+
+def resnet34(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=BasicBlock,
+        layers=[3, 4, 6, 3],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet34'])
+    return model
+
+
+def resnet50(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model
+
+
+def resnet101(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 23, 3],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet101'])
+    return model
+
+
+def resnet152(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 8, 36, 3],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet152'])
+    return model
+
+
+"""ResNeXt"""
+
+
+def resnext50_32x4d(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        groups=32,
+        width_per_group=4,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnext50_32x4d'])
+    return model
+
+
+def resnext101_32x8d(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 23, 3],
+        last_stride=2,
+        fc_dims=None,
+        dropout_p=None,
+        groups=32,
+        width_per_group=8,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnext101_32x8d'])
+    return model
+
+
+"""
+ResNet + FC
+"""
+
+
+def resnet50_fc512(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNet(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        last_stride=1,
+        fc_dims=[512],
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model

trackers/strongsort/deep/models/resnet_ibn_a.py

@@ -0,0 +1,289 @@
+"""
+Credit to https://github.com/XingangPan/IBN-Net.
+"""
+from __future__ import division, absolute_import
+import math
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['resnet50_ibn_a']
+
+model_urls = {
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class IBN(nn.Module):
+
+    def __init__(self, planes):
+        super(IBN, self).__init__()
+        half1 = int(planes / 2)
+        self.half = half1
+        half2 = planes - half1
+        self.IN = nn.InstanceNorm2d(half1, affine=True)
+        self.BN = nn.BatchNorm2d(half2)
+
+    def forward(self, x):
+        split = torch.split(x, self.half, 1)
+        out1 = self.IN(split[0].contiguous())
+        out2 = self.BN(split[1].contiguous())
+        out = torch.cat((out1, out2), 1)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        if ibn:
+            self.bn1 = IBN(planes)
+        else:
+            self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, planes * self.expansion, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    """Residual network + IBN layer.
+    
+    Reference:
+        - He et al. Deep Residual Learning for Image Recognition. CVPR 2016.
+        - Pan et al. Two at Once: Enhancing Learning and Generalization
+          Capacities via IBN-Net. ECCV 2018.
+    """
+
+    def __init__(
+        self,
+        block,
+        layers,
+        num_classes=1000,
+        loss='softmax',
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    ):
+        scale = 64
+        self.inplanes = scale
+        super(ResNet, self).__init__()
+        self.loss = loss
+        self.feature_dim = scale * 8 * block.expansion
+
+        self.conv1 = nn.Conv2d(
+            3, scale, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(scale)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, scale, layers[0])
+        self.layer2 = self._make_layer(block, scale * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, scale * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, scale * 8, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = self._construct_fc_layer(
+            fc_dims, scale * 8 * block.expansion, dropout_p
+        )
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.InstanceNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        ibn = True
+        if planes == 512:
+            ibn = False
+        layers.append(block(self.inplanes, planes, ibn, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, ibn))
+
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.avgpool(f)
+        v = v.view(v.size(0), -1)
+        if self.fc is not None:
+            v = self.fc(v)
+        if not self.training:
+            return v
+        y = self.classifier(v)
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def resnet50_ibn_a(num_classes, loss='softmax', pretrained=False, **kwargs):
+    model = ResNet(
+        Bottleneck, [3, 4, 6, 3], num_classes=num_classes, loss=loss, **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model

trackers/strongsort/deep/models/resnet_ibn_b.py

@@ -0,0 +1,274 @@
+"""
+Credit to https://github.com/XingangPan/IBN-Net.
+"""
+from __future__ import division, absolute_import
+import math
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['resnet50_ibn_b']
+
+model_urls = {
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, IN=False):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, planes * self.expansion, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.IN = None
+        if IN:
+            self.IN = nn.InstanceNorm2d(planes * 4, affine=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        if self.IN is not None:
+            out = self.IN(out)
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    """Residual network + IBN layer.
+    
+    Reference:
+        - He et al. Deep Residual Learning for Image Recognition. CVPR 2016.
+        - Pan et al. Two at Once: Enhancing Learning and Generalization
+          Capacities via IBN-Net. ECCV 2018.
+    """
+
+    def __init__(
+        self,
+        block,
+        layers,
+        num_classes=1000,
+        loss='softmax',
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    ):
+        scale = 64
+        self.inplanes = scale
+        super(ResNet, self).__init__()
+        self.loss = loss
+        self.feature_dim = scale * 8 * block.expansion
+
+        self.conv1 = nn.Conv2d(
+            3, scale, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = nn.InstanceNorm2d(scale, affine=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(
+            block, scale, layers[0], stride=1, IN=True
+        )
+        self.layer2 = self._make_layer(
+            block, scale * 2, layers[1], stride=2, IN=True
+        )
+        self.layer3 = self._make_layer(block, scale * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, scale * 8, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = self._construct_fc_layer(
+            fc_dims, scale * 8 * block.expansion, dropout_p
+        )
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.InstanceNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1, IN=False):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks - 1):
+            layers.append(block(self.inplanes, planes))
+        layers.append(block(self.inplanes, planes, IN=IN))
+
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.avgpool(f)
+        v = v.view(v.size(0), -1)
+        if self.fc is not None:
+            v = self.fc(v)
+        if not self.training:
+            return v
+        y = self.classifier(v)
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def resnet50_ibn_b(num_classes, loss='softmax', pretrained=False, **kwargs):
+    model = ResNet(
+        Bottleneck, [3, 4, 6, 3], num_classes=num_classes, loss=loss, **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model

trackers/strongsort/deep/models/resnetmid.py

@@ -0,0 +1,307 @@
+from __future__ import division, absolute_import
+import torch
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+
+__all__ = ['resnet50mid']
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(
+            planes, planes * self.expansion, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNetMid(nn.Module):
+    """Residual network + mid-level features.
+    
+    Reference:
+        Yu et al. The Devil is in the Middle: Exploiting Mid-level Representations for
+        Cross-Domain Instance Matching. arXiv:1711.08106.
+
+    Public keys:
+        - ``resnet50mid``: ResNet50 + mid-level feature fusion.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        block,
+        layers,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    ):
+        self.inplanes = 64
+        super(ResNetMid, self).__init__()
+        self.loss = loss
+        self.feature_dim = 512 * block.expansion
+
+        # backbone network
+        self.conv1 = nn.Conv2d(
+            3, 64, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(
+            block, 512, layers[3], stride=last_stride
+        )
+
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        assert fc_dims is not None
+        self.fc_fusion = self._construct_fc_layer(
+            fc_dims, 512 * block.expansion * 2
+        )
+        self.feature_dim += 512 * block.expansion
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x4a = self.layer4[0](x)
+        x4b = self.layer4[1](x4a)
+        x4c = self.layer4[2](x4b)
+        return x4a, x4b, x4c
+
+    def forward(self, x):
+        x4a, x4b, x4c = self.featuremaps(x)
+
+        v4a = self.global_avgpool(x4a)
+        v4b = self.global_avgpool(x4b)
+        v4c = self.global_avgpool(x4c)
+        v4ab = torch.cat([v4a, v4b], 1)
+        v4ab = v4ab.view(v4ab.size(0), -1)
+        v4ab = self.fc_fusion(v4ab)
+        v4c = v4c.view(v4c.size(0), -1)
+        v = torch.cat([v4ab, v4c], 1)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+"""
+Residual network configurations:
+--
+resnet18: block=BasicBlock, layers=[2, 2, 2, 2]
+resnet34: block=BasicBlock, layers=[3, 4, 6, 3]
+resnet50: block=Bottleneck, layers=[3, 4, 6, 3]
+resnet101: block=Bottleneck, layers=[3, 4, 23, 3]
+resnet152: block=Bottleneck, layers=[3, 8, 36, 3]
+"""
+
+
+def resnet50mid(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ResNetMid(
+        num_classes=num_classes,
+        loss=loss,
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        last_stride=2,
+        fc_dims=[1024],
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['resnet50'])
+    return model

trackers/strongsort/deep/models/senet.py

@@ -0,0 +1,688 @@
+from __future__ import division, absolute_import
+import math
+from collections import OrderedDict
+import torch.nn as nn
+from torch.utils import model_zoo
+
+__all__ = [
+    'senet154', 'se_resnet50', 'se_resnet101', 'se_resnet152',
+    'se_resnext50_32x4d', 'se_resnext101_32x4d', 'se_resnet50_fc512'
+]
+"""
+Code imported from https://github.com/Cadene/pretrained-models.pytorch
+"""
+
+pretrained_settings = {
+    'senet154': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/senet154-c7b49a05.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.485, 0.456, 0.406],
+            'std': [0.229, 0.224, 0.225],
+            'num_classes': 1000
+        }
+    },
+    'se_resnet50': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet50-ce0d4300.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.485, 0.456, 0.406],
+            'std': [0.229, 0.224, 0.225],
+            'num_classes': 1000
+        }
+    },
+    'se_resnet101': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet101-7e38fcc6.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.485, 0.456, 0.406],
+            'std': [0.229, 0.224, 0.225],
+            'num_classes': 1000
+        }
+    },
+    'se_resnet152': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet152-d17c99b7.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.485, 0.456, 0.406],
+            'std': [0.229, 0.224, 0.225],
+            'num_classes': 1000
+        }
+    },
+    'se_resnext50_32x4d': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/se_resnext50_32x4d-a260b3a4.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.485, 0.456, 0.406],
+            'std': [0.229, 0.224, 0.225],
+            'num_classes': 1000
+        }
+    },
+    'se_resnext101_32x4d': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/se_resnext101_32x4d-3b2fe3d8.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 224, 224],
+            'input_range': [0, 1],
+            'mean': [0.485, 0.456, 0.406],
+            'std': [0.229, 0.224, 0.225],
+            'num_classes': 1000
+        }
+    },
+}
+
+
+class SEModule(nn.Module):
+
+    def __init__(self, channels, reduction):
+        super(SEModule, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc1 = nn.Conv2d(
+            channels, channels // reduction, kernel_size=1, padding=0
+        )
+        self.relu = nn.ReLU(inplace=True)
+        self.fc2 = nn.Conv2d(
+            channels // reduction, channels, kernel_size=1, padding=0
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        module_input = x
+        x = self.avg_pool(x)
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.sigmoid(x)
+        return module_input * x
+
+
+class Bottleneck(nn.Module):
+    """
+    Base class for bottlenecks that implements `forward()` method.
+    """
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out = self.se_module(out) + residual
+        out = self.relu(out)
+
+        return out
+
+
+class SEBottleneck(Bottleneck):
+    """
+    Bottleneck for SENet154.
+    """
+    expansion = 4
+
+    def __init__(
+        self, inplanes, planes, groups, reduction, stride=1, downsample=None
+    ):
+        super(SEBottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes * 2, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes * 2)
+        self.conv2 = nn.Conv2d(
+            planes * 2,
+            planes * 4,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            groups=groups,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes * 4)
+        self.conv3 = nn.Conv2d(
+            planes * 4, planes * 4, kernel_size=1, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.se_module = SEModule(planes * 4, reduction=reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+
+class SEResNetBottleneck(Bottleneck):
+    """
+    ResNet bottleneck with a Squeeze-and-Excitation module. It follows Caffe
+    implementation and uses `stride=stride` in `conv1` and not in `conv2`
+    (the latter is used in the torchvision implementation of ResNet).
+    """
+    expansion = 4
+
+    def __init__(
+        self, inplanes, planes, groups, reduction, stride=1, downsample=None
+    ):
+        super(SEResNetBottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(
+            inplanes, planes, kernel_size=1, bias=False, stride=stride
+        )
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(
+            planes,
+            planes,
+            kernel_size=3,
+            padding=1,
+            groups=groups,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.se_module = SEModule(planes * 4, reduction=reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+
+class SEResNeXtBottleneck(Bottleneck):
+    """ResNeXt bottleneck type C with a Squeeze-and-Excitation module"""
+    expansion = 4
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        groups,
+        reduction,
+        stride=1,
+        downsample=None,
+        base_width=4
+    ):
+        super(SEResNeXtBottleneck, self).__init__()
+        width = int(math.floor(planes * (base_width/64.)) * groups)
+        self.conv1 = nn.Conv2d(
+            inplanes, width, kernel_size=1, bias=False, stride=1
+        )
+        self.bn1 = nn.BatchNorm2d(width)
+        self.conv2 = nn.Conv2d(
+            width,
+            width,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            groups=groups,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(width)
+        self.conv3 = nn.Conv2d(width, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.se_module = SEModule(planes * 4, reduction=reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+
+class SENet(nn.Module):
+    """Squeeze-and-excitation network.
+    
+    Reference:
+        Hu et al. Squeeze-and-Excitation Networks. CVPR 2018.
+
+    Public keys:
+        - ``senet154``: SENet154.
+        - ``se_resnet50``: ResNet50 + SE.
+        - ``se_resnet101``: ResNet101 + SE.
+        - ``se_resnet152``: ResNet152 + SE.
+        - ``se_resnext50_32x4d``: ResNeXt50 (groups=32, width=4) + SE.
+        - ``se_resnext101_32x4d``: ResNeXt101 (groups=32, width=4) + SE.
+        - ``se_resnet50_fc512``: (ResNet50 + SE) + FC.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        block,
+        layers,
+        groups,
+        reduction,
+        dropout_p=0.2,
+        inplanes=128,
+        input_3x3=True,
+        downsample_kernel_size=3,
+        downsample_padding=1,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    ):
+        """
+        Parameters
+        ----------
+        block (nn.Module): Bottleneck class.
+            - For SENet154: SEBottleneck
+            - For SE-ResNet models: SEResNetBottleneck
+            - For SE-ResNeXt models:  SEResNeXtBottleneck
+        layers (list of ints): Number of residual blocks for 4 layers of the
+            network (layer1...layer4).
+        groups (int): Number of groups for the 3x3 convolution in each
+            bottleneck block.
+            - For SENet154: 64
+            - For SE-ResNet models: 1
+            - For SE-ResNeXt models:  32
+        reduction (int): Reduction ratio for Squeeze-and-Excitation modules.
+            - For all models: 16
+        dropout_p (float or None): Drop probability for the Dropout layer.
+            If `None` the Dropout layer is not used.
+            - For SENet154: 0.2
+            - For SE-ResNet models: None
+            - For SE-ResNeXt models: None
+        inplanes (int):  Number of input channels for layer1.
+            - For SENet154: 128
+            - For SE-ResNet models: 64
+            - For SE-ResNeXt models: 64
+        input_3x3 (bool): If `True`, use three 3x3 convolutions instead of
+            a single 7x7 convolution in layer0.
+            - For SENet154: True
+            - For SE-ResNet models: False
+            - For SE-ResNeXt models: False
+        downsample_kernel_size (int): Kernel size for downsampling convolutions
+            in layer2, layer3 and layer4.
+            - For SENet154: 3
+            - For SE-ResNet models: 1
+            - For SE-ResNeXt models: 1
+        downsample_padding (int): Padding for downsampling convolutions in
+            layer2, layer3 and layer4.
+            - For SENet154: 1
+            - For SE-ResNet models: 0
+            - For SE-ResNeXt models: 0
+        num_classes (int): Number of outputs in `classifier` layer.
+        """
+        super(SENet, self).__init__()
+        self.inplanes = inplanes
+        self.loss = loss
+
+        if input_3x3:
+            layer0_modules = [
+                (
+                    'conv1',
+                    nn.Conv2d(3, 64, 3, stride=2, padding=1, bias=False)
+                ),
+                ('bn1', nn.BatchNorm2d(64)),
+                ('relu1', nn.ReLU(inplace=True)),
+                (
+                    'conv2',
+                    nn.Conv2d(64, 64, 3, stride=1, padding=1, bias=False)
+                ),
+                ('bn2', nn.BatchNorm2d(64)),
+                ('relu2', nn.ReLU(inplace=True)),
+                (
+                    'conv3',
+                    nn.Conv2d(
+                        64, inplanes, 3, stride=1, padding=1, bias=False
+                    )
+                ),
+                ('bn3', nn.BatchNorm2d(inplanes)),
+                ('relu3', nn.ReLU(inplace=True)),
+            ]
+        else:
+            layer0_modules = [
+                (
+                    'conv1',
+                    nn.Conv2d(
+                        3,
+                        inplanes,
+                        kernel_size=7,
+                        stride=2,
+                        padding=3,
+                        bias=False
+                    )
+                ),
+                ('bn1', nn.BatchNorm2d(inplanes)),
+                ('relu1', nn.ReLU(inplace=True)),
+            ]
+        # To preserve compatibility with Caffe weights `ceil_mode=True`
+        # is used instead of `padding=1`.
+        layer0_modules.append(
+            ('pool', nn.MaxPool2d(3, stride=2, ceil_mode=True))
+        )
+        self.layer0 = nn.Sequential(OrderedDict(layer0_modules))
+        self.layer1 = self._make_layer(
+            block,
+            planes=64,
+            blocks=layers[0],
+            groups=groups,
+            reduction=reduction,
+            downsample_kernel_size=1,
+            downsample_padding=0
+        )
+        self.layer2 = self._make_layer(
+            block,
+            planes=128,
+            blocks=layers[1],
+            stride=2,
+            groups=groups,
+            reduction=reduction,
+            downsample_kernel_size=downsample_kernel_size,
+            downsample_padding=downsample_padding
+        )
+        self.layer3 = self._make_layer(
+            block,
+            planes=256,
+            blocks=layers[2],
+            stride=2,
+            groups=groups,
+            reduction=reduction,
+            downsample_kernel_size=downsample_kernel_size,
+            downsample_padding=downsample_padding
+        )
+        self.layer4 = self._make_layer(
+            block,
+            planes=512,
+            blocks=layers[3],
+            stride=last_stride,
+            groups=groups,
+            reduction=reduction,
+            downsample_kernel_size=downsample_kernel_size,
+            downsample_padding=downsample_padding
+        )
+
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = self._construct_fc_layer(
+            fc_dims, 512 * block.expansion, dropout_p
+        )
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+    def _make_layer(
+        self,
+        block,
+        planes,
+        blocks,
+        groups,
+        reduction,
+        stride=1,
+        downsample_kernel_size=1,
+        downsample_padding=0
+    ):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=downsample_kernel_size,
+                    stride=stride,
+                    padding=downsample_padding,
+                    bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes, planes, groups, reduction, stride, downsample
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups, reduction))
+
+        return nn.Sequential(*layers)
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """
+        Construct fully connected layer
+
+        - fc_dims (list or tuple): dimensions of fc layers, if None,
+                                   no fc layers are constructed
+        - input_dim (int): input dimension
+        - dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def featuremaps(self, x):
+        x = self.layer0(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if self.fc is not None:
+            v = self.fc(v)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def senet154(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEBottleneck,
+        layers=[3, 8, 36, 3],
+        groups=64,
+        reduction=16,
+        dropout_p=0.2,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['senet154']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model
+
+
+def se_resnet50(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEResNetBottleneck,
+        layers=[3, 4, 6, 3],
+        groups=1,
+        reduction=16,
+        dropout_p=None,
+        inplanes=64,
+        input_3x3=False,
+        downsample_kernel_size=1,
+        downsample_padding=0,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['se_resnet50']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model
+
+
+def se_resnet50_fc512(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEResNetBottleneck,
+        layers=[3, 4, 6, 3],
+        groups=1,
+        reduction=16,
+        dropout_p=None,
+        inplanes=64,
+        input_3x3=False,
+        downsample_kernel_size=1,
+        downsample_padding=0,
+        last_stride=1,
+        fc_dims=[512],
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['se_resnet50']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model
+
+
+def se_resnet101(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEResNetBottleneck,
+        layers=[3, 4, 23, 3],
+        groups=1,
+        reduction=16,
+        dropout_p=None,
+        inplanes=64,
+        input_3x3=False,
+        downsample_kernel_size=1,
+        downsample_padding=0,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['se_resnet101']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model
+
+
+def se_resnet152(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEResNetBottleneck,
+        layers=[3, 8, 36, 3],
+        groups=1,
+        reduction=16,
+        dropout_p=None,
+        inplanes=64,
+        input_3x3=False,
+        downsample_kernel_size=1,
+        downsample_padding=0,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['se_resnet152']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model
+
+
+def se_resnext50_32x4d(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEResNeXtBottleneck,
+        layers=[3, 4, 6, 3],
+        groups=32,
+        reduction=16,
+        dropout_p=None,
+        inplanes=64,
+        input_3x3=False,
+        downsample_kernel_size=1,
+        downsample_padding=0,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['se_resnext50_32x4d']['imagenet']['url'
+                                                                          ]
+        init_pretrained_weights(model, model_url)
+    return model
+
+
+def se_resnext101_32x4d(
+    num_classes, loss='softmax', pretrained=True, **kwargs
+):
+    model = SENet(
+        num_classes=num_classes,
+        loss=loss,
+        block=SEResNeXtBottleneck,
+        layers=[3, 4, 23, 3],
+        groups=32,
+        reduction=16,
+        dropout_p=None,
+        inplanes=64,
+        input_3x3=False,
+        downsample_kernel_size=1,
+        downsample_padding=0,
+        last_stride=2,
+        fc_dims=None,
+        **kwargs
+    )
+    if pretrained:
+        model_url = pretrained_settings['se_resnext101_32x4d']['imagenet'][
+            'url']
+        init_pretrained_weights(model, model_url)
+    return model

trackers/strongsort/deep/models/shufflenet.py

@@ -0,0 +1,198 @@
+from __future__ import division, absolute_import
+import torch
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+from torch.nn import functional as F
+
+__all__ = ['shufflenet']
+
+model_urls = {
+    # training epoch = 90, top1 = 61.8
+    'imagenet':
+    'https://mega.nz/#!RDpUlQCY!tr_5xBEkelzDjveIYBBcGcovNCOrgfiJO9kiidz9fZM',
+}
+
+
+class ChannelShuffle(nn.Module):
+
+    def __init__(self, num_groups):
+        super(ChannelShuffle, self).__init__()
+        self.g = num_groups
+
+    def forward(self, x):
+        b, c, h, w = x.size()
+        n = c // self.g
+        # reshape
+        x = x.view(b, self.g, n, h, w)
+        # transpose
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+        # flatten
+        x = x.view(b, c, h, w)
+        return x
+
+
+class Bottleneck(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        num_groups,
+        group_conv1x1=True
+    ):
+        super(Bottleneck, self).__init__()
+        assert stride in [1, 2], 'Warning: stride must be either 1 or 2'
+        self.stride = stride
+        mid_channels = out_channels // 4
+        if stride == 2:
+            out_channels -= in_channels
+        # group conv is not applied to first conv1x1 at stage 2
+        num_groups_conv1x1 = num_groups if group_conv1x1 else 1
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            mid_channels,
+            1,
+            groups=num_groups_conv1x1,
+            bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(mid_channels)
+        self.shuffle1 = ChannelShuffle(num_groups)
+        self.conv2 = nn.Conv2d(
+            mid_channels,
+            mid_channels,
+            3,
+            stride=stride,
+            padding=1,
+            groups=mid_channels,
+            bias=False
+        )
+        self.bn2 = nn.BatchNorm2d(mid_channels)
+        self.conv3 = nn.Conv2d(
+            mid_channels, out_channels, 1, groups=num_groups, bias=False
+        )
+        self.bn3 = nn.BatchNorm2d(out_channels)
+        if stride == 2:
+            self.shortcut = nn.AvgPool2d(3, stride=2, padding=1)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.shuffle1(out)
+        out = self.bn2(self.conv2(out))
+        out = self.bn3(self.conv3(out))
+        if self.stride == 2:
+            res = self.shortcut(x)
+            out = F.relu(torch.cat([res, out], 1))
+        else:
+            out = F.relu(x + out)
+        return out
+
+
+# configuration of (num_groups: #out_channels) based on Table 1 in the paper
+cfg = {
+    1: [144, 288, 576],
+    2: [200, 400, 800],
+    3: [240, 480, 960],
+    4: [272, 544, 1088],
+    8: [384, 768, 1536],
+}
+
+
+class ShuffleNet(nn.Module):
+    """ShuffleNet.
+
+    Reference:
+        Zhang et al. ShuffleNet: An Extremely Efficient Convolutional Neural
+        Network for Mobile Devices. CVPR 2018.
+
+    Public keys:
+        - ``shufflenet``: ShuffleNet (groups=3).
+    """
+
+    def __init__(self, num_classes, loss='softmax', num_groups=3, **kwargs):
+        super(ShuffleNet, self).__init__()
+        self.loss = loss
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(3, 24, 3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(24),
+            nn.ReLU(),
+            nn.MaxPool2d(3, stride=2, padding=1),
+        )
+
+        self.stage2 = nn.Sequential(
+            Bottleneck(
+                24, cfg[num_groups][0], 2, num_groups, group_conv1x1=False
+            ),
+            Bottleneck(cfg[num_groups][0], cfg[num_groups][0], 1, num_groups),
+            Bottleneck(cfg[num_groups][0], cfg[num_groups][0], 1, num_groups),
+            Bottleneck(cfg[num_groups][0], cfg[num_groups][0], 1, num_groups),
+        )
+
+        self.stage3 = nn.Sequential(
+            Bottleneck(cfg[num_groups][0], cfg[num_groups][1], 2, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),
+        )
+
+        self.stage4 = nn.Sequential(
+            Bottleneck(cfg[num_groups][1], cfg[num_groups][2], 2, num_groups),
+            Bottleneck(cfg[num_groups][2], cfg[num_groups][2], 1, num_groups),
+            Bottleneck(cfg[num_groups][2], cfg[num_groups][2], 1, num_groups),
+            Bottleneck(cfg[num_groups][2], cfg[num_groups][2], 1, num_groups),
+        )
+
+        self.classifier = nn.Linear(cfg[num_groups][2], num_classes)
+        self.feat_dim = cfg[num_groups][2]
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.stage4(x)
+        x = F.avg_pool2d(x, x.size()[2:]).view(x.size(0), -1)
+
+        if not self.training:
+            return x
+
+        y = self.classifier(x)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, x
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def shufflenet(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ShuffleNet(num_classes, loss, **kwargs)
+    if pretrained:
+        # init_pretrained_weights(model, model_urls['imagenet'])
+        import warnings
+        warnings.warn(
+            'The imagenet pretrained weights need to be manually downloaded from {}'
+            .format(model_urls['imagenet'])
+        )
+    return model

trackers/strongsort/deep/models/shufflenetv2.py

@@ -0,0 +1,262 @@
+"""
+Code source: https://github.com/pytorch/vision
+"""
+from __future__ import division, absolute_import
+import torch
+import torch.utils.model_zoo as model_zoo
+from torch import nn
+
+__all__ = [
+    'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0', 'shufflenet_v2_x1_5',
+    'shufflenet_v2_x2_0'
+]
+
+model_urls = {
+    'shufflenetv2_x0.5':
+    'https://download.pytorch.org/models/shufflenetv2_x0.5-f707e7126e.pth',
+    'shufflenetv2_x1.0':
+    'https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth',
+    'shufflenetv2_x1.5': None,
+    'shufflenetv2_x2.0': None,
+}
+
+
+def channel_shuffle(x, groups):
+    batchsize, num_channels, height, width = x.data.size()
+    channels_per_group = num_channels // groups
+
+    # reshape
+    x = x.view(batchsize, groups, channels_per_group, height, width)
+
+    x = torch.transpose(x, 1, 2).contiguous()
+
+    # flatten
+    x = x.view(batchsize, -1, height, width)
+
+    return x
+
+
+class InvertedResidual(nn.Module):
+
+    def __init__(self, inp, oup, stride):
+        super(InvertedResidual, self).__init__()
+
+        if not (1 <= stride <= 3):
+            raise ValueError('illegal stride value')
+        self.stride = stride
+
+        branch_features = oup // 2
+        assert (self.stride != 1) or (inp == branch_features << 1)
+
+        if self.stride > 1:
+            self.branch1 = nn.Sequential(
+                self.depthwise_conv(
+                    inp, inp, kernel_size=3, stride=self.stride, padding=1
+                ),
+                nn.BatchNorm2d(inp),
+                nn.Conv2d(
+                    inp,
+                    branch_features,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    bias=False
+                ),
+                nn.BatchNorm2d(branch_features),
+                nn.ReLU(inplace=True),
+            )
+
+        self.branch2 = nn.Sequential(
+            nn.Conv2d(
+                inp if (self.stride > 1) else branch_features,
+                branch_features,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=False
+            ),
+            nn.BatchNorm2d(branch_features),
+            nn.ReLU(inplace=True),
+            self.depthwise_conv(
+                branch_features,
+                branch_features,
+                kernel_size=3,
+                stride=self.stride,
+                padding=1
+            ),
+            nn.BatchNorm2d(branch_features),
+            nn.Conv2d(
+                branch_features,
+                branch_features,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=False
+            ),
+            nn.BatchNorm2d(branch_features),
+            nn.ReLU(inplace=True),
+        )
+
+    @staticmethod
+    def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
+        return nn.Conv2d(
+            i, o, kernel_size, stride, padding, bias=bias, groups=i
+        )
+
+    def forward(self, x):
+        if self.stride == 1:
+            x1, x2 = x.chunk(2, dim=1)
+            out = torch.cat((x1, self.branch2(x2)), dim=1)
+        else:
+            out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
+
+        out = channel_shuffle(out, 2)
+
+        return out
+
+
+class ShuffleNetV2(nn.Module):
+    """ShuffleNetV2.
+    
+    Reference:
+        Ma et al. ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design. ECCV 2018.
+
+    Public keys:
+        - ``shufflenet_v2_x0_5``: ShuffleNetV2 x0.5.
+        - ``shufflenet_v2_x1_0``: ShuffleNetV2 x1.0.
+        - ``shufflenet_v2_x1_5``: ShuffleNetV2 x1.5.
+        - ``shufflenet_v2_x2_0``: ShuffleNetV2 x2.0.
+    """
+
+    def __init__(
+        self, num_classes, loss, stages_repeats, stages_out_channels, **kwargs
+    ):
+        super(ShuffleNetV2, self).__init__()
+        self.loss = loss
+
+        if len(stages_repeats) != 3:
+            raise ValueError(
+                'expected stages_repeats as list of 3 positive ints'
+            )
+        if len(stages_out_channels) != 5:
+            raise ValueError(
+                'expected stages_out_channels as list of 5 positive ints'
+            )
+        self._stage_out_channels = stages_out_channels
+
+        input_channels = 3
+        output_channels = self._stage_out_channels[0]
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False),
+            nn.BatchNorm2d(output_channels),
+            nn.ReLU(inplace=True),
+        )
+        input_channels = output_channels
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        stage_names = ['stage{}'.format(i) for i in [2, 3, 4]]
+        for name, repeats, output_channels in zip(
+            stage_names, stages_repeats, self._stage_out_channels[1:]
+        ):
+            seq = [InvertedResidual(input_channels, output_channels, 2)]
+            for i in range(repeats - 1):
+                seq.append(
+                    InvertedResidual(output_channels, output_channels, 1)
+                )
+            setattr(self, name, nn.Sequential(*seq))
+            input_channels = output_channels
+
+        output_channels = self._stage_out_channels[-1]
+        self.conv5 = nn.Sequential(
+            nn.Conv2d(input_channels, output_channels, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(output_channels),
+            nn.ReLU(inplace=True),
+        )
+        self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        self.classifier = nn.Linear(output_channels, num_classes)
+
+    def featuremaps(self, x):
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        x = self.stage2(x)
+        x = self.stage3(x)
+        x = self.stage4(x)
+        x = self.conv5(x)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError("Unsupported loss: {}".format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    if model_url is None:
+        import warnings
+        warnings.warn(
+            'ImageNet pretrained weights are unavailable for this model'
+        )
+        return
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def shufflenet_v2_x0_5(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ShuffleNetV2(
+        num_classes, loss, [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['shufflenetv2_x0.5'])
+    return model
+
+
+def shufflenet_v2_x1_0(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ShuffleNetV2(
+        num_classes, loss, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['shufflenetv2_x1.0'])
+    return model
+
+
+def shufflenet_v2_x1_5(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ShuffleNetV2(
+        num_classes, loss, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['shufflenetv2_x1.5'])
+    return model
+
+
+def shufflenet_v2_x2_0(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = ShuffleNetV2(
+        num_classes, loss, [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['shufflenetv2_x2.0'])
+    return model

trackers/strongsort/deep/models/squeezenet.py

@@ -0,0 +1,236 @@
+"""
+Code source: https://github.com/pytorch/vision
+"""
+from __future__ import division, absolute_import
+import torch
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['squeezenet1_0', 'squeezenet1_1', 'squeezenet1_0_fc512']
+
+model_urls = {
+    'squeezenet1_0':
+    'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth',
+    'squeezenet1_1':
+    'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth',
+}
+
+
+class Fire(nn.Module):
+
+    def __init__(
+        self, inplanes, squeeze_planes, expand1x1_planes, expand3x3_planes
+    ):
+        super(Fire, self).__init__()
+        self.inplanes = inplanes
+        self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)
+        self.squeeze_activation = nn.ReLU(inplace=True)
+        self.expand1x1 = nn.Conv2d(
+            squeeze_planes, expand1x1_planes, kernel_size=1
+        )
+        self.expand1x1_activation = nn.ReLU(inplace=True)
+        self.expand3x3 = nn.Conv2d(
+            squeeze_planes, expand3x3_planes, kernel_size=3, padding=1
+        )
+        self.expand3x3_activation = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.squeeze_activation(self.squeeze(x))
+        return torch.cat(
+            [
+                self.expand1x1_activation(self.expand1x1(x)),
+                self.expand3x3_activation(self.expand3x3(x))
+            ], 1
+        )
+
+
+class SqueezeNet(nn.Module):
+    """SqueezeNet.
+
+    Reference:
+        Iandola et al. SqueezeNet: AlexNet-level accuracy with 50x fewer parameters
+        and< 0.5 MB model size. arXiv:1602.07360.
+
+    Public keys:
+        - ``squeezenet1_0``: SqueezeNet (version=1.0).
+        - ``squeezenet1_1``: SqueezeNet (version=1.1).
+        - ``squeezenet1_0_fc512``: SqueezeNet (version=1.0) + FC.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        loss,
+        version=1.0,
+        fc_dims=None,
+        dropout_p=None,
+        **kwargs
+    ):
+        super(SqueezeNet, self).__init__()
+        self.loss = loss
+        self.feature_dim = 512
+
+        if version not in [1.0, 1.1]:
+            raise ValueError(
+                'Unsupported SqueezeNet version {version}:'
+                '1.0 or 1.1 expected'.format(version=version)
+            )
+
+        if version == 1.0:
+            self.features = nn.Sequential(
+                nn.Conv2d(3, 96, kernel_size=7, stride=2),
+                nn.ReLU(inplace=True),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(96, 16, 64, 64),
+                Fire(128, 16, 64, 64),
+                Fire(128, 32, 128, 128),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(256, 32, 128, 128),
+                Fire(256, 48, 192, 192),
+                Fire(384, 48, 192, 192),
+                Fire(384, 64, 256, 256),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(512, 64, 256, 256),
+            )
+        else:
+            self.features = nn.Sequential(
+                nn.Conv2d(3, 64, kernel_size=3, stride=2),
+                nn.ReLU(inplace=True),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(64, 16, 64, 64),
+                Fire(128, 16, 64, 64),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(128, 32, 128, 128),
+                Fire(256, 32, 128, 128),
+                nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+                Fire(256, 48, 192, 192),
+                Fire(384, 48, 192, 192),
+                Fire(384, 64, 256, 256),
+                Fire(512, 64, 256, 256),
+            )
+
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = self._construct_fc_layer(fc_dims, 512, dropout_p)
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        f = self.features(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if self.fc is not None:
+            v = self.fc(v)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initializes model with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url, map_location=None)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def squeezenet1_0(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SqueezeNet(
+        num_classes, loss, version=1.0, fc_dims=None, dropout_p=None, **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['squeezenet1_0'])
+    return model
+
+
+def squeezenet1_0_fc512(
+    num_classes, loss='softmax', pretrained=True, **kwargs
+):
+    model = SqueezeNet(
+        num_classes,
+        loss,
+        version=1.0,
+        fc_dims=[512],
+        dropout_p=None,
+        **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['squeezenet1_0'])
+    return model
+
+
+def squeezenet1_1(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = SqueezeNet(
+        num_classes, loss, version=1.1, fc_dims=None, dropout_p=None, **kwargs
+    )
+    if pretrained:
+        init_pretrained_weights(model, model_urls['squeezenet1_1'])
+    return model

trackers/strongsort/deep/models/xception.py

@@ -0,0 +1,344 @@
+from __future__ import division, absolute_import
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ['xception']
+
+pretrained_settings = {
+    'xception': {
+        'imagenet': {
+            'url':
+            'http://data.lip6.fr/cadene/pretrainedmodels/xception-43020ad28.pth',
+            'input_space': 'RGB',
+            'input_size': [3, 299, 299],
+            'input_range': [0, 1],
+            'mean': [0.5, 0.5, 0.5],
+            'std': [0.5, 0.5, 0.5],
+            'num_classes': 1000,
+            'scale':
+            0.8975 # The resize parameter of the validation transform should be 333, and make sure to center crop at 299x299
+        }
+    }
+}
+
+
+class SeparableConv2d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=1,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=False
+    ):
+        super(SeparableConv2d, self).__init__()
+
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups=in_channels,
+            bias=bias
+        )
+        self.pointwise = nn.Conv2d(
+            in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.pointwise(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(
+        self,
+        in_filters,
+        out_filters,
+        reps,
+        strides=1,
+        start_with_relu=True,
+        grow_first=True
+    ):
+        super(Block, self).__init__()
+
+        if out_filters != in_filters or strides != 1:
+            self.skip = nn.Conv2d(
+                in_filters, out_filters, 1, stride=strides, bias=False
+            )
+            self.skipbn = nn.BatchNorm2d(out_filters)
+        else:
+            self.skip = None
+
+        self.relu = nn.ReLU(inplace=True)
+        rep = []
+
+        filters = in_filters
+        if grow_first:
+            rep.append(self.relu)
+            rep.append(
+                SeparableConv2d(
+                    in_filters,
+                    out_filters,
+                    3,
+                    stride=1,
+                    padding=1,
+                    bias=False
+                )
+            )
+            rep.append(nn.BatchNorm2d(out_filters))
+            filters = out_filters
+
+        for i in range(reps - 1):
+            rep.append(self.relu)
+            rep.append(
+                SeparableConv2d(
+                    filters, filters, 3, stride=1, padding=1, bias=False
+                )
+            )
+            rep.append(nn.BatchNorm2d(filters))
+
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(
+                SeparableConv2d(
+                    in_filters,
+                    out_filters,
+                    3,
+                    stride=1,
+                    padding=1,
+                    bias=False
+                )
+            )
+            rep.append(nn.BatchNorm2d(out_filters))
+
+        if not start_with_relu:
+            rep = rep[1:]
+        else:
+            rep[0] = nn.ReLU(inplace=False)
+
+        if strides != 1:
+            rep.append(nn.MaxPool2d(3, strides, 1))
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, inp):
+        x = self.rep(inp)
+
+        if self.skip is not None:
+            skip = self.skip(inp)
+            skip = self.skipbn(skip)
+        else:
+            skip = inp
+
+        x += skip
+        return x
+
+
+class Xception(nn.Module):
+    """Xception.
+    
+    Reference:
+        Chollet. Xception: Deep Learning with Depthwise
+        Separable Convolutions. CVPR 2017.
+
+    Public keys:
+        - ``xception``: Xception.
+    """
+
+    def __init__(
+        self, num_classes, loss, fc_dims=None, dropout_p=None, **kwargs
+    ):
+        super(Xception, self).__init__()
+        self.loss = loss
+
+        self.conv1 = nn.Conv2d(3, 32, 3, 2, 0, bias=False)
+        self.bn1 = nn.BatchNorm2d(32)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, bias=False)
+        self.bn2 = nn.BatchNorm2d(64)
+
+        self.block1 = Block(
+            64, 128, 2, 2, start_with_relu=False, grow_first=True
+        )
+        self.block2 = Block(
+            128, 256, 2, 2, start_with_relu=True, grow_first=True
+        )
+        self.block3 = Block(
+            256, 728, 2, 2, start_with_relu=True, grow_first=True
+        )
+
+        self.block4 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+        self.block5 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+        self.block6 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+        self.block7 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+
+        self.block8 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+        self.block9 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+        self.block10 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+        self.block11 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True
+        )
+
+        self.block12 = Block(
+            728, 1024, 2, 2, start_with_relu=True, grow_first=False
+        )
+
+        self.conv3 = SeparableConv2d(1024, 1536, 3, 1, 1)
+        self.bn3 = nn.BatchNorm2d(1536)
+
+        self.conv4 = SeparableConv2d(1536, 2048, 3, 1, 1)
+        self.bn4 = nn.BatchNorm2d(2048)
+
+        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
+        self.feature_dim = 2048
+        self.fc = self._construct_fc_layer(fc_dims, 2048, dropout_p)
+        self.classifier = nn.Linear(self.feature_dim, num_classes)
+
+        self._init_params()
+
+    def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):
+        """Constructs fully connected layer.
+
+        Args:
+            fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed
+            input_dim (int): input dimension
+            dropout_p (float): dropout probability, if None, dropout is unused
+        """
+        if fc_dims is None:
+            self.feature_dim = input_dim
+            return None
+
+        assert isinstance(
+            fc_dims, (list, tuple)
+        ), 'fc_dims must be either list or tuple, but got {}'.format(
+            type(fc_dims)
+        )
+
+        layers = []
+        for dim in fc_dims:
+            layers.append(nn.Linear(input_dim, dim))
+            layers.append(nn.BatchNorm1d(dim))
+            layers.append(nn.ReLU(inplace=True))
+            if dropout_p is not None:
+                layers.append(nn.Dropout(p=dropout_p))
+            input_dim = dim
+
+        self.feature_dim = fc_dims[-1]
+
+        return nn.Sequential(*layers)
+
+    def _init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='relu'
+                )
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def featuremaps(self, input):
+        x = self.conv1(input)
+        x = self.bn1(x)
+        x = F.relu(x, inplace=True)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = F.relu(x, inplace=True)
+
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        x = self.block4(x)
+        x = self.block5(x)
+        x = self.block6(x)
+        x = self.block7(x)
+        x = self.block8(x)
+        x = self.block9(x)
+        x = self.block10(x)
+        x = self.block11(x)
+        x = self.block12(x)
+
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = F.relu(x, inplace=True)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = F.relu(x, inplace=True)
+        return x
+
+    def forward(self, x):
+        f = self.featuremaps(x)
+        v = self.global_avgpool(f)
+        v = v.view(v.size(0), -1)
+
+        if self.fc is not None:
+            v = self.fc(v)
+
+        if not self.training:
+            return v
+
+        y = self.classifier(v)
+
+        if self.loss == 'softmax':
+            return y
+        elif self.loss == 'triplet':
+            return y, v
+        else:
+            raise KeyError('Unsupported loss: {}'.format(self.loss))
+
+
+def init_pretrained_weights(model, model_url):
+    """Initialize models with pretrained weights.
+    
+    Layers that don't match with pretrained layers in name or size are kept unchanged.
+    """
+    pretrain_dict = model_zoo.load_url(model_url)
+    model_dict = model.state_dict()
+    pretrain_dict = {
+        k: v
+        for k, v in pretrain_dict.items()
+        if k in model_dict and model_dict[k].size() == v.size()
+    }
+    model_dict.update(pretrain_dict)
+    model.load_state_dict(model_dict)
+
+
+def xception(num_classes, loss='softmax', pretrained=True, **kwargs):
+    model = Xception(num_classes, loss, fc_dims=None, dropout_p=None, **kwargs)
+    if pretrained:
+        model_url = pretrained_settings['xception']['imagenet']['url']
+        init_pretrained_weights(model, model_url)
+    return model

trackers/strongsort/deep/reid_model_factory.py

@@ -0,0 +1,215 @@
+import torch
+from collections import OrderedDict
+
+
+
+__model_types = [
+    'resnet50', 'mlfn', 'hacnn', 'mobilenetv2_x1_0', 'mobilenetv2_x1_4',
+    'osnet_x1_0', 'osnet_x0_75', 'osnet_x0_5', 'osnet_x0_25',
+    'osnet_ibn_x1_0', 'osnet_ain_x1_0']
+
+__trained_urls = {
+
+    # market1501 models ########################################################
+    'resnet50_market1501.pt':
+    'https://drive.google.com/uc?id=1dUUZ4rHDWohmsQXCRe2C_HbYkzz94iBV',
+    'resnet50_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=17ymnLglnc64NRvGOitY3BqMRS9UWd1wg',
+    'resnet50_msmt17.pt':
+    'https://drive.google.com/uc?id=1ep7RypVDOthCRIAqDnn4_N-UhkkFHJsj',
+
+    'resnet50_fc512_market1501.pt':
+    'https://drive.google.com/uc?id=1kv8l5laX_YCdIGVCetjlNdzKIA3NvsSt',
+    'resnet50_fc512_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=13QN8Mp3XH81GK4BPGXobKHKyTGH50Rtx',
+    'resnet50_fc512_msmt17.pt':
+    'https://drive.google.com/uc?id=1fDJLcz4O5wxNSUvImIIjoaIF9u1Rwaud',
+
+    'mlfn_market1501.pt':
+    'https://drive.google.com/uc?id=1wXcvhA_b1kpDfrt9s2Pma-MHxtj9pmvS',
+    'mlfn_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1rExgrTNb0VCIcOnXfMsbwSUW1h2L1Bum',
+    'mlfn_msmt17.pt':
+    'https://drive.google.com/uc?id=18JzsZlJb3Wm7irCbZbZ07TN4IFKvR6p-',
+
+    'hacnn_market1501.pt':
+    'https://drive.google.com/uc?id=1LRKIQduThwGxMDQMiVkTScBwR7WidmYF',
+    'hacnn_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1zNm6tP4ozFUCUQ7Sv1Z98EAJWXJEhtYH',
+    'hacnn_msmt17.pt':
+    'https://drive.google.com/uc?id=1MsKRtPM5WJ3_Tk2xC0aGOO7pM3VaFDNZ',
+
+    'mobilenetv2_x1_0_market1501.pt':
+    'https://drive.google.com/uc?id=18DgHC2ZJkjekVoqBWszD8_Xiikz-fewp',
+    'mobilenetv2_x1_0_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1q1WU2FETRJ3BXcpVtfJUuqq4z3psetds',
+    'mobilenetv2_x1_0_msmt17.pt':
+    'https://drive.google.com/uc?id=1j50Hv14NOUAg7ZeB3frzfX-WYLi7SrhZ',
+
+    'mobilenetv2_x1_4_market1501.pt':
+    'https://drive.google.com/uc?id=1t6JCqphJG-fwwPVkRLmGGyEBhGOf2GO5',
+    'mobilenetv2_x1_4_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=12uD5FeVqLg9-AFDju2L7SQxjmPb4zpBN',
+    'mobilenetv2_x1_4_msmt17.pt':
+    'https://drive.google.com/uc?id=1ZY5P2Zgm-3RbDpbXM0kIBMPvspeNIbXz',
+
+    'osnet_x1_0_market1501.pt':
+    'https://drive.google.com/uc?id=1vduhq5DpN2q1g4fYEZfPI17MJeh9qyrA',
+    'osnet_x1_0_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1QZO_4sNf4hdOKKKzKc-TZU9WW1v6zQbq',
+    'osnet_x1_0_msmt17.pt':
+    'https://drive.google.com/uc?id=112EMUfBPYeYg70w-syK6V6Mx8-Qb9Q1M',
+
+    'osnet_x0_75_market1501.pt':
+    'https://drive.google.com/uc?id=1ozRaDSQw_EQ8_93OUmjDbvLXw9TnfPer',
+    'osnet_x0_75_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1IE3KRaTPp4OUa6PGTFL_d5_KQSJbP0Or',
+    'osnet_x0_75_msmt17.pt':
+    'https://drive.google.com/uc?id=1QEGO6WnJ-BmUzVPd3q9NoaO_GsPNlmWc',
+
+    'osnet_x0_5_market1501.pt':
+    'https://drive.google.com/uc?id=1PLB9rgqrUM7blWrg4QlprCuPT7ILYGKT',
+    'osnet_x0_5_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1KoUVqmiST175hnkALg9XuTi1oYpqcyTu',
+    'osnet_x0_5_msmt17.pt':
+    'https://drive.google.com/uc?id=1UT3AxIaDvS2PdxzZmbkLmjtiqq7AIKCv',
+
+    'osnet_x0_25_market1501.pt':
+    'https://drive.google.com/uc?id=1z1UghYvOTtjx7kEoRfmqSMu-z62J6MAj',
+    'osnet_x0_25_dukemtmcreid.pt':
+    'https://drive.google.com/uc?id=1eumrtiXT4NOspjyEV4j8cHmlOaaCGk5l',
+    'osnet_x0_25_msmt17.pt':
+    'https://drive.google.com/uc?id=1sSwXSUlj4_tHZequ_iZ8w_Jh0VaRQMqF',
+
+    ####### market1501 models ##################################################
+    'resnet50_msmt17.pt':
+    'https://drive.google.com/uc?id=1yiBteqgIZoOeywE8AhGmEQl7FTVwrQmf',
+    'osnet_x1_0_msmt17.pt':
+    'https://drive.google.com/uc?id=1IosIFlLiulGIjwW3H8uMRmx3MzPwf86x',
+    'osnet_x0_75_msmt17.pt':
+    'https://drive.google.com/uc?id=1fhjSS_7SUGCioIf2SWXaRGPqIY9j7-uw',
+
+    'osnet_x0_5_msmt17.pt':
+    'https://drive.google.com/uc?id=1DHgmb6XV4fwG3n-CnCM0zdL9nMsZ9_RF',
+    'osnet_x0_25_msmt17.pt':
+    'https://drive.google.com/uc?id=1Kkx2zW89jq_NETu4u42CFZTMVD5Hwm6e',
+    'osnet_ibn_x1_0_msmt17.pt':
+    'https://drive.google.com/uc?id=1q3Sj2ii34NlfxA4LvmHdWO_75NDRmECJ',
+    'osnet_ain_x1_0_msmt17.pt':
+    'https://drive.google.com/uc?id=1SigwBE6mPdqiJMqhuIY4aqC7--5CsMal',
+}
+
+
+def show_downloadeable_models():
+    print('\nAvailable .pt ReID models for automatic download')
+    print(list(__trained_urls.keys()))
+
+
+def get_model_url(model):
+    if model.name in __trained_urls:
+        return __trained_urls[model.name]
+    else:
+        None
+
+
+def is_model_in_model_types(model):
+    if model.name in __model_types:
+        return True
+    else:
+        return False
+
+
+def get_model_name(model):
+    for x in __model_types:
+        if x in model.name:
+            return x
+    return None
+
+
+def download_url(url, dst):
+    """Downloads file from a url to a destination.
+
+    Args:
+        url (str): url to download file.
+        dst (str): destination path.
+    """
+    from six.moves import urllib
+    print('* url="{}"'.format(url))
+    print('* destination="{}"'.format(dst))
+
+    def _reporthook(count, block_size, total_size):
+        global start_time
+        if count == 0:
+            start_time = time.time()
+            return
+        duration = time.time() - start_time
+        progress_size = int(count * block_size)
+        speed = int(progress_size / (1024*duration))
+        percent = int(count * block_size * 100 / total_size)
+        sys.stdout.write(
+            '\r...%d%%, %d MB, %d KB/s, %d seconds passed' %
+            (percent, progress_size / (1024*1024), speed, duration)
+        )
+        sys.stdout.flush()
+
+    urllib.request.urlretrieve(url, dst, _reporthook)
+    sys.stdout.write('\n')
+
+
+def load_pretrained_weights(model, weight_path):
+    r"""Loads pretrianed weights to model.
+
+    Features::
+        - Incompatible layers (unmatched in name or size) will be ignored.
+        - Can automatically deal with keys containing "module.".
+
+    Args:
+        model (nn.Module): network model.
+        weight_path (str): path to pretrained weights.
+
+    Examples::
+        >>> from torchreid.utils import load_pretrained_weights
+        >>> weight_path = 'log/my_model/model-best.pth.tar'
+        >>> load_pretrained_weights(model, weight_path)
+    """
+    checkpoint = torch.load(weight_path)
+    if 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+
+    model_dict = model.state_dict()
+    new_state_dict = OrderedDict()
+    matched_layers, discarded_layers = [], []
+
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            k = k[7:] # discard module.
+
+        if k in model_dict and model_dict[k].size() == v.size():
+            new_state_dict[k] = v
+            matched_layers.append(k)
+        else:
+            discarded_layers.append(k)
+
+    model_dict.update(new_state_dict)
+    model.load_state_dict(model_dict)
+
+    if len(matched_layers) == 0:
+        warnings.warn(
+            'The pretrained weights "{}" cannot be loaded, '
+            'please check the key names manually '
+            '(** ignored and continue **)'.format(weight_path)
+        )
+    else:
+        print(
+            'Successfully loaded pretrained weights from "{}"'.
+            format(weight_path)
+        )
+        if len(discarded_layers) > 0:
+            print(
+                '** The following layers are discarded '
+                'due to unmatched keys or layer size: {}'.
+                format(discarded_layers)
+            )
+

trackers/strongsort/reid_multibackend.py

@@ -0,0 +1,237 @@
+import torch.nn as nn
+import torch
+from pathlib import Path
+import numpy as np
+from itertools import islice
+import torchvision.transforms as transforms
+import cv2
+import sys
+import torchvision.transforms as T
+from collections import OrderedDict, namedtuple
+import gdown
+from os.path import exists as file_exists
+
+
+from ultralytics.yolo.utils.checks import check_requirements, check_version
+from ultralytics.yolo.utils import LOGGER
+from trackers.strongsort.deep.reid_model_factory import (show_downloadeable_models, get_model_url, get_model_name,
+                                                          download_url, load_pretrained_weights)
+from trackers.strongsort.deep.models import build_model
+
+
+def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''):
+    # Check file(s) for acceptable suffix
+    if file and suffix:
+        if isinstance(suffix, str):
+            suffix = [suffix]
+        for f in file if isinstance(file, (list, tuple)) else [file]:
+            s = Path(f).suffix.lower()  # file suffix
+            if len(s):
+                assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}"
+
+
+class ReIDDetectMultiBackend(nn.Module):
+    # ReID models MultiBackend class for python inference on various backends
+    def __init__(self, weights='osnet_x0_25_msmt17.pt', device=torch.device('cpu'), fp16=False):
+        super().__init__()
+
+        w = weights[0] if isinstance(weights, list) else weights
+        self.pt, self.jit, self.onnx, self.xml, self.engine, self.tflite = self.model_type(w)  # get backend
+        self.fp16 = fp16
+        self.fp16 &= self.pt or self.jit or self.engine  # FP16
+
+        # Build transform functions
+        self.device = device
+        self.image_size=(256, 128)
+        self.pixel_mean=[0.485, 0.456, 0.406]
+        self.pixel_std=[0.229, 0.224, 0.225]
+        self.transforms = []
+        self.transforms += [T.Resize(self.image_size)]
+        self.transforms += [T.ToTensor()]
+        self.transforms += [T.Normalize(mean=self.pixel_mean, std=self.pixel_std)]
+        self.preprocess = T.Compose(self.transforms)
+        self.to_pil = T.ToPILImage()
+
+        model_name = get_model_name(w)
+
+        if w.suffix == '.pt':
+            model_url = get_model_url(w)
+            if not file_exists(w) and model_url is not None:
+                gdown.download(model_url, str(w), quiet=False)
+            elif file_exists(w):
+                pass
+            else:
+                print(f'No URL associated to the chosen StrongSORT weights ({w}). Choose between:')
+                show_downloadeable_models()
+                exit()
+
+        # Build model
+        self.model = build_model(
+            model_name,
+            num_classes=1,
+            pretrained=not (w and w.is_file()),
+            use_gpu=device
+        )
+
+        if self.pt:  # PyTorch
+            # populate model arch with weights
+            if w and w.is_file() and w.suffix == '.pt':
+                load_pretrained_weights(self.model, w)
+                
+            self.model.to(device).eval()
+            self.model.half() if self.fp16 else  self.model.float()
+        elif self.jit:
+            LOGGER.info(f'Loading {w} for TorchScript inference...')
+            self.model = torch.jit.load(w)
+            self.model.half() if self.fp16 else self.model.float()
+        elif self.onnx:  # ONNX Runtime
+            LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
+            cuda = torch.cuda.is_available() and device.type != 'cpu'
+            #check_requirements(('onnx', 'onnxruntime-gpu' if cuda else 'onnxruntime'))
+            import onnxruntime
+            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
+            self.session = onnxruntime.InferenceSession(str(w), providers=providers)
+        elif self.engine:  # TensorRT
+            LOGGER.info(f'Loading {w} for TensorRT inference...')
+            import tensorrt as trt  # https://developer.nvidia.com/nvidia-tensorrt-download
+            check_version(trt.__version__, '7.0.0', hard=True)  # require tensorrt>=7.0.0
+            if device.type == 'cpu':
+                device = torch.device('cuda:0')
+            Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
+            logger = trt.Logger(trt.Logger.INFO)
+            with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
+                self.model_ = runtime.deserialize_cuda_engine(f.read())
+            self.context = self.model_.create_execution_context()
+            self.bindings = OrderedDict()
+            self.fp16 = False  # default updated below
+            dynamic = False
+            for index in range(self.model_.num_bindings):
+                name = self.model_.get_binding_name(index)
+                dtype = trt.nptype(self.model_.get_binding_dtype(index))
+                if self.model_.binding_is_input(index):
+                    if -1 in tuple(self.model_.get_binding_shape(index)):  # dynamic
+                        dynamic = True
+                        self.context.set_binding_shape(index, tuple(self.model_.get_profile_shape(0, index)[2]))
+                    if dtype == np.float16:
+                        self.fp16 = True
+                shape = tuple(self.context.get_binding_shape(index))
+                im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
+                self.bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
+            self.binding_addrs = OrderedDict((n, d.ptr) for n, d in self.bindings.items())
+            batch_size = self.bindings['images'].shape[0]  # if dynamic, this is instead max batch size
+        elif self.xml:  # OpenVINO
+            LOGGER.info(f'Loading {w} for OpenVINO inference...')
+            check_requirements(('openvino',))  # requires openvino-dev: https://pypi.org/project/openvino-dev/
+            from openvino.runtime import Core, Layout, get_batch
+            ie = Core()
+            if not Path(w).is_file():  # if not *.xml
+                w = next(Path(w).glob('*.xml'))  # get *.xml file from *_openvino_model dir
+            network = ie.read_model(model=w, weights=Path(w).with_suffix('.bin'))
+            if network.get_parameters()[0].get_layout().empty:
+                network.get_parameters()[0].set_layout(Layout("NCWH"))
+            batch_dim = get_batch(network)
+            if batch_dim.is_static:
+                batch_size = batch_dim.get_length()
+            self.executable_network = ie.compile_model(network, device_name="CPU")  # device_name="MYRIAD" for Intel NCS2
+            self.output_layer = next(iter(self.executable_network.outputs))
+        
+        elif self.tflite:
+            LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
+            try:  # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
+                from tflite_runtime.interpreter import Interpreter, load_delegate
+            except ImportError:
+                import tensorflow as tf
+                Interpreter, load_delegate = tf.lite.Interpreter, tf.lite.experimental.load_delegate,
+            self.interpreter = tf.lite.Interpreter(model_path=w)
+            self.interpreter.allocate_tensors()
+            # Get input and output tensors.
+            self.input_details = self.interpreter.get_input_details()
+            self.output_details = self.interpreter.get_output_details()
+            
+            # Test model on random input data.
+            input_data = np.array(np.random.random_sample((1,256,128,3)), dtype=np.float32)
+            self.interpreter.set_tensor(self.input_details[0]['index'], input_data)
+            
+            self.interpreter.invoke()
+
+            # The function `get_tensor()` returns a copy of the tensor data.
+            output_data = self.interpreter.get_tensor(self.output_details[0]['index'])
+        else:
+            print('This model framework is not supported yet!')
+            exit()
+        
+        
+    @staticmethod
+    def model_type(p='path/to/model.pt'):
+        # Return model type from model path, i.e. path='path/to/model.onnx' -> type=onnx
+        from trackers.reid_export import export_formats
+        sf = list(export_formats().Suffix)  # export suffixes
+        check_suffix(p, sf)  # checks
+        types = [s in Path(p).name for s in sf]
+        return types
+
+    def _preprocess(self, im_batch):
+
+        images = []
+        for element in im_batch:
+            image = self.to_pil(element)
+            image = self.preprocess(image)
+            images.append(image)
+
+        images = torch.stack(images, dim=0)
+        images = images.to(self.device)
+
+        return images
+    
+    
+    def forward(self, im_batch):
+        
+        # preprocess batch
+        im_batch = self._preprocess(im_batch)
+
+        # batch to half
+        if self.fp16 and im_batch.dtype != torch.float16:
+           im_batch = im_batch.half()
+
+        # batch processing
+        features = []
+        if self.pt:
+            features = self.model(im_batch)
+        elif self.jit:  # TorchScript
+            features = self.model(im_batch)
+        elif self.onnx:  # ONNX Runtime
+            im_batch = im_batch.cpu().numpy()  # torch to numpy
+            features = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im_batch})[0]
+        elif self.engine:  # TensorRT
+            if True and im_batch.shape != self.bindings['images'].shape:
+                i_in, i_out = (self.model_.get_binding_index(x) for x in ('images', 'output'))
+                self.context.set_binding_shape(i_in, im_batch.shape)  # reshape if dynamic
+                self.bindings['images'] = self.bindings['images']._replace(shape=im_batch.shape)
+                self.bindings['output'].data.resize_(tuple(self.context.get_binding_shape(i_out)))
+            s = self.bindings['images'].shape
+            assert im_batch.shape == s, f"input size {im_batch.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
+            self.binding_addrs['images'] = int(im_batch.data_ptr())
+            self.context.execute_v2(list(self.binding_addrs.values()))
+            features = self.bindings['output'].data
+        elif self.xml:  # OpenVINO
+            im_batch = im_batch.cpu().numpy()  # FP32
+            features = self.executable_network([im_batch])[self.output_layer]
+        else:
+            print('Framework not supported at the moment, we are working on it...')
+            exit()
+
+        if isinstance(features, (list, tuple)):
+            return self.from_numpy(features[0]) if len(features) == 1 else [self.from_numpy(x) for x in features]
+        else:
+            return self.from_numpy(features)
+
+    def from_numpy(self, x):
+        return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
+
+    def warmup(self, imgsz=[(256, 128, 3)]):
+        # Warmup model by running inference once
+        warmup_types = self.pt, self.jit, self.onnx, self.engine, self.tflite
+        if any(warmup_types) and self.device.type != 'cpu':
+            im = [np.empty(*imgsz).astype(np.uint8)]  # input
+            for _ in range(2 if self.jit else 1):  #
+                self.forward(im)  # warmup

trackers/strongsort/sort/detection.py

@@ -0,0 +1,58 @@
+# vim: expandtab:ts=4:sw=4
+import numpy as np
+
+
+class Detection(object):
+    """
+    This class represents a bounding box detection in a single image.
+
+    Parameters
+    ----------
+    tlwh : array_like
+        Bounding box in format `(x, y, w, h)`.
+    confidence : float
+        Detector confidence score.
+    feature : array_like
+        A feature vector that describes the object contained in this image.
+
+    Attributes
+    ----------
+    tlwh : ndarray
+        Bounding box in format `(top left x, top left y, width, height)`.
+    confidence : ndarray
+        Detector confidence score.
+    feature : ndarray | NoneType
+        A feature vector that describes the object contained in this image.
+
+    """
+
+    def __init__(self, tlwh, confidence, feature):
+        self.tlwh = np.asarray(tlwh, dtype=np.float32)
+        self.confidence = float(confidence)
+        self.feature = np.asarray(feature.cpu(), dtype=np.float32)
+
+    def to_tlbr(self):
+        """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
+        `(top left, bottom right)`.
+        """
+        ret = self.tlwh.copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    def to_xyah(self):
+        """Convert bounding box to format `(center x, center y, aspect ratio,
+        height)`, where the aspect ratio is `width / height`.
+        """
+        ret = self.tlwh.copy()
+        ret[:2] += ret[2:] / 2
+        ret[2] /= ret[3]
+        return ret
+    
+def to_xyah_ext(bbox):
+    """Convert bounding box to format `(center x, center y, aspect ratio,
+    height)`, where the aspect ratio is `width / height`.
+    """
+    ret = bbox.copy()
+    ret[:2] += ret[2:] / 2
+    ret[2] /= ret[3]
+    return ret

trackers/strongsort/sort/iou_matching.py

@@ -0,0 +1,82 @@
+# vim: expandtab:ts=4:sw=4
+from __future__ import absolute_import
+import numpy as np
+from . import linear_assignment
+
+
+def iou(bbox, candidates):
+    """Computer intersection over union.
+
+    Parameters
+    ----------
+    bbox : ndarray
+        A bounding box in format `(top left x, top left y, width, height)`.
+    candidates : ndarray
+        A matrix of candidate bounding boxes (one per row) in the same format
+        as `bbox`.
+
+    Returns
+    -------
+    ndarray
+        The intersection over union in [0, 1] between the `bbox` and each
+        candidate. A higher score means a larger fraction of the `bbox` is
+        occluded by the candidate.
+
+    """
+    bbox_tl, bbox_br = bbox[:2], bbox[:2] + bbox[2:]
+    candidates_tl = candidates[:, :2]
+    candidates_br = candidates[:, :2] + candidates[:, 2:]
+
+    tl = np.c_[np.maximum(bbox_tl[0], candidates_tl[:, 0])[:, np.newaxis],
+               np.maximum(bbox_tl[1], candidates_tl[:, 1])[:, np.newaxis]]
+    br = np.c_[np.minimum(bbox_br[0], candidates_br[:, 0])[:, np.newaxis],
+               np.minimum(bbox_br[1], candidates_br[:, 1])[:, np.newaxis]]
+    wh = np.maximum(0., br - tl)
+
+    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)
+
+
+def iou_cost(tracks, detections, track_indices=None,
+             detection_indices=None):
+    """An intersection over union distance metric.
+
+    Parameters
+    ----------
+    tracks : List[deep_sort.track.Track]
+        A list of tracks.
+    detections : List[deep_sort.detection.Detection]
+        A list of detections.
+    track_indices : Optional[List[int]]
+        A list of indices to tracks that should be matched. Defaults to
+        all `tracks`.
+    detection_indices : Optional[List[int]]
+        A list of indices to detections that should be matched. Defaults
+        to all `detections`.
+
+    Returns
+    -------
+    ndarray
+        Returns a cost matrix of shape
+        len(track_indices), len(detection_indices) where entry (i, j) is
+        `1 - iou(tracks[track_indices[i]], detections[detection_indices[j]])`.
+
+    """
+    if track_indices is None:
+        track_indices = np.arange(len(tracks))
+    if detection_indices is None:
+        detection_indices = np.arange(len(detections))
+
+    cost_matrix = np.zeros((len(track_indices), len(detection_indices)))
+    for row, track_idx in enumerate(track_indices):
+        if tracks[track_idx].time_since_update > 1:
+            cost_matrix[row, :] = linear_assignment.INFTY_COST
+            continue
+
+        bbox = tracks[track_idx].to_tlwh()
+        candidates = np.asarray(
+            [detections[i].tlwh for i in detection_indices])
+        cost_matrix[row, :] = 1. - iou(bbox, candidates)
+    return cost_matrix

trackers/strongsort/sort/kalman_filter.py

@@ -0,0 +1,214 @@
+# vim: expandtab:ts=4:sw=4
+import numpy as np
+import scipy.linalg
+"""
+Table for the 0.95 quantile of the chi-square distribution with N degrees of
+freedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv
+function and used as Mahalanobis gating threshold.
+"""
+chi2inv95 = {
+    1: 3.8415,
+    2: 5.9915,
+    3: 7.8147,
+    4: 9.4877,
+    5: 11.070,
+    6: 12.592,
+    7: 14.067,
+    8: 15.507,
+    9: 16.919}
+
+
+class KalmanFilter(object):
+    """
+    A simple Kalman filter for tracking bounding boxes in image space.
+    The 8-dimensional state space
+        x, y, a, h, vx, vy, va, vh
+    contains the bounding box center position (x, y), aspect ratio a, height h,
+    and their respective velocities.
+    Object motion follows a constant velocity model. The bounding box location
+    (x, y, a, h) is taken as direct observation of the state space (linear
+    observation model).
+    """
+
+    def __init__(self):
+        ndim, dt = 4, 1.
+
+        # Create Kalman filter model matrices.
+        self._motion_mat = np.eye(2 * ndim, 2 * ndim)
+        for i in range(ndim):
+            self._motion_mat[i, ndim + i] = dt
+
+        self._update_mat = np.eye(ndim, 2 * ndim)
+
+        # Motion and observation uncertainty are chosen relative to the current
+        # state estimate. These weights control the amount of uncertainty in
+        # the model. This is a bit hacky.
+        self._std_weight_position = 1. / 20
+        self._std_weight_velocity = 1. / 160
+
+    def initiate(self, measurement):
+        """Create track from unassociated measurement.
+        Parameters
+        ----------
+        measurement : ndarray
+            Bounding box coordinates (x, y, a, h) with center position (x, y),
+            aspect ratio a, and height h.
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the mean vector (8 dimensional) and covariance matrix (8x8
+            dimensional) of the new track. Unobserved velocities are initialized
+            to 0 mean.
+        """
+        mean_pos = measurement
+        mean_vel = np.zeros_like(mean_pos)
+        mean = np.r_[mean_pos, mean_vel]
+
+        std = [
+            2 * self._std_weight_position * measurement[0],   # the center point x
+            2 * self._std_weight_position * measurement[1],   # the center point y
+            1 * measurement[2],                               # the ratio of width/height
+            2 * self._std_weight_position * measurement[3],   # the height
+            10 * self._std_weight_velocity * measurement[0],
+            10 * self._std_weight_velocity * measurement[1],
+            0.1 * measurement[2],
+            10 * self._std_weight_velocity * measurement[3]]
+        covariance = np.diag(np.square(std))
+        return mean, covariance
+
+    def predict(self, mean, covariance):
+        """Run Kalman filter prediction step.
+        Parameters
+        ----------
+        mean : ndarray
+            The 8 dimensional mean vector of the object state at the previous
+            time step.
+        covariance : ndarray
+            The 8x8 dimensional covariance matrix of the object state at the
+            previous time step.
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the mean vector and covariance matrix of the predicted
+            state. Unobserved velocities are initialized to 0 mean.
+        """
+        std_pos = [
+            self._std_weight_position * mean[0],
+            self._std_weight_position * mean[1],
+            1 * mean[2],
+            self._std_weight_position * mean[3]]
+        std_vel = [
+            self._std_weight_velocity * mean[0],
+            self._std_weight_velocity * mean[1],
+            0.1 * mean[2],
+            self._std_weight_velocity * mean[3]]
+        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
+
+        mean = np.dot(self._motion_mat, mean)
+        covariance = np.linalg.multi_dot((
+            self._motion_mat, covariance, self._motion_mat.T)) + motion_cov
+
+        return mean, covariance
+
+    def project(self, mean, covariance, confidence=.0):
+        """Project state distribution to measurement space.
+        Parameters
+        ----------
+        mean : ndarray
+            The state's mean vector (8 dimensional array).
+        covariance : ndarray
+            The state's covariance matrix (8x8 dimensional).
+        confidence: (dyh) 检测框置信度
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the projected mean and covariance matrix of the given state
+            estimate.
+        """
+        std = [
+            self._std_weight_position * mean[3],
+            self._std_weight_position * mean[3],
+            1e-1,
+            self._std_weight_position * mean[3]]
+
+
+        std = [(1 - confidence) * x for x in std]
+
+        innovation_cov = np.diag(np.square(std))
+
+        mean = np.dot(self._update_mat, mean)
+        covariance = np.linalg.multi_dot((
+            self._update_mat, covariance, self._update_mat.T))
+        return mean, covariance + innovation_cov
+
+    def update(self, mean, covariance, measurement, confidence=.0):
+        """Run Kalman filter correction step.
+        Parameters
+        ----------
+        mean : ndarray
+            The predicted state's mean vector (8 dimensional).
+        covariance : ndarray
+            The state's covariance matrix (8x8 dimensional).
+        measurement : ndarray
+            The 4 dimensional measurement vector (x, y, a, h), where (x, y)
+            is the center position, a the aspect ratio, and h the height of the
+            bounding box.
+        confidence: (dyh)检测框置信度
+        Returns
+        -------
+        (ndarray, ndarray)
+            Returns the measurement-corrected state distribution.
+        """
+        projected_mean, projected_cov = self.project(mean, covariance, confidence)
+
+        chol_factor, lower = scipy.linalg.cho_factor(
+            projected_cov, lower=True, check_finite=False)
+        kalman_gain = scipy.linalg.cho_solve(
+            (chol_factor, lower), np.dot(covariance, self._update_mat.T).T,
+            check_finite=False).T
+        innovation = measurement - projected_mean
+
+        new_mean = mean + np.dot(innovation, kalman_gain.T)
+        new_covariance = covariance - np.linalg.multi_dot((
+            kalman_gain, projected_cov, kalman_gain.T))
+        return new_mean, new_covariance
+
+    def gating_distance(self, mean, covariance, measurements,
+                        only_position=False):
+        """Compute gating distance between state distribution and measurements.
+        A suitable distance threshold can be obtained from `chi2inv95`. If
+        `only_position` is False, the chi-square distribution has 4 degrees of
+        freedom, otherwise 2.
+        Parameters
+        ----------
+        mean : ndarray
+            Mean vector over the state distribution (8 dimensional).
+        covariance : ndarray
+            Covariance of the state distribution (8x8 dimensional).
+        measurements : ndarray
+            An Nx4 dimensional matrix of N measurements, each in
+            format (x, y, a, h) where (x, y) is the bounding box center
+            position, a the aspect ratio, and h the height.
+        only_position : Optional[bool]
+            If True, distance computation is done with respect to the bounding
+            box center position only.
+        Returns
+        -------
+        ndarray
+            Returns an array of length N, where the i-th element contains the
+            squared Mahalanobis distance between (mean, covariance) and
+            `measurements[i]`.
+        """
+        mean, covariance = self.project(mean, covariance)
+
+        if only_position:
+            mean, covariance = mean[:2], covariance[:2, :2]
+            measurements = measurements[:, :2]
+
+        cholesky_factor = np.linalg.cholesky(covariance)
+        d = measurements - mean
+        z = scipy.linalg.solve_triangular(
+            cholesky_factor, d.T, lower=True, check_finite=False,
+            overwrite_b=True)
+        squared_maha = np.sum(z * z, axis=0)
+        return squared_maha