add file to build StrongSORT

app.py

@@ -1,38 +1,65 @@
 import gradio as gr
-#import spaces
-#from huggingface_hub import hf_hub_download
-from detect import run
-
-def yolov9_inference(model_id, image_size, conf_threshold, iou_threshold, input_path = None):
-
-    # if img_path is not None:
-    #     # Load the model
-    #     # model_path = download_models(model_id)
-    #     model = load_model(model_id)
-    #     # Set model parameters
-    #     model.conf = conf_threshold
-    #     model.iou = iou_threshold
-    #     # Perform inference
-    #     results = model(img_path, size=image_size)
-    #     # Optionally, show detection bounding boxes on image
-    #     output = results.render()
-    #     return output[0]
-    # else:  
-    model = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, hide_conf=True, device='cpu')
-    return model, model 
+from detect_strongsort import run
+import os
+import threading
+
+should_continue = True
+
+
+def yolov9_inference(model_id, image_size, conf_threshold, iou_threshold, img_path=None, vid_path=None):
+    global should_continue
+    img_extensions = ['.jpg', '.jpeg', '.png', '.gif']  # Add more image extensions if needed
+    vid_extensions = ['.mp4', '.avi', '.mov', '.mkv']  # Add more video extensions if needed
+
+    input_path = None
+    if img_path is not None:
+        _, img_extension = os.path.splitext(img_path)
+        if img_extension.lower() in img_extensions:
+            input_path = img_path
+    elif vid_path is not None:
+        _, vid_extension = os.path.splitext(vid_path)
+        if vid_extension.lower() in vid_extensions:
+            input_path = vid_path
+
+    output_path = run(yolo_weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', strong_sort_weights = "osnet_x0_25_msmt17.pt", hide_conf= True)
+    # Assuming output_path is the path to the output file
+    _, output_extension = os.path.splitext(output_path)
+    if output_extension.lower() in img_extensions:
+        output_image = output_path  # Load the image file here
+        output_video = None
+    elif output_extension.lower() in vid_extensions:
+        output_image = None
+        output_video = output_path  # Load the video file here
+
+    return output_image, output_video, output_path
+
+def inference(model_id, image_size, conf_threshold, iou_threshold, img_path=None, vid_path=None):
+    global should_continue
+    should_continue = True
+    output_image, output_video, output_path = yolov9_inference(model_id, image_size, conf_threshold, iou_threshold, img_path, vid_path)
+    return output_image, output_video, output_path
+
+
+def stop_processing():
+    global should_continue
+    should_continue = False 
+    return "Stop..."
 
 def app():
     with gr.Blocks():
         with gr.Row():
             with gr.Column():
-                img_path = gr.Image(type="filepath", label="Image")
-                input_path = gr.Video(label="Input video")
+                gr.HTML("<h2>Input Parameters</h2>")
+                img_path = gr.File(label="Image")
+                vid_path = gr.File(label="Video")
                 model_id = gr.Dropdown(
                     label="Model",
                     choices=[
                         "last_best_model.pt",
+                        "best_model-converted.pt"
                     ],
-                    value="./last_best_model.pt",
+                    value="./last_best_model.pt"
+                    
                 )
                 image_size = gr.Slider(
                     label="Image Size",
@@ -56,21 +83,26 @@ def app():
                     value=0.5,
                 )
                 yolov9_infer = gr.Button(value="Inference")
-
+                stop_button = gr.Button(value="Stop")
             with gr.Column():
-                output = gr.Video(label="Output")
+                gr.HTML("<h2>Output</h2>")
+                output_image = gr.Image(type="numpy",label="Output Image")
+                output_video = gr.Video(label="Output Video")
                 output_path = gr.Textbox(label="Output path")
+
         yolov9_infer.click(
-            fn=yolov9_inference,
+            fn=inference,
             inputs=[
                 model_id,
                 image_size,
                 conf_threshold,
                 iou_threshold,
-                input_path
+                img_path,
+                vid_path
             ],
-            outputs=[output, output_path],
+            outputs=[output_image, output_video, output_path],
         )
+        stop_button.click(stop_processing)
 
 
 gradio_app = gr.Blocks()
@@ -81,8 +113,18 @@ with gradio_app:
     YOLOv9: Real-time Object Detection
     </h1>
     """)
+    css = """
+    body {
+        background-color: #f0f0f0;
+    }
+    h1 {
+        color: #4CAF50;
+    }
+    """
     with gr.Row():
         with gr.Column():
             app()
 
 gradio_app.launch(debug=True)
+
+

detect_strongsort.py

@@ -0,0 +1,392 @@
+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 platform
+import sys
+import numpy as np
+from pathlib import Path
+import torch
+import torch.backends.cudnn as cudnn
+from numpy import random
+from time import time
+
+
+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 / 'yolov9') not in sys.path:
+    sys.path.append(str(ROOT / 'yolov9'))  # add yolov5 ROOT to PATH
+if str(ROOT / 'strong_sort') not in sys.path:
+    sys.path.append(str(ROOT / 'strong_sort'))  # add strong_sort ROOT to PATH
+ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
+from models.experimental import attempt_load
+from models.common import DetectMultiBackend
+from utils.dataloaders import LoadImages, LoadStreams, LoadScreenshots
+from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
+                           increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
+from utils.torch_utils import select_device, time_sync, smart_inference_mode
+from utils.plots import Annotator, colors, save_one_box
+from strong_sort.utils.parser import get_config
+from strong_sort.strong_sort import StrongSORT
+
+
+VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv'  # include video suffixes
+
+
+def plot_one_box(x, img, color=None, label=None, line_thickness=3):
+    # Plots one bounding box on image img
+    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
+    color = color or [random.randint(0, 255) for _ in range(3)]
+    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
+    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
+    if label:
+        tf = max(tl - 1, 1)  # font thickness
+        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
+        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
+        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
+        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+
+
+
+@smart_inference_mode()
+def run(
+        source='0',
+        data = ROOT / 'data/coco.yaml',  # data.yaml path
+        yolo_weights=WEIGHTS / 'yolo.pt',  # model.pt path(s),
+        strong_sort_weights=WEIGHTS / 'osnet_x0_25_msmt17.pt',  # model.pt path,
+        config_strongsort=ROOT / 'strong_sort/configs/strong_sort.yaml',
+        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
+        view_img=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
+        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
+        name='exp',  # save results to project/name
+        exist_ok=False,  # existing project/name ok, do not increment
+        line_thickness=3,  # bounding box thickness (pixels)
+        hide_labels=False,  # hide labels
+        hide_conf=False,  # hide confidences
+        half=False,  # use FP16 half-precision inference
+        dnn=False,  # use OpenCV DNN for ONNX inference
+        vid_stride=1,  # video frame-rate stride
+):
+
+    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)
+    screenshot = source.lower().startswith('screen')
+
+    if is_url and is_file:
+        source = check_file(source)  # download
+
+    # Directories
+    if not isinstance(yolo_weights, list):  # single yolo model
+        exp_name = Path(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
+        yolo_weights = Path(yolo_weights[0])
+    else:  # multiple models after --yolo_weights
+        exp_name = 'ensemble'
+    exp_name = name if name else exp_name + "_" + Path(strong_sort_weights).stem
+    save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok)  # increment run
+    save_dir = Path(save_dir)
+    (save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
+
+    # Load model
+    device = select_device(device)
+    model = DetectMultiBackend(yolo_weights, device=device, dnn=dnn, data=data, fp16=half)
+    stride, names, pt = model.stride, model.names, model.pt
+    imgsz = check_img_size(imgsz, s=stride)  # check image size
+
+    # Dataloader
+    
+    # Dataloader
+    bs = 1  # batch_size
+    if webcam:
+        view_img = check_imshow(warn=True)
+        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
+        bs = len(dataset)
+    elif screenshot:
+        dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
+    else:
+        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
+    vid_path, vid_writer,txt_path = [None] * bs, [None] * bs, [None] * bs
+    
+    
+    # initialize StrongSORT
+    cfg = get_config()
+    cfg.merge_from_file(config_strongsort)
+
+    # Create as many strong sort instances as there are video sources
+    strongsort_list = []
+    for i in range(bs):
+        strongsort_list.append(
+            StrongSORT(
+                strong_sort_weights,
+                device,
+                half,
+                #max_dist=cfg.STRONGSORT.MAX_DIST,
+                max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE,
+                max_age=cfg.STRONGSORT.MAX_AGE,
+                n_init=cfg.STRONGSORT.N_INIT,
+                nn_budget=cfg.STRONGSORT.NN_BUDGET,
+                mc_lambda=cfg.STRONGSORT.MC_LAMBDA,
+                ema_alpha=cfg.STRONGSORT.EMA_ALPHA,
+
+            )
+        )
+        strongsort_list[i].model.warmup()
+    outputs = [None] * bs
+    
+    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
+
+    # Run tracking
+    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
+    seen, windows, dt,sdt = 0, [], (Profile(), Profile(), Profile(), Profile()),[0.0, 0.0, 0.0, 0.0]
+    curr_frames, prev_frames = [None] * bs, [None] * bs
+    for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
+        # s = ''
+        t1 = time_sync()
+        with dt[0]:
+            im = torch.from_numpy(im).to(model.device)
+            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
+            im /= 255  # 0 - 255 to 0.0 - 1.0
+            if len(im.shape) == 3:
+                im = im[None]  # expand for batch dim
+        t2 = time_sync()
+        sdt[0] += t2 - t1
+
+        # Inference
+        with dt[1]:
+            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
+            pred = model(im, augment=augment, visualize=visualize)
+            pred = pred[0][1]
+        t3 = time_sync()
+        sdt[1] += t3 - t2
+
+        # Apply NMS
+        with dt[2]:
+            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
+        sdt[2] += time_sync() - t3
+        
+        # Second-stage classifier (optional)
+        # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
+
+
+        # Process detections
+        for i, det in enumerate(pred):  # 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
+                txt_file_name = p.stem + f'_{i}' # Unique text file name
+                # save_path = str(save_dir / p.name) + str(i)  # im.jpg, vid.mp4, ...
+                save_path = str(save_dir / p.stem) + f'_{i}'  # Unique video file name
+
+            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 / p.name)  # im.jpg, vid.mp4, ...
+                # 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
+            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
+            imc = im0.copy() if save_crop else im0  # for save_crop
+            annotator = Annotator(im0, line_width=line_thickness, example=str(names))
+
+            if cfg.STRONGSORT.ECC:  # camera motion compensation
+                strongsort_list[i].tracker.camera_update(prev_frames[i], curr_frames[i])
+
+            if det is not None and len(det):
+                # Rescale boxes from img_size to im0 size
+                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
+
+                # Print results
+                for c in det[:, -1].unique():
+                    n = (det[:, -1] == c).sum()  # detections per class
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+
+                xywhs = xyxy2xywh(det[:, 0:4])
+                confs = det[:, 4]
+                clss = det[:, 5]
+
+                # pass detections to strongsort
+                t4 = time_sync()
+                outputs[i] = strongsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
+                t5 = time_sync()
+                sdt[3] += t5 - t4
+
+                # Write results
+                for j, (output, conf) in enumerate(zip(outputs[i], confs)):
+                    xyxy = output[0:4]
+                    id = output[4]
+                    cls = output[5]
+                # for *xyxy, conf, cls in reversed(det):
+                    if save_txt:  # Write to file
+                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
+                        # line = (id , cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
+                        line = ( int(p.stem), frame_idx, id , cls, *xywh, conf) if save_conf else ( p.stem, frame_idx, cls, *xywh)  # label format
+                        with open(txt_path + '.txt', 'a') as file:
+                            file.write(('%g ' * len(line) + '\n') % line)
+
+                    if save_img or save_crop or view_img:  # Add bbox to image
+                        c = int(cls)  # integer class
+                        label = None if hide_labels else (names[c] if hide_conf else f' { id } {names[c]} {conf:.2f}')
+                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=2)
+                    if save_crop:
+                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
+
+
+                # # draw boxes for visualization
+                # if len(outputs[i]) > 0:
+                #     for j, (output, conf) in enumerate(zip(outputs[i], confs)):
+    
+                #         bboxes = output[0:4]
+                #         id = output[4]
+                #         cls = output[5]
+
+                #         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]
+                #             # format video_name frame id xmin ymin width height score class 
+                #             with open(txt_path + '.txt', 'a') as file:
+                #                 file.write(f'{p.stem} {frame_idx} {id} {bbox_left} {bbox_top} {bbox_w} {bbox_h} {conf:.2f} {cls}\n')
+
+                #         if save_img or save_crop or view_img:  # Add bbox to image
+                #             c = int(cls)  # integer class
+                #             id = int(id)  # integer id
+                #             label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
+                #             plot_one_box(bboxes, im0, label=label, color=colors[int(cls)], line_thickness=2)
+                #             if save_crop:
+                #                 txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''
+                #                 save_one_box(bboxes, imc, file=save_dir / 'crops' / txt_file_name / names[c] / f'{id}' / f'{p.stem}.jpg', BGR=True)
+
+                print(f'{s}Done. YOLO:({t3 - t2:.3f}s), StrongSORT:({t5 - t4:.3f}s)')
+
+            else:
+                strongsort_list[i].increment_ages()
+                print('No detections')
+
+            # Stream results
+            im0 = annotator.result()
+            if view_img:
+                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)
+                cv2.waitKey(1)  # 1 millisecond
+
+            # Save results (image with detections)
+            if save_img:
+                if dataset.mode == 'image':
+                    cv2.imwrite(save_path, im0)
+                else:  # 'video' or 'stream'
+                    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
+                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc('m','p','4','v'), fps, (w, h))
+                    vid_writer[i].write(im0)
+
+            prev_frames[i] = curr_frames[i]
+
+    # Print time (inference-only)
+    LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
+    # Print results
+    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape, %.1fms StrongSORT' % tuple(1E3 * x / seen for x in sdt))
+    if save_txt or save_img:
+        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
+        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
+    if update:
+        strip_optimizer(yolo_weights[0])  # update model (to fix SourceChangeWarning)
+    return save_path
+def parse_opt():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--yolo-weights', nargs='+', type=str, default=WEIGHTS / 'yolov9.pt', help='model.pt path(s)')
+    parser.add_argument('--strong-sort-weights', type=str, default=WEIGHTS / 'osnet_x0_25_msmt17.pt')
+    parser.add_argument('--config-strongsort', type=str, default='strong_sort/configs/strong_sort.yaml')
+    parser.add_argument('--source', type=str, default='0', help='file/dir/URL/glob, 0 for webcam')  
+    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')
+    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('--view-img', action='store_true', help='show 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('--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 / 'runs/track', help='save results to project/name')
+    parser.add_argument('--name', default='exp', help='save results to project/name')
+    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument('--line-thickness', default=3, 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('--half', action='store_true', help='use FP16 half-precision inference')
+    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
+    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
+    opt = parser.parse_args()
+    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
+
+    return opt
+
+
+def main(opt):
+    # check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    run(**vars(opt))
+
+
+if __name__ == "__main__":
+    opt = parse_opt()
+    main(opt)

strong_sort/.gitignore

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

strong_sort/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Ziqiang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

strong_sort/README.md

@@ -0,0 +1,137 @@
+# Deep Sort with PyTorch
+
+![](demo/demo.gif)
+
+## Update(1-1-2020)
+Changes
+- fix bugs
+- refactor code
+- accerate detection by adding nms on gpu
+
+## Latest Update(07-22)
+Changes
+- bug fix (Thanks @JieChen91 and @yingsen1 for bug reporting).  
+- using batch for feature extracting for each frame, which lead to a small speed up.  
+- code improvement.
+
+Futher improvement direction  
+- Train detector on specific dataset rather than the official one.
+- Retrain REID model on pedestrain dataset for better performance.
+- Replace YOLOv3 detector with advanced ones.
+
+**Any contributions to this repository is welcome!**
+
+
+## Introduction
+This is an implement of MOT tracking algorithm deep sort. Deep sort is basicly the same with sort but added a CNN model to extract features in image of human part bounded by a detector. This CNN model is indeed a RE-ID model and the detector used in [PAPER](https://arxiv.org/abs/1703.07402) is FasterRCNN , and the original source code is [HERE](https://github.com/nwojke/deep_sort).  
+However in original code, the CNN model is implemented with tensorflow, which I'm not familier with. SO I re-implemented the CNN feature extraction model with PyTorch, and changed the CNN model a little bit. Also, I use **YOLOv3** to generate bboxes instead of FasterRCNN.
+
+## Dependencies
+- python 3 (python2 not sure)
+- numpy
+- scipy
+- opencv-python
+- sklearn
+- torch >= 0.4
+- torchvision >= 0.1
+- pillow
+- vizer
+- edict
+
+## Quick Start
+0. Check all dependencies installed
+```bash
+pip install -r requirements.txt
+```
+for user in china, you can specify pypi source to accelerate install like:
+```bash
+pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+```
+
+1. Clone this repository
+```
+git clone git@github.com:ZQPei/deep_sort_pytorch.git
+```
+
+2. Download YOLOv3 parameters
+```
+cd detector/YOLOv3/weight/
+wget https://pjreddie.com/media/files/yolov3.weights
+wget https://pjreddie.com/media/files/yolov3-tiny.weights
+cd ../../../
+```
+
+3. Download deepsort parameters ckpt.t7
+```
+cd deep_sort/deep/checkpoint
+# download ckpt.t7 from
+https://drive.google.com/drive/folders/1xhG0kRH1EX5B9_Iz8gQJb7UNnn_riXi6 to this folder
+cd ../../../
+```  
+
+4. Compile nms module
+```bash
+cd detector/YOLOv3/nms
+sh build.sh
+cd ../../..
+```
+
+Notice:
+If compiling failed, the simplist way is to **Upgrade your pytorch >= 1.1 and torchvision >= 0.3" and you can avoid the troublesome compiling problems which are most likely caused by either `gcc version too low` or `libraries missing`.
+
+5. Run demo
+```
+usage: python yolov3_deepsort.py VIDEO_PATH
+                                [--help]
+                                [--frame_interval FRAME_INTERVAL]
+                                [--config_detection CONFIG_DETECTION]
+                                [--config_deepsort CONFIG_DEEPSORT]
+                                [--display]
+                                [--display_width DISPLAY_WIDTH]
+                                [--display_height DISPLAY_HEIGHT]
+                                [--save_path SAVE_PATH]          
+                                [--cpu]          
+
+# yolov3 + deepsort
+python yolov3_deepsort.py [VIDEO_PATH]
+
+# yolov3_tiny + deepsort
+python yolov3_deepsort.py [VIDEO_PATH] --config_detection ./configs/yolov3_tiny.yaml
+
+# yolov3 + deepsort on webcam
+python3 yolov3_deepsort.py /dev/video0 --camera 0
+
+# yolov3_tiny + deepsort on webcam
+python3 yolov3_deepsort.py /dev/video0 --config_detection ./configs/yolov3_tiny.yaml --camera 0
+```
+Use `--display` to enable display.  
+Results will be saved to `./output/results.avi` and `./output/results.txt`.
+
+All files above can also be accessed from BaiduDisk!  
+linker：[BaiduDisk](https://pan.baidu.com/s/1YJ1iPpdFTlUyLFoonYvozg)
+passwd：fbuw
+
+## Training the RE-ID model
+The original model used in paper is in original_model.py, and its parameter here [original_ckpt.t7](https://drive.google.com/drive/folders/1xhG0kRH1EX5B9_Iz8gQJb7UNnn_riXi6).  
+
+To train the model, first you need download [Market1501](http://www.liangzheng.com.cn/Project/project_reid.html) dataset or [Mars](http://www.liangzheng.com.cn/Project/project_mars.html) dataset.  
+
+Then you can try [train.py](deep_sort/deep/train.py) to train your own parameter and evaluate it using [test.py](deep_sort/deep/test.py) and [evaluate.py](deep_sort/deep/evalute.py).
+![train.jpg](deep_sort/deep/train.jpg)
+
+## Demo videos and images
+[demo.avi](https://drive.google.com/drive/folders/1xhG0kRH1EX5B9_Iz8gQJb7UNnn_riXi6)
+[demo2.avi](https://drive.google.com/drive/folders/1xhG0kRH1EX5B9_Iz8gQJb7UNnn_riXi6)
+
+![1.jpg](demo/1.jpg)
+![2.jpg](demo/2.jpg)
+
+
+## References
+- paper: [Simple Online and Realtime Tracking with a Deep Association Metric](https://arxiv.org/abs/1703.07402)
+
+- code: [nwojke/deep_sort](https://github.com/nwojke/deep_sort)
+
+- paper: [YOLOv3](https://pjreddie.com/media/files/papers/YOLOv3.pdf)
+
+- code: [Joseph Redmon/yolov3](https://pjreddie.com/darknet/yolo/)

strong_sort/__init__.py

@@ -0,0 +1,11 @@
+from .strong_sort import StrongSORT
+
+
+__all__ = ['StrongSORT', 'build_tracker']
+
+
+def build_tracker(cfg, use_cuda):
+    return StrongSORT(cfg.STRONGSORT.REID_CKPT, 
+                max_dist=cfg.STRONGSORT.MAX_DIST, min_confidence=cfg.STRONGSORT.MIN_CONFIDENCE, 
+                nms_max_overlap=cfg.STRONGSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE, 
+                max_age=cfg.STRONGSORT.MAX_AGE, n_init=cfg.STRONGSORT.N_INIT, nn_budget=cfg.STRONGSORT.NN_BUDGET, use_cuda=use_cuda)

strong_sort/configs/strong_sort.yaml

@@ -0,0 +1,10 @@
+STRONGSORT:
+  ECC: True              # activate camera motion compensation
+  MC_LAMBDA: 0.995       # matching with both appearance (1 - MC_LAMBDA) and motion cost
+  EMA_ALPHA: 0.9         # updates  appearance  state in  an exponential moving average manner
+  MAX_DIST: 0.2          # The matching threshold. Samples with larger distance are considered an invalid match
+  MAX_IOU_DISTANCE: 0.7  # Gating threshold. Associations with cost larger than this value are disregarded.
+  MAX_AGE: 30            # Maximum number of missed misses before a track is deleted
+  N_INIT: 1              # Number of frames that a track remains in initialization phase
+  NN_BUDGET: 100         # Maximum size of the appearance descriptors gallery
+  

strong_sort/deep/reid/.flake8

@@ -0,0 +1,18 @@
+[flake8]
+ignore =
+    # At least two spaces before inline comment
+    E261,
+    # Line lengths are recommended to be no greater than 79 characters
+    E501,
+    # Missing whitespace around arithmetic operator 
+    E226,
+    # Blank line contains whitespace
+    W293,
+    # Do not use bare 'except'
+    E722,
+    # Line break after binary operator
+    W504,
+    # isort found an import in the wrong position
+    I001
+max-line-length = 79
+exclude = __init__.py, build, torchreid/metrics/rank_cylib/

strong_sort/deep/reid/.gitignore

@@ -0,0 +1,140 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Cython eval code
+*.c
+*.html
+
+# OS X
+.DS_Store
+.Spotlight-V100
+.Trashes
+._*
+
+# ReID
+reid-data/
+log/
+saved-models/
+model-zoo/
+debug*

strong_sort/deep/reid/.isort.cfg

@@ -0,0 +1,10 @@
+[isort]
+line_length=79
+multi_line_output=3
+length_sort=true
+known_standard_library=numpy,setuptools
+known_myself=torchreid
+known_third_party=matplotlib,cv2,torch,torchvision,PIL,yacs
+no_lines_before=STDLIB,THIRDPARTY
+sections=FUTURE,STDLIB,THIRDPARTY,myself,FIRSTPARTY,LOCALFOLDER
+default_section=FIRSTPARTY

strong_sort/deep/reid/.style.yapf

@@ -0,0 +1,7 @@
+[style]
+BASED_ON_STYLE = pep8
+BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true
+SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true
+DEDENT_CLOSING_BRACKETS = true
+SPACES_BEFORE_COMMENT = 1
+ARITHMETIC_PRECEDENCE_INDICATION = true

strong_sort/deep/reid/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Kaiyang Zhou
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

strong_sort/deep/reid/README.rst

@@ -0,0 +1,317 @@
+Torchreid
+===========
+Torchreid is a library for deep-learning person re-identification, written in `PyTorch <https://pytorch.org/>`_ and developed for our ICCV'19 project, `Omni-Scale Feature Learning for Person Re-Identification <https://arxiv.org/abs/1905.00953>`_.
+
+It features:
+
+- multi-GPU training
+- support both image- and video-reid
+- end-to-end training and evaluation
+- incredibly easy preparation of reid datasets
+- multi-dataset training
+- cross-dataset evaluation
+- standard protocol used by most research papers
+- highly extensible (easy to add models, datasets, training methods, etc.)
+- implementations of state-of-the-art deep reid models
+- access to pretrained reid models
+- advanced training techniques
+- visualization tools (tensorboard, ranks, etc.)
+
+
+Code: https://github.com/KaiyangZhou/deep-person-reid.
+
+Documentation: https://kaiyangzhou.github.io/deep-person-reid/.
+
+How-to instructions: https://kaiyangzhou.github.io/deep-person-reid/user_guide.
+
+Model zoo: https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.
+
+Tech report: https://arxiv.org/abs/1910.10093.
+
+You can find some research projects that are built on top of Torchreid `here <https://github.com/KaiyangZhou/deep-person-reid/tree/master/projects>`_.
+
+
+What's new
+------------
+- [Aug 2021] We have released the ImageNet-pretrained models of ``osnet_ain_x0_75``, ``osnet_ain_x0_5`` and ``osnet_ain_x0_25``. The pretraining setup follows `pycls <https://github.com/facebookresearch/pycls/blob/master/configs/archive/imagenet/resnet/R-50-1x64d_step_8gpu.yaml>`_.
+- [Apr 2021] We have updated the appendix in the `TPAMI version of OSNet <https://arxiv.org/abs/1910.06827v5>`_ to include results in the multi-source domain generalization setting. The trained models can be found in the `Model Zoo <https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.html>`_.
+- [Apr 2021] We have added a script to automate the process of calculating average results over multiple splits. For more details please see ``tools/parse_test_res.py``.
+- [Apr 2021] ``v1.4.0``: We added the person search dataset, `CUHK-SYSU <http://www.ee.cuhk.edu.hk/~xgwang/PS/dataset.html>`_.  Please see the `documentation <https://kaiyangzhou.github.io/deep-person-reid/>`_ regarding how to download the dataset (it contains cropped person images).
+- [Apr 2021] All models in the model zoo have been moved to google drive. Please raise an issue if any model's performance is inconsistent with the numbers shown in the model zoo page (could be caused by wrong links).
+- [Mar 2021] `OSNet <https://arxiv.org/abs/1910.06827>`_ will appear in the TPAMI journal! Compared with the conference version, which focuses on discriminative feature learning using the omni-scale building block, this journal extension further considers generalizable feature learning by integrating `instance normalization layers <https://arxiv.org/abs/1607.08022>`_ with the OSNet architecture. We hope this journal paper can motivate more future work to taclke the generalization issue in cross-dataset re-ID.
+- [Mar 2021] Generalization across domains (datasets) in person re-ID is crucial in real-world applications, which is closely related to the topic of *domain generalization*. Interested in learning how the field of domain generalization has developed over the last decade? Check our recent survey in this topic at https://arxiv.org/abs/2103.02503, with coverage on the history, datasets, related problems, methodologies, potential directions, and so on (*methods designed for generalizable re-ID are also covered*!).
+- [Feb 2021] ``v1.3.6`` Added `University-1652 <https://dl.acm.org/doi/abs/10.1145/3394171.3413896>`_, a new dataset for multi-view multi-source geo-localization (credit to `Zhedong Zheng <https://github.com/layumi>`_).
+- [Feb 2021] ``v1.3.5``: Now the `cython code <https://github.com/KaiyangZhou/deep-person-reid/pull/412>`_ works on Windows (credit to `lablabla <https://github.com/lablabla>`_).
+- [Jan 2021] Our recent work, `MixStyle <https://openreview.net/forum?id=6xHJ37MVxxp>`_ (mixing instance-level feature statistics of samples of different domains for improving domain generalization), has been accepted to ICLR'21. The code has been released at https://github.com/KaiyangZhou/mixstyle-release where the person re-ID part is based on Torchreid.
+- [Jan 2021] A new evaluation metric called `mean Inverse Negative Penalty (mINP)` for person re-ID has been introduced in `Deep Learning for Person Re-identification: A Survey and Outlook (TPAMI 2021) <https://arxiv.org/abs/2001.04193>`_. Their code can be accessed at `<https://github.com/mangye16/ReID-Survey>`_.
+- [Aug 2020] ``v1.3.3``: Fixed bug in ``visrank`` (caused by not unpacking ``dsetid``).
+- [Aug 2020] ``v1.3.2``: Added ``_junk_pids`` to ``grid`` and ``prid``. This avoids using mislabeled gallery images for training when setting ``combineall=True``.
+- [Aug 2020] ``v1.3.0``: (1) Added ``dsetid`` to the existing 3-tuple data source, resulting in ``(impath, pid, camid, dsetid)``. This variable denotes the dataset ID and is useful when combining multiple datasets for training (as a dataset indicator). E.g., when combining ``market1501`` and ``cuhk03``, the former will be assigned ``dsetid=0`` while the latter will be assigned ``dsetid=1``. (2) Added ``RandomDatasetSampler``. Analogous to ``RandomDomainSampler``, ``RandomDatasetSampler`` samples a certain number of images (``batch_size // num_datasets``) from each of specified datasets (the amount is determined by ``num_datasets``).
+- [Aug 2020] ``v1.2.6``: Added ``RandomDomainSampler`` (it samples ``num_cams`` cameras each with ``batch_size // num_cams`` images to form a mini-batch).
+- [Jun 2020] ``v1.2.5``: (1) Dataloader's output from ``__getitem__`` has been changed from ``list`` to ``dict``. Previously, an element, e.g. image tensor, was fetched with ``imgs=data[0]``. Now it should be obtained by ``imgs=data['img']``. See this `commit <https://github.com/KaiyangZhou/deep-person-reid/commit/aefe335d68f39a20160860e6d14c2d34f539b8a5>`_ for detailed changes. (2) Added ``k_tfm`` as an option to image data loader, which allows data augmentation to be applied ``k_tfm`` times *independently* to an image. If ``k_tfm > 1``, ``imgs=data['img']`` returns a list with ``k_tfm`` image tensors.
+- [May 2020] Added the person attribute recognition code used in `Omni-Scale Feature Learning for Person Re-Identification (ICCV'19) <https://arxiv.org/abs/1905.00953>`_. See ``projects/attribute_recognition/``.
+- [May 2020] ``v1.2.1``: Added a simple API for feature extraction (``torchreid/utils/feature_extractor.py``). See the `documentation <https://kaiyangzhou.github.io/deep-person-reid/user_guide.html>`_ for the instruction.
+- [Apr 2020] Code for reproducing the experiments of `deep mutual learning <https://zpascal.net/cvpr2018/Zhang_Deep_Mutual_Learning_CVPR_2018_paper.pdf>`_ in the `OSNet paper <https://arxiv.org/pdf/1905.00953v6.pdf>`__ (Supp. B) has been released at ``projects/DML``.
+- [Apr 2020] Upgraded to ``v1.2.0``. The engine class has been made more model-agnostic to improve extensibility. See `Engine <torchreid/engine/engine.py>`_ and `ImageSoftmaxEngine <torchreid/engine/image/softmax.py>`_ for more details. Credit to `Dassl.pytorch <https://github.com/KaiyangZhou/Dassl.pytorch>`_.
+- [Dec 2019] Our `OSNet paper <https://arxiv.org/pdf/1905.00953v6.pdf>`_ has been updated, with additional experiments (in section B of the supplementary) showing some useful techniques for improving OSNet's performance in practice.
+- [Nov 2019] ``ImageDataManager`` can load training data from target datasets by setting ``load_train_targets=True``, and the train-loader can be accessed with ``train_loader_t = datamanager.train_loader_t``. This feature is useful for domain adaptation research.
+
+
+Installation
+---------------
+
+Make sure `conda <https://www.anaconda.com/distribution/>`_ is installed.
+
+
+.. code-block:: bash
+
+    # cd to your preferred directory and clone this repo
+    git clone https://github.com/KaiyangZhou/deep-person-reid.git
+
+    # create environment
+    cd deep-person-reid/
+    conda create --name torchreid python=3.7
+    conda activate torchreid
+
+    # install dependencies
+    # make sure `which python` and `which pip` point to the correct path
+    pip install -r requirements.txt
+
+    # install torch and torchvision (select the proper cuda version to suit your machine)
+    conda install pytorch torchvision cudatoolkit=9.0 -c pytorch
+
+    # install torchreid (don't need to re-build it if you modify the source code)
+    python setup.py develop
+
+
+Get started: 30 seconds to Torchreid
+-------------------------------------
+1. Import ``torchreid``
+
+.. code-block:: python
+    
+    import torchreid
+
+2. Load data manager
+
+.. code-block:: python
+    
+    datamanager = torchreid.data.ImageDataManager(
+        root="reid-data",
+        sources="market1501",
+        targets="market1501",
+        height=256,
+        width=128,
+        batch_size_train=32,
+        batch_size_test=100,
+        transforms=["random_flip", "random_crop"]
+    )
+
+3 Build model, optimizer and lr_scheduler
+
+.. code-block:: python
+    
+    model = torchreid.models.build_model(
+        name="resnet50",
+        num_classes=datamanager.num_train_pids,
+        loss="softmax",
+        pretrained=True
+    )
+
+    model = model.cuda()
+
+    optimizer = torchreid.optim.build_optimizer(
+        model,
+        optim="adam",
+        lr=0.0003
+    )
+
+    scheduler = torchreid.optim.build_lr_scheduler(
+        optimizer,
+        lr_scheduler="single_step",
+        stepsize=20
+    )
+
+4. Build engine
+
+.. code-block:: python
+    
+    engine = torchreid.engine.ImageSoftmaxEngine(
+        datamanager,
+        model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        label_smooth=True
+    )
+
+5. Run training and test
+
+.. code-block:: python
+    
+    engine.run(
+        save_dir="log/resnet50",
+        max_epoch=60,
+        eval_freq=10,
+        print_freq=10,
+        test_only=False
+    )
+
+
+A unified interface
+-----------------------
+In "deep-person-reid/scripts/", we provide a unified interface to train and test a model. See "scripts/main.py" and "scripts/default_config.py" for more details. The folder "configs/" contains some predefined configs which you can use as a starting point.
+
+Below we provide an example to train and test `OSNet (Zhou et al. ICCV'19) <https://arxiv.org/abs/1905.00953>`_. Assume :code:`PATH_TO_DATA` is the directory containing reid datasets. The environmental variable :code:`CUDA_VISIBLE_DEVICES` is omitted, which you need to specify if you have a pool of gpus and want to use a specific set of them.
+
+Conventional setting
+^^^^^^^^^^^^^^^^^^^^^
+
+To train OSNet on Market1501, do
+
+.. code-block:: bash
+
+    python scripts/main.py \
+    --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \
+    --transforms random_flip random_erase \
+    --root $PATH_TO_DATA
+
+
+The config file sets Market1501 as the default dataset. If you wanna use DukeMTMC-reID, do
+
+.. code-block:: bash
+
+    python scripts/main.py \
+    --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \
+    -s dukemtmcreid \
+    -t dukemtmcreid \
+    --transforms random_flip random_erase \
+    --root $PATH_TO_DATA \
+    data.save_dir log/osnet_x1_0_dukemtmcreid_softmax_cosinelr
+
+The code will automatically (download and) load the ImageNet pretrained weights. After the training is done, the model will be saved as "log/osnet_x1_0_market1501_softmax_cosinelr/model.pth.tar-250". Under the same folder, you can find the `tensorboard <https://pytorch.org/docs/stable/tensorboard.html>`_ file. To visualize the learning curves using tensorboard, you can run :code:`tensorboard --logdir=log/osnet_x1_0_market1501_softmax_cosinelr` in the terminal and visit :code:`http://localhost:6006/` in your web browser.
+
+Evaluation is automatically performed at the end of training. To run the test again using the trained model, do
+
+.. code-block:: bash
+
+    python scripts/main.py \
+    --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \
+    --root $PATH_TO_DATA \
+    model.load_weights log/osnet_x1_0_market1501_softmax_cosinelr/model.pth.tar-250 \
+    test.evaluate True
+
+
+Cross-domain setting
+^^^^^^^^^^^^^^^^^^^^^
+
+Suppose you wanna train OSNet on DukeMTMC-reID and test its performance on Market1501, you can do
+
+.. code-block:: bash
+
+    python scripts/main.py \
+    --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad.yaml \
+    -s dukemtmcreid \
+    -t market1501 \
+    --transforms random_flip color_jitter \
+    --root $PATH_TO_DATA
+
+Here we only test the cross-domain performance. However, if you also want to test the performance on the source dataset, i.e. DukeMTMC-reID, you can set :code:`-t dukemtmcreid market1501`, which will evaluate the model on the two datasets separately.
+
+Different from the same-domain setting, here we replace :code:`random_erase` with :code:`color_jitter`. This can improve the generalization performance on the unseen target dataset.
+
+Pretrained models are available in the `Model Zoo <https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.html>`_.
+
+
+Datasets
+--------
+
+Image-reid datasets
+^^^^^^^^^^^^^^^^^^^^^
+- `Market1501 <https://www.cv-foundation.org/openaccess/content_iccv_2015/papers/Zheng_Scalable_Person_Re-Identification_ICCV_2015_paper.pdf>`_
+- `CUHK03 <https://www.cv-foundation.org/openaccess/content_cvpr_2014/papers/Li_DeepReID_Deep_Filter_2014_CVPR_paper.pdf>`_
+- `DukeMTMC-reID <https://arxiv.org/abs/1701.07717>`_
+- `MSMT17 <https://arxiv.org/abs/1711.08565>`_
+- `VIPeR <http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.331.7285&rep=rep1&type=pdf>`_
+- `GRID <http://www.eecs.qmul.ac.uk/~txiang/publications/LoyXiangGong_cvpr_2009.pdf>`_
+- `CUHK01 <http://www.ee.cuhk.edu.hk/~xgwang/papers/liZWaccv12.pdf>`_
+- `SenseReID <http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhao_Spindle_Net_Person_CVPR_2017_paper.pdf>`_
+- `QMUL-iLIDS <http://www.eecs.qmul.ac.uk/~sgg/papers/ZhengGongXiang_BMVC09.pdf>`_
+- `PRID <https://pdfs.semanticscholar.org/4c1b/f0592be3e535faf256c95e27982db9b3d3d3.pdf>`_
+
+Geo-localization datasets
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- `University-1652 <https://dl.acm.org/doi/abs/10.1145/3394171.3413896>`_
+
+Video-reid datasets
+^^^^^^^^^^^^^^^^^^^^^^^
+- `MARS <http://www.liangzheng.org/1320.pdf>`_
+- `iLIDS-VID <https://www.eecs.qmul.ac.uk/~sgg/papers/WangEtAl_ECCV14.pdf>`_
+- `PRID2011 <https://pdfs.semanticscholar.org/4c1b/f0592be3e535faf256c95e27982db9b3d3d3.pdf>`_
+- `DukeMTMC-VideoReID <http://openaccess.thecvf.com/content_cvpr_2018/papers/Wu_Exploit_the_Unknown_CVPR_2018_paper.pdf>`_
+
+
+Models
+-------
+
+ImageNet classification models
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- `ResNet <https://arxiv.org/abs/1512.03385>`_
+- `ResNeXt <https://arxiv.org/abs/1611.05431>`_
+- `SENet <https://arxiv.org/abs/1709.01507>`_
+- `DenseNet <https://arxiv.org/abs/1608.06993>`_
+- `Inception-ResNet-V2 <https://arxiv.org/abs/1602.07261>`_
+- `Inception-V4 <https://arxiv.org/abs/1602.07261>`_
+- `Xception <https://arxiv.org/abs/1610.02357>`_
+- `IBN-Net <https://arxiv.org/abs/1807.09441>`_
+
+Lightweight models
+^^^^^^^^^^^^^^^^^^^
+- `NASNet <https://arxiv.org/abs/1707.07012>`_
+- `MobileNetV2 <https://arxiv.org/abs/1801.04381>`_
+- `ShuffleNet <https://arxiv.org/abs/1707.01083>`_
+- `ShuffleNetV2 <https://arxiv.org/abs/1807.11164>`_
+- `SqueezeNet <https://arxiv.org/abs/1602.07360>`_
+
+ReID-specific models
+^^^^^^^^^^^^^^^^^^^^^^
+- `MuDeep <https://arxiv.org/abs/1709.05165>`_
+- `ResNet-mid <https://arxiv.org/abs/1711.08106>`_
+- `HACNN <https://arxiv.org/abs/1802.08122>`_
+- `PCB <https://arxiv.org/abs/1711.09349>`_
+- `MLFN <https://arxiv.org/abs/1803.09132>`_
+- `OSNet <https://arxiv.org/abs/1905.00953>`_
+- `OSNet-AIN <https://arxiv.org/abs/1910.06827>`_
+
+
+Useful links
+-------------
+- `OSNet-IBN1-Lite (test-only code with lite docker container) <https://github.com/RodMech/OSNet-IBN1-Lite>`_
+- `Deep Learning for Person Re-identification: A Survey and Outlook <https://github.com/mangye16/ReID-Survey>`_
+
+
+Citation
+---------
+If you use this code or the models in your research, please give credit to the following papers:
+
+.. code-block:: bash
+
+    @article{torchreid,
+      title={Torchreid: A Library for Deep Learning Person Re-Identification in Pytorch},
+      author={Zhou, Kaiyang and Xiang, Tao},
+      journal={arXiv preprint arXiv:1910.10093},
+      year={2019}
+    }
+    
+    @inproceedings{zhou2019osnet,
+      title={Omni-Scale Feature Learning for Person Re-Identification},
+      author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
+      booktitle={ICCV},
+      year={2019}
+    }
+
+    @article{zhou2021osnet,
+      title={Learning Generalisable Omni-Scale Representations for Person Re-Identification},
+      author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
+      journal={TPAMI},
+      year={2021}
+    }

strong_sort/deep/reid/configs/im_osnet_ain_x1_0_softmax_256x128_amsgrad_cosine.yaml

@@ -0,0 +1,35 @@
+model:
+  name: 'osnet_ain_x1_0'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501', 'dukemtmcreid']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip', 'color_jitter']
+  save_dir: 'log/osnet_ain_x1_0_market1501_softmax_cosinelr'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0015
+  max_epoch: 100
+  batch_size: 64
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'cosine'
+
+test:
+  batch_size: 300
+  dist_metric: 'cosine'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_osnet_ibn_x1_0_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'osnet_ibn_x1_0'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['dukemtmcreid']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip', 'color_jitter']
+  save_dir: 'log/osnet_ibn_x1_0_market2duke_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0015
+  max_epoch: 150
+  batch_size: 64
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [60]
+
+test:
+  batch_size: 300
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_osnet_x0_25_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'osnet_x0_25'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/osnet_x0_25_market1501_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.003
+  max_epoch: 180
+  batch_size: 128
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [80]
+
+test:
+  batch_size: 300
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_osnet_x0_5_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'osnet_x0_5'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/osnet_x0_5_market1501_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.003
+  max_epoch: 180
+  batch_size: 128
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [80]
+
+test:
+  batch_size: 300
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_osnet_x0_75_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'osnet_x0_75'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/osnet_x0_75_market1501_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0015
+  max_epoch: 150
+  batch_size: 64
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [60]
+
+test:
+  batch_size: 300
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_osnet_x1_0_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'osnet_x1_0'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/osnet_x1_0_market1501_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0015
+  max_epoch: 150
+  batch_size: 64
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [60]
+
+test:
+  batch_size: 300
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml

@@ -0,0 +1,35 @@
+model:
+  name: 'osnet_x1_0'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/osnet_x1_0_market1501_softmax_cosinelr'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0015
+  max_epoch: 250
+  batch_size: 64
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'cosine'
+
+test:
+  batch_size: 300
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_r50_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'resnet50_fc512'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/resnet50_market1501_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0003
+  max_epoch: 60
+  batch_size: 32
+  fixbase_epoch: 5
+  open_layers: ['classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [20]
+
+test:
+  batch_size: 100
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/configs/im_r50fc512_softmax_256x128_amsgrad.yaml

@@ -0,0 +1,36 @@
+model:
+  name: 'resnet50_fc512'
+  pretrained: True
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip']
+  save_dir: 'log/resnet50_fc512_market1501_softmax'
+
+loss:
+  name: 'softmax'
+  softmax:
+    label_smooth: True
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0003
+  max_epoch: 60
+  batch_size: 32
+  fixbase_epoch: 5
+  open_layers: ['fc', 'classifier']
+  lr_scheduler: 'single_step'
+  stepsize: [20]
+
+test:
+  batch_size: 100
+  dist_metric: 'euclidean'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/docs/AWESOME_REID.md

@@ -0,0 +1,69 @@
+# Awesome-ReID
+Here is a collection of ReID-related research with links to papers and code. You are welcome to submit [PR](https://help.github.com/articles/creating-a-pull-request/)s if you find something missing.
+
+
+- [TPAMI21] Learning Generalisable Omni-Scale Representations for Person Re-Identification [[paper](https://arxiv.org/abs/1910.06827)][[code](https://github.com/KaiyangZhou/deep-person-reid)]
+
+- [TPAMI21] Deep Learning for Person Re-identification: A Survey and Outlook [[paper](https://arxiv.org/abs/2001.04193)] [[code](https://github.com/mangye16/ReID-Survey)]
+
+- [ICCV19] RGB-Infrared Cross-Modality Person Re-Identification via Joint Pixel and Feature Alignment. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Wang_RGB-Infrared_Cross-Modality_Person_Re-Identification_via_Joint_Pixel_and_Feature_Alignment_ICCV_2019_paper.pdf)] [[code](https://github.com/wangguanan/AlignGAN)]
+
+- [ICCV19] Unsupervised Graph Association for Person Re-identification. [[paper](https://github.com/yichuan9527/Unsupervised-Graph-Association-for-Person-Re-identification)] [[code](https://github.com/yichuan9527/Unsupervised-Graph-Association-for-Person-Re-identification)]
+
+- [ICCV19] Self-similarity Grouping: A Simple Unsupervised Cross Domain Adaptation Approach for Person Re-identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Fu_Self-Similarity_Grouping_A_Simple_Unsupervised_Cross_Domain_Adaptation_Approach_for_ICCV_2019_paper.pdf)] [[code](https://github.com/OasisYang/SSG)]
+
+- [ICCV19] Spectral Feature Transformation for Person Re-Identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Luo_Spectral_Feature_Transformation_for_Person_Re-Identification_ICCV_2019_paper.pdf)] [[code](https://github.com/LuckyDC/SFT_REID)]
+
+- [ICCV19] Beyond Human Parts: Dual Part-Aligned Representations for Person Re-Identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Guo_Beyond_Human_Parts_Dual_Part-Aligned_Representations_for_Person_Re-Identification_ICCV_2019_paper.pdf)] [[code](https://github.com/ggjy/P2Net.pytorch)]
+
+- [ICCV19] Co-segmentation Inspired Attention Networks for Video-based Person Re-identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Subramaniam_Co-Segmentation_Inspired_Attention_Networks_for_Video-Based_Person_Re-Identification_ICCV_2019_paper.pdf)][[code](https://github.com/InnovArul/vidreid_cosegmentation)]
+
+- [ICCV19] Mixed High-Order Attention Network for Person Re-Identification. [[paper](https://arxiv.org/abs/1908.05819)][[code](https://github.com/chenbinghui1/MHN)]
+
+- [ICCV19] ABD-Net: Attentive but Diverse Person Re-Identification. [[paper](https://arxiv.org/abs/1908.01114)] [[code](https://github.com/TAMU-VITA/ABD-Net)]
+
+- [ICCV19] Omni-Scale Feature Learning for Person Re-Identification. [[paper](https://arxiv.org/abs/1905.00953)] [[code](https://github.com/KaiyangZhou/deep-person-reid)]
+
+- [CVPR19] Joint Discriminative and Generative Learning for Person Re-identification. [[paper](https://arxiv.org/abs/1904.07223)][[code](https://github.com/NVlabs/DG-Net)]
+- [CVPR19] Invariance Matters: Exemplar Memory for Domain Adaptive Person Re-identification. [[paper](https://arxiv.org/abs/1904.01990)][[code](https://github.com/zhunzhong07/ECN)]
+- [CVPR19] Dissecting Person Re-identification from the Viewpoint of Viewpoint. [[paper](https://arxiv.org/abs/1812.02162)][[code](https://github.com/sxzrt/Dissecting-Person-Re-ID-from-the-Viewpoint-of-Viewpoint)]
+- [CVPR19] Unsupervised Person Re-identification by Soft Multilabel Learning. [[paper](https://arxiv.org/abs/1903.06325)][[code](https://github.com/KovenYu/MAR)]
+- [CVPR19] Patch-based Discriminative Feature Learning for Unsupervised Person Re-identification. [[paper](https://kovenyu.com/publication/2019-cvpr-pedal/)][[code](https://github.com/QizeYang/PAUL)]
+
+- [AAAI19] Spatial and Temporal Mutual Promotion for Video-based Person Re-identification. [[paper](https://arxiv.org/abs/1812.10305)][[code](https://github.com/yolomax/person-reid-lib)]
+
+- [AAAI19] Spatial-Temporal Person Re-identification. [[paper](https://arxiv.org/abs/1812.03282)][[code](https://github.com/Wanggcong/Spatial-Temporal-Re-identification)]
+
+- [AAAI19] Horizontal Pyramid Matching for Person Re-identification. [[paper](https://arxiv.org/abs/1804.05275)][[code](https://github.com/OasisYang/HPM)]
+
+- [AAAI19] Backbone Can Not be Trained at Once: Rolling Back to Pre-trained Network for Person Re-identification. [[paper](https://arxiv.org/abs/1901.06140)][[code](https://github.com/youngminPIL/rollback)]
+
+- [AAAI19] A Bottom-Up Clustering Approach to Unsupervised Person Re-identification. [[paper](https://vana77.github.io/vana77.github.io/images/AAAI19.pdf)][[code](https://github.com/vana77/Bottom-up-Clustering-Person-Re-identification)]
+
+- [NIPS18] FD-GAN: Pose-guided Feature Distilling GAN for Robust Person Re-identification. [[paper](https://arxiv.org/abs/1810.02936)][[code](https://github.com/yxgeee/FD-GAN)]
+
+- [ECCV18] Generalizing A Person Retrieval Model Hetero- and Homogeneously. [[paper](http://openaccess.thecvf.com/content_ECCV_2018/papers/Zhun_Zhong_Generalizing_A_Person_ECCV_2018_paper.pdf)][[code](https://github.com/zhunzhong07/HHL)]
+
+- [ECCV18] Pose-Normalized Image Generation for Person Re-identification. [[paper](https://arxiv.org/abs/1712.02225)][[code](https://github.com/naiq/PN_GAN)]
+
+- [CVPR18] Camera Style Adaptation for Person Re-Identification. [[paper](https://arxiv.org/abs/1711.10295)][[code](https://github.com/zhunzhong07/CamStyle)]
+
+- [CVPR18] Deep Group-Shuffling Random Walk for Person Re-Identification. [[paper](https://arxiv.org/abs/1807.11178)][[code](https://github.com/YantaoShen/kpm_rw_person_reid)]
+
+- [CVPR18] End-to-End Deep Kronecker-Product Matching for Person Re-identification. [[paper](https://arxiv.org/abs/1807.11182)][[code](https://github.com/YantaoShen/kpm_rw_person_reid)]
+
+- [CVPR18] Features for Multi-Target Multi-Camera Tracking and Re-Identification. [[paper](https://arxiv.org/abs/1803.10859)][[code](https://github.com/ergysr/DeepCC)]
+
+- [CVPR18] Group Consistent Similarity Learning via Deep CRF for Person Re-Identification. [[paper](http://openaccess.thecvf.com/content_cvpr_2018/papers/Chen_Group_Consistent_Similarity_CVPR_2018_paper.pdf)][[code](https://github.com/dapengchen123/crf_affinity)]
+
+- [CVPR18] Harmonious Attention Network for Person Re-Identification. [[paper](https://arxiv.org/abs/1802.08122)][[code](https://github.com/KaiyangZhou/deep-person-reid)]
+
+- [CVPR18] Human Semantic Parsing for Person Re-Identification. [[paper](https://arxiv.org/abs/1804.00216)][[code](https://github.com/emrahbasaran/SPReID)]
+
+- [CVPR18] Multi-Level Factorisation Net for Person Re-Identification. [[paper](https://arxiv.org/abs/1803.09132)][[code](https://github.com/KaiyangZhou/deep-person-reid)]
+
+- [CVPR18] Resource Aware Person Re-identification across Multiple Resolutions. [[paper](https://arxiv.org/abs/1805.08805)][[code](https://github.com/mileyan/DARENet)]
+
+- [CVPR18] Exploit the Unknown Gradually: One-Shot Video-Based Person Re-Identification by Stepwise Learning. [[paper](https://yu-wu.net/pdf/CVPR2018_Exploit-Unknown-Gradually.pdf)][[code](https://github.com/Yu-Wu/Exploit-Unknown-Gradually)]
+
+- [ArXiv18] Revisiting Temporal Modeling for Video-based Person ReID. [[paper](https://arxiv.org/abs/1805.02104)][[code](https://github.com/jiyanggao/Video-Person-ReID)]

strong_sort/deep/reid/docs/MODEL_ZOO.md

@@ -0,0 +1,93 @@
+# Model Zoo
+
+- Results are presented in the format of *<Rank-1 (mAP)>*.
+- When computing model size and FLOPs, only layers that are used at test time are considered (see `torchreid.utils.compute_model_complexity`).
+- Asterisk (\*) means the model is trained from scratch.
+- `combineall=True` means all images in the dataset are used for model training.
+- Why not use heavy data augmentation like [random erasing](https://arxiv.org/abs/1708.04896) for model training? It's because heavy data augmentation might harm the cross-dataset generalization performance (see [this paper](https://arxiv.org/abs/1708.04896)).
+
+
+## ImageNet pretrained models
+
+
+| Model | Download |
+| :--- | :---: |
+| shufflenet | [model](https://drive.google.com/file/d/1RFnYcHK1TM-yt3yLsNecaKCoFO4Yb6a-/view?usp=sharing) |
+| mobilenetv2_x1_0 | [model](https://drive.google.com/file/d/1K7_CZE_L_Tf-BRY6_vVm0G-0ZKjVWh3R/view?usp=sharing) |
+| mobilenetv2_x1_4 | [model](https://drive.google.com/file/d/10c0ToIGIVI0QZTx284nJe8QfSJl5bIta/view?usp=sharing) |
+| mlfn | [model](https://drive.google.com/file/d/1PP8Eygct5OF4YItYRfA3qypYY9xiqHuV/view?usp=sharing) |
+| osnet_x1_0 | [model](https://drive.google.com/file/d/1LaG1EJpHrxdAxKnSCJ_i0u-nbxSAeiFY/view?usp=sharing) |
+| osnet_x0_75 | [model](https://drive.google.com/file/d/1uwA9fElHOk3ZogwbeY5GkLI6QPTX70Hq/view?usp=sharing) |
+| osnet_x0_5 | [model](https://drive.google.com/file/d/16DGLbZukvVYgINws8u8deSaOqjybZ83i/view?usp=sharing) |
+| osnet_x0_25 | [model](https://drive.google.com/file/d/1rb8UN5ZzPKRc_xvtHlyDh-cSz88YX9hs/view?usp=sharing) |
+| osnet_ibn_x1_0 | [model](https://drive.google.com/file/d/1sr90V6irlYYDd4_4ISU2iruoRG8J__6l/view?usp=sharing) |
+| osnet_ain_x1_0 | [model](https://drive.google.com/file/d/1-CaioD9NaqbHK_kzSMW8VE4_3KcsRjEo/view?usp=sharing) |
+| osnet_ain_x0_75 | [model](https://drive.google.com/file/d/1apy0hpsMypqstfencdH-jKIUEFOW4xoM/view?usp=sharing) |
+| osnet_ain_x0_5 | [model](https://drive.google.com/file/d/1KusKvEYyKGDTUBVRxRiz55G31wkihB6l/view?usp=sharing) |
+| osnet_ain_x0_25 | [model](https://drive.google.com/file/d/1SxQt2AvmEcgWNhaRb2xC4rP6ZwVDP0Wt/view?usp=sharing) |
+
+
+## Same-domain ReID
+
+
+| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | market1501  | dukemtmcreid | msmt17 |
+| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| resnet50 | 23.5 | 2.7 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [87.9 (70.4)](https://drive.google.com/file/d/1dUUZ4rHDWohmsQXCRe2C_HbYkzz94iBV/view?usp=sharing) | [78.3 (58.9)](https://drive.google.com/file/d/17ymnLglnc64NRvGOitY3BqMRS9UWd1wg/view?usp=sharing) | [63.2 (33.9)](https://drive.google.com/file/d/1ep7RypVDOthCRIAqDnn4_N-UhkkFHJsj/view?usp=sharing) |
+| resnet50_fc512 | 24.6 | 4.1 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [90.8 (75.3)](https://drive.google.com/file/d/1kv8l5laX_YCdIGVCetjlNdzKIA3NvsSt/view?usp=sharing) | [81.0 (64.0)](https://drive.google.com/file/d/13QN8Mp3XH81GK4BPGXobKHKyTGH50Rtx/view?usp=sharing) | [69.6 (38.4)](https://drive.google.com/file/d/1fDJLcz4O5wxNSUvImIIjoaIF9u1Rwaud/view?usp=sharing) |
+| mlfn | 32.5 | 2.8 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [90.1 (74.3)](https://drive.google.com/file/d/1wXcvhA_b1kpDfrt9s2Pma-MHxtj9pmvS/view?usp=sharing) | [81.1 (63.2)](https://drive.google.com/file/d/1rExgrTNb0VCIcOnXfMsbwSUW1h2L1Bum/view?usp=sharing) | [66.4 (37.2)](https://drive.google.com/file/d/18JzsZlJb3Wm7irCbZbZ07TN4IFKvR6p-/view?usp=sharing) |
+| hacnn<sup>*</sup> | 4.5 | 0.5 | softmax | (160, 64) | `random_flip`, `random_crop` | `euclidean` | [90.9 (75.6)](https://drive.google.com/file/d/1LRKIQduThwGxMDQMiVkTScBwR7WidmYF/view?usp=sharing) | [80.1 (63.2)](https://drive.google.com/file/d/1zNm6tP4ozFUCUQ7Sv1Z98EAJWXJEhtYH/view?usp=sharing) | [64.7 (37.2)](https://drive.google.com/file/d/1MsKRtPM5WJ3_Tk2xC0aGOO7pM3VaFDNZ/view?usp=sharing) |
+| mobilenetv2_x1_0 | 2.2 | 0.2 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [85.6 (67.3)](https://drive.google.com/file/d/18DgHC2ZJkjekVoqBWszD8_Xiikz-fewp/view?usp=sharing) | [74.2 (54.7)](https://drive.google.com/file/d/1q1WU2FETRJ3BXcpVtfJUuqq4z3psetds/view?usp=sharing) | [57.4 (29.3)](https://drive.google.com/file/d/1j50Hv14NOUAg7ZeB3frzfX-WYLi7SrhZ/view?usp=sharing) |
+| mobilenetv2_x1_4 | 4.3 | 0.4 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [87.0 (68.5)](https://drive.google.com/file/d/1t6JCqphJG-fwwPVkRLmGGyEBhGOf2GO5/view?usp=sharing) | [76.2 (55.8)](https://drive.google.com/file/d/12uD5FeVqLg9-AFDju2L7SQxjmPb4zpBN/view?usp=sharing) | [60.1 (31.5)](https://drive.google.com/file/d/1ZY5P2Zgm-3RbDpbXM0kIBMPvspeNIbXz/view?usp=sharing) |
+| osnet_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip` | `euclidean` | [94.2 (82.6)](https://drive.google.com/file/d/1vduhq5DpN2q1g4fYEZfPI17MJeh9qyrA/view?usp=sharing) | [87.0 (70.2)](https://drive.google.com/file/d/1QZO_4sNf4hdOKKKzKc-TZU9WW1v6zQbq/view?usp=sharing) | [74.9 (43.8)](https://drive.google.com/file/d/112EMUfBPYeYg70w-syK6V6Mx8-Qb9Q1M/view?usp=sharing) |
+| osnet_x0_75 | 1.3 | 0.57 | softmax | (256, 128) | `random_flip` | `euclidean` | [93.7 (81.2)](https://drive.google.com/file/d/1ozRaDSQw_EQ8_93OUmjDbvLXw9TnfPer/view?usp=sharing) | [85.8 (69.8)](https://drive.google.com/file/d/1IE3KRaTPp4OUa6PGTFL_d5_KQSJbP0Or/view?usp=sharing) | [72.8 (41.4)](https://drive.google.com/file/d/1QEGO6WnJ-BmUzVPd3q9NoaO_GsPNlmWc/view?usp=sharing) |
+| osnet_x0_5 | 0.6 | 0.27 | softmax | (256, 128) | `random_flip` | `euclidean` | [92.5 (79.8)](https://drive.google.com/file/d/1PLB9rgqrUM7blWrg4QlprCuPT7ILYGKT/view?usp=sharing) | [85.1 (67.4)](https://drive.google.com/file/d/1KoUVqmiST175hnkALg9XuTi1oYpqcyTu/view?usp=sharing) | [69.7 (37.5)](https://drive.google.com/file/d/1UT3AxIaDvS2PdxzZmbkLmjtiqq7AIKCv/view?usp=sharing) |
+| osnet_x0_25 | 0.2 | 0.08 | softmax | (256, 128) | `random_flip` | `euclidean` | [91.2 (75.0)](https://drive.google.com/file/d/1z1UghYvOTtjx7kEoRfmqSMu-z62J6MAj/view?usp=sharing) | [82.0 (61.4)](https://drive.google.com/file/d/1eumrtiXT4NOspjyEV4j8cHmlOaaCGk5l/view?usp=sharing) | [61.4 (29.5)](https://drive.google.com/file/d/1sSwXSUlj4_tHZequ_iZ8w_Jh0VaRQMqF/view?usp=sharing) |
+
+
+## Cross-domain ReID
+
+#### Market1501 -> DukeMTMC-reID
+
+
+| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance  | Rank-1 | Rank-5 | Rank-10 | mAP | Download |
+| :--- | :---: | :---: | :---: | :---: |  :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| osnet_ibn_x1_0 | 2.2 | 0.98  | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 48.5 | 62.3 | 67.4 | 26.7 | [model](https://drive.google.com/file/d/1uWW7_z_IcUmRNPqQOrEBdsvic94fWH37/view?usp=sharing) |
+| osnet_ain_x1_0 | 2.2 | 0.98  | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | 52.4 | 66.1 | 71.2 | 30.5 | [model](https://drive.google.com/file/d/14bNFGm0FhwHEkEpYKqKiDWjLNhXywFAd/view?usp=sharing) |
+
+
+#### DukeMTMC-reID -> Market1501
+
+
+| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance  | Rank-1 | Rank-5 | Rank-10 | mAP | Download |
+| :--- | :---: | :---: | :---: | :---: |  :---: | :---: | :---: | :---: | :---: | :---: | :---: |
+| osnet_ibn_x1_0 | 2.2 | 0.98  | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 57.7 | 73.7 | 80.0 | 26.1 | [model](https://drive.google.com/file/d/1CNxL1IP0BjcE1TSttiVOID1VNipAjiF3/view?usp=sharing) |
+| osnet_ain_x1_0 | 2.2 | 0.98  | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | 61.0 | 77.0 | 82.5 | 30.6 | [model](https://drive.google.com/file/d/1hypJvq8G04SOby6jvF337GEkg5K_bmCw/view?usp=sharing) |
+
+
+#### MSMT17 (`combineall=True`) -> Market1501 & DukeMTMC-reID
+
+
+| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | msmt17 -> market1501 | msmt17 -> dukemtmcreid | Download |
+| :--- | :---: | :---: | :---: | :---: |  :---: | :---: | :---: | :---: | :---: |
+| resnet50 | 23.5 | 2.7 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 46.3 (22.8) | 52.3 (32.1) | [model](https://drive.google.com/file/d/1yiBteqgIZoOeywE8AhGmEQl7FTVwrQmf/view?usp=sharing) |
+| osnet_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 66.6 (37.5) | 66.0 (45.3) | [model](https://drive.google.com/file/d/1IosIFlLiulGIjwW3H8uMRmx3MzPwf86x/view?usp=sharing) |
+| osnet_x0_75 | 1.3 | 0.57 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 63.6 (35.5) | 65.3 (44.5) | [model](https://drive.google.com/file/d/1fhjSS_7SUGCioIf2SWXaRGPqIY9j7-uw/view?usp=sharing) |
+| osnet_x0_5 | 0.6 | 0.27 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 64.3 (34.9) | 65.2 (43.3) | [model](https://drive.google.com/file/d/1DHgmb6XV4fwG3n-CnCM0zdL9nMsZ9_RF/view?usp=sharing) |
+| osnet_x0_25 | 0.2 | 0.08 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 59.9 (31.0) | 61.5 (39.6) | [model](https://drive.google.com/file/d/1Kkx2zW89jq_NETu4u42CFZTMVD5Hwm6e/view?usp=sharing) |
+| osnet_ibn_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 66.5 (37.2) | 67.4 (45.6) | [model](https://drive.google.com/file/d/1q3Sj2ii34NlfxA4LvmHdWO_75NDRmECJ/view?usp=sharing) |
+| osnet_ain_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | 70.1 (43.3) | 71.1 (52.7) | [model](https://drive.google.com/file/d/1SigwBE6mPdqiJMqhuIY4aqC7--5CsMal/view?usp=sharing) |
+
+
+#### Multi-source domain generalization
+
+The models below are trained using multiple source datasets, as described in [Zhou et al. TPAMI'21](https://arxiv.org/abs/1910.06827).
+
+Regarding the abbreviations, MS is MSMT17; M is Market1501; D is DukeMTMC-reID; and C is CUHK03.
+
+All models were trained with [im_osnet_ain_x1_0_softmax_256x128_amsgrad_cosine.yaml](https://github.com/KaiyangZhou/deep-person-reid/blob/master/configs/im_osnet_ain_x1_0_softmax_256x128_amsgrad_cosine.yaml) and `max_epoch=50`.
+
+| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | MS+D+C->M | MS+M+C->D | MS+D+M->C |D+M+C->MS |
+| :--- | :---: | :---: | :---: | :---: |  :---: | :---: | :---: | :---: | :---: | :---: |
+| osnet_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | [72.5 (44.2)](https://drive.google.com/file/d/1tuYY1vQXReEd8N8_npUkc7npPDDmjNCV/view?usp=sharing) | [65.2 (47.0)](https://drive.google.com/file/d/1UxUI4NsE108UCvcy3O1Ufe73nIVPKCiu/view?usp=sharing) | [23.9 (23.3)](https://drive.google.com/file/d/1kAA6qHJvbaJtyh1b39ZyEqWROwUgWIhl/view?usp=sharing) | [33.2 (12.6)](https://drive.google.com/file/d/1wAHuYVTzj8suOwqCNcEmu6YdbVnHDvA2/view?usp=sharing) |
+| osnet_ibn_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | [73.0 (44.9)](https://drive.google.com/file/d/14sH6yZwuNHPTElVoEZ26zozOOZIej5Mf/view?usp=sharing) | [64.6 (45.7)](https://drive.google.com/file/d/1Sk-2SSwKAF8n1Z4p_Lm_pl0E6v2WlIBn/view?usp=sharing) | [25.7 (25.4)](https://drive.google.com/file/d/1actHP7byqWcK4eBE1ojnspSMdo7k2W4G/view?usp=sharing) | [39.8 (16.2)](https://drive.google.com/file/d/1BGOSdLdZgqHe2qFafatb-5sPY40JlYfp/view?usp=sharing) |
+| osnet_ain_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | [73.3 (45.8)](https://drive.google.com/file/d/1nIrszJVYSHf3Ej8-j6DTFdWz8EnO42PB/view?usp=sharing) | [65.6 (47.2)](https://drive.google.com/file/d/1YjJ1ZprCmaKG6MH2P9nScB9FL_Utf9t1/view?usp=sharing) | [27.4 (27.1)](https://drive.google.com/file/d/1IxIg5P0cei3KPOJQ9ZRWDE_Mdrz01ha2/view?usp=sharing) | [40.2 (16.2)](https://drive.google.com/file/d/1KcoUKzLmsUoGHI7B6as_Z2fXL50gzexS/view?usp=sharing) |

strong_sort/deep/reid/docs/Makefile

@@ -0,0 +1,19 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+SPHINXOPTS    =
+SPHINXBUILD   = sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

strong_sort/deep/reid/docs/conf.py

@@ -0,0 +1,181 @@
+# -*- coding: utf-8 -*-
+#
+# Configuration file for the Sphinx documentation builder.
+#
+# This file does only contain a selection of the most common options. For a
+# full list see the documentation:
+# http://www.sphinx-doc.org/en/master/config
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath('..'))
+
+# -- Project information -----------------------------------------------------
+
+project = u'torchreid'
+copyright = u'2019, Kaiyang Zhou'
+author = u'Kaiyang Zhou'
+
+version_file = '../torchreid/__init__.py'
+with open(version_file, 'r') as f:
+    exec(compile(f.read(), version_file, 'exec'))
+__version__ = locals()['__version__']
+
+# The short X.Y version
+version = __version__
+# The full version, including alpha/beta/rc tags
+release = __version__
+
+# -- General configuration ---------------------------------------------------
+
+# If your documentation needs a minimal Sphinx version, state it here.
+#
+# needs_sphinx = '1.0'
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinxcontrib.napoleon',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.githubpages',
+    'sphinx_markdown_tables',
+]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ['_templates']
+
+# The suffix(es) of source filenames.
+# You can specify multiple suffix as a list of string:
+#
+source_suffix = ['.rst', '.md']
+# source_suffix = '.rst'
+source_parsers = {'.md': 'recommonmark.parser.CommonMarkParser'}
+
+# The master toctree document.
+master_doc = 'index'
+
+# The language for content autogenerated by Sphinx. Refer to documentation
+# for a list of supported languages.
+#
+# This is also used if you do content translation via gettext catalogs.
+# Usually you set "language" from the command line for these cases.
+language = None
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = [u'_build', 'Thumbs.db', '.DS_Store']
+
+# The name of the Pygments (syntax highlighting) style to use.
+pygments_style = None
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = 'sphinx_rtd_theme'
+
+# Theme options are theme-specific and customize the look and feel of a theme
+# further.  For a list of options available for each theme, see the
+# documentation.
+#
+# html_theme_options = {}
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ['_static']
+
+# Custom sidebar templates, must be a dictionary that maps document names
+# to template names.
+#
+# The default sidebars (for documents that don't match any pattern) are
+# defined by theme itself.  Builtin themes are using these templates by
+# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',
+# 'searchbox.html']``.
+#
+# html_sidebars = {}
+
+# -- Options for HTMLHelp output ---------------------------------------------
+
+# Output file base name for HTML help builder.
+htmlhelp_basename = 'torchreiddoc'
+
+# -- Options for LaTeX output ------------------------------------------------
+
+latex_elements = {
+    # The paper size ('letterpaper' or 'a4paper').
+    #
+    # 'papersize': 'letterpaper',
+
+    # The font size ('10pt', '11pt' or '12pt').
+    #
+    # 'pointsize': '10pt',
+
+    # Additional stuff for the LaTeX preamble.
+    #
+    # 'preamble': '',
+
+    # Latex figure (float) alignment
+    #
+    # 'figure_align': 'htbp',
+}
+
+# Grouping the document tree into LaTeX files. List of tuples
+# (source start file, target name, title,
+#  author, documentclass [howto, manual, or own class]).
+latex_documents = [
+    (
+        master_doc, 'torchreid.tex', u'torchreid Documentation',
+        u'Kaiyang Zhou', 'manual'
+    ),
+]
+
+# -- Options for manual page output ------------------------------------------
+
+# One entry per manual page. List of tuples
+# (source start file, name, description, authors, manual section).
+man_pages = [
+    (master_doc, 'torchreid', u'torchreid Documentation', [author], 1)
+]
+
+# -- Options for Texinfo output ----------------------------------------------
+
+# Grouping the document tree into Texinfo files. List of tuples
+# (source start file, target name, title, author,
+#  dir menu entry, description, category)
+texinfo_documents = [
+    (
+        master_doc, 'torchreid', u'torchreid Documentation', author,
+        'torchreid', 'One line description of project.', 'Miscellaneous'
+    ),
+]
+
+# -- Options for Epub output -------------------------------------------------
+
+# Bibliographic Dublin Core info.
+epub_title = project
+
+# The unique identifier of the text. This can be a ISBN number
+# or the project homepage.
+#
+# epub_identifier = ''
+
+# A unique identification for the text.
+#
+# epub_uid = ''
+
+# A list of files that should not be packed into the epub file.
+epub_exclude_files = ['search.html']
+
+# -- Extension configuration -------------------------------------------------

strong_sort/deep/reid/docs/datasets.rst

@@ -0,0 +1,264 @@
+.. _datasets:
+
+Datasets
+=========
+
+Here we provide a comprehensive guide on how to prepare the datasets.
+
+Suppose you want to store the reid data in a directory called "path/to/reid-data/", you need to specify the ``root`` as *root='path/to/reid-data/'* when initializing ``DataManager``. Below we use ``$REID`` to denote "path/to/reid-data".
+
+Please refer to :ref:`torchreid_data` for details regarding the arguments.
+
+
+.. note::
+    Dataset with a :math:`\dagger` symbol means that the process is automated, so you can directly call the dataset in ``DataManager`` (which automatically downloads the dataset and organizes the data structure). However, we also provide a way below to help the manual setup in case the automation fails.
+
+
+.. note::
+    The keys to use specific datasets are enclosed in the parantheses beside the datasets' names.
+
+
+.. note::
+    You are suggested to use the provided names for dataset folders such as "market1501" for Market1501 and "dukemtmcreid" for DukeMTMC-reID when doing the manual setup, otherwise you need to modify the source code accordingly (i.e. the ``dataset_dir`` attribute).
+
+.. note::
+    Some download links provided by the original authors might not work. You can email `Kaiyang Zhou <https://kaiyangzhou.github.io/>`_ to reqeust new links. Please do provide your full name, institution, and purpose of using the data in the email (best use your work email address).
+
+.. contents::
+   :local:
+
+
+Image Datasets
+--------------
+
+Market1501 :math:`^\dagger` (``market1501``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create a directory named "market1501" under ``$REID``.
+- Download the dataset to "market1501" from http://www.liangzheng.org/Project/project_reid.html and extract the files.
+- The data structure should look like
+
+.. code-block:: none
+    
+    market1501/
+        Market-1501-v15.09.15/
+            query/
+            bounding_box_train/
+            bounding_box_test/
+
+- To use the extra 500K distractors (i.e. Market1501 + 500K), go to the **Market-1501+500k Dataset** section at http://www.liangzheng.org/Project/project_reid.html, download the zip file "distractors_500k.zip" and extract it under "market1501/Market-1501-v15.09.15". The argument to use these 500K distrctors is ``market1501_500k`` in ``ImageDataManager``.
+
+
+CUHK03 (``cuhk03``)
+^^^^^^^^^^^^^^^^^^^^^
+- Create a folder named "cuhk03" under ``$REID``.
+- Download the dataset to "cuhk03/" from http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html and extract "cuhk03_release.zip", resulting in "cuhk03/cuhk03_release/".
+- Download the new split (767/700) from `person-re-ranking <https://github.com/zhunzhong07/person-re-ranking/tree/master/evaluation/data/CUHK03>`_. What you need are "cuhk03_new_protocol_config_detected.mat" and "cuhk03_new_protocol_config_labeled.mat". Put these two mat files under "cuhk03/".
+- The data structure should look like
+
+.. code-block:: none
+    
+    cuhk03/
+        cuhk03_release/
+        cuhk03_new_protocol_config_detected.mat
+        cuhk03_new_protocol_config_labeled.mat
+
+
+- In the default mode, we load data using the new split (767/700). If you wanna use the original (20) splits (1367/100), please set ``cuhk03_classic_split`` to True in ``ImageDataManager``. As the CMC is computed differently from Market1501 for the 1367/100 split (see `here <http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html>`_), you need to enable ``use_metric_cuhk03`` in ``ImageDataManager`` to activate the *single-gallery-shot* metric for fair comparison with some methods that adopt the old splits (*do not need to report mAP*). In addition, we support both *labeled* and *detected* modes. The default mode loads *detected* images. Enable ``cuhk03_labeled`` in ``ImageDataManager`` if you wanna train and test on *labeled* images.
+
+.. note::
+    The code will extract images in "cuhk-03.mat" and save them under "cuhk03/images_detected" and "cuhk03/images_labeled". Also, four json files will be automatically generated, i.e. "splits_classic_detected.json", "splits_classic_labeled.json", "splits_new_detected.json" and "splits_new_labeled.json". If the parent path of ``$REID`` is changed, these json files should be manually deleted. The code can automatically generate new json files to match the new path.
+    
+
+DukeMTMC-reID :math:`^\dagger` (``dukemtmcreid``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create a directory called "dukemtmc-reid" under ``$REID``.
+- Download "DukeMTMC-reID" from http://vision.cs.duke.edu/DukeMTMC/ and extract it under "dukemtmc-reid".
+- The data structure should look like
+
+.. code-block:: none
+    
+    dukemtmc-reid/
+        DukeMTMC-reID/
+            query/
+            bounding_box_train/
+            bounding_box_test/
+            ...
+
+MSMT17 (``msmt17``)
+^^^^^^^^^^^^^^^^^^^^^
+- Create a directory called "msmt17" under ``$REID``.
+- Download the dataset from http://www.pkuvmc.com/publications/msmt17.html to "msmt17" and extract the files.
+- The data structure should look like
+
+.. code-block:: none
+    
+    msmt17/
+        MSMT17_V1/ # or MSMT17_V2
+            train/
+            test/
+            list_train.txt
+            list_query.txt
+            list_gallery.txt
+            list_val.txt
+
+VIPeR :math:`^\dagger` (``viper``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- The download link is http://users.soe.ucsc.edu/~manduchi/VIPeR.v1.0.zip.
+- Organize the dataset in a folder named "viper" as follows
+
+.. code-block:: none
+    
+    viper/
+        VIPeR/
+            cam_a/
+            cam_b/
+
+GRID :math:`^\dagger` (``grid``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- The download link is http://personal.ie.cuhk.edu.hk/~ccloy/files/datasets/underground_reid.zip.
+- Organize the dataset in a folder named "grid" as follows
+
+.. code-block:: none
+    
+    grid/
+        underground_reid/
+            probe/
+            gallery/
+            ...
+
+CUHK01 (``cuhk01``)
+^^^^^^^^^^^^^^^^^^^^^^^^
+- Create a folder named "cuhk01" under ``$REID``.
+- Download "CUHK01.zip" from http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html and place it under "cuhk01/".
+- The code can automatically extract the files, or you can do it yourself.
+- The data structure should look like
+
+.. code-block:: none
+    
+    cuhk01/
+        campus/
+
+SenseReID (``sensereid``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create "sensereid" under ``$REID``.
+- Download the dataset from this `link <https://drive.google.com/file/d/0B56OfSrVI8hubVJLTzkwV2VaOWM/view>`_ and extract it to "sensereid".
+- Organize the data to be like
+
+.. code-block:: none
+
+    sensereid/
+        SenseReID/
+            test_probe/
+            test_gallery/
+
+QMUL-iLIDS :math:`^\dagger` (``ilids``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create a folder named "ilids" under ``$REID``.
+- Download the dataset from http://www.eecs.qmul.ac.uk/~jason/data/i-LIDS_Pedestrian.tgz and organize it to look like
+
+.. code-block:: none
+    
+    ilids/
+        i-LIDS_Pedestrian/
+            Persons/
+
+PRID (``prid``)
+^^^^^^^^^^^^^^^^^^^
+- Create a directory named "prid2011" under ``$REID``.
+- Download the dataset from https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/PRID11/ and extract it under "prid2011".
+- The data structure should end up with
+
+.. code-block:: none
+
+    prid2011/
+        prid_2011/
+            single_shot/
+            multi_shot/
+
+CUHK02 (``cuhk02``)
+^^^^^^^^^^^^^^^^^^^^^
+- Create a folder named "cuhk02" under ``$REID``.
+- Download the data from http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html and put it under "cuhk02/".
+- Extract the file so the data structure looks like
+
+.. code-block:: none
+    
+    cuhk02/
+        Dataset/
+            P1/
+            P2/
+            P3/
+            P4/
+            P5/
+
+CUHKSYSU (``cuhksysu``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create a folder named "cuhksysu" under ``$REID``.
+- Download the data to "cuhksysu/" from this `google drive link <https://drive.google.com/file/d/1XmiNVrfK2ZmI0ZZ2HHT80HHbDrnE4l3W/view?usp=sharing>`_.
+- Extract the zip file under "cuhksysu/".
+- The data structure should look like
+
+.. code-block:: none
+    
+    cuhksysu/
+        cropped_images
+
+
+Video Datasets
+--------------
+
+MARS (``mars``)
+^^^^^^^^^^^^^^^^^
+- Create "mars/" under ``$REID``.
+- Download the dataset from http://www.liangzheng.com.cn/Project/project_mars.html and place it in "mars/".
+- Extract "bbox_train.zip" and "bbox_test.zip".
+- Download the split metadata from https://github.com/liangzheng06/MARS-evaluation/tree/master/info and put "info/" in "mars/".
+- The data structure should end up with
+
+.. code-block:: none
+    
+    mars/
+        bbox_test/
+        bbox_train/
+        info/
+
+iLIDS-VID :math:`^\dagger` (``ilidsvid``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create "ilids-vid" under ``$REID``.
+- Download the dataset from https://xiatian-zhu.github.io/downloads_qmul_iLIDS-VID_ReID_dataset.html to "ilids-vid".
+- Organize the data structure to match
+
+.. code-block:: none
+    
+    ilids-vid/
+        i-LIDS-VID/
+        train-test people splits/
+
+PRID2011 (``prid2011``)
+^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create a directory named "prid2011" under ``$REID``.
+- Download the dataset from https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/PRID11/ and extract it under "prid2011".
+- Download the split created by *iLIDS-VID* from `this google drive <https://drive.google.com/open?id=1qw7SI7YdIgfHetIQO7LLW4SHpL_qkieT>`_ and put it under "prid2011/". Following the standard protocol, only 178 persons whose sequences are more than a threshold are used.
+- The data structure should end up with
+
+.. code-block:: none
+
+    prid2011/
+        splits_prid2011.json
+        prid_2011/
+            single_shot/
+            multi_shot/
+
+DukeMTMC-VideoReID :math:`^\dagger` (``dukemtmcvidreid``)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+- Create "dukemtmc-vidreid" under ``$REID``.
+- Download "DukeMTMC-VideoReID" from http://vision.cs.duke.edu/DukeMTMC/ and unzip the file to "dukemtmc-vidreid/".
+- The data structure should look like
+
+.. code-block:: none
+    
+    dukemtmc-vidreid/
+        DukeMTMC-VideoReID/
+            train/
+            query/
+            gallery/

strong_sort/deep/reid/docs/evaluation.rst

@@ -0,0 +1,21 @@
+Evaluation
+==========
+
+Image ReID
+-----------
+- **Market1501**, **DukeMTMC-reID**, **CUHK03 (767/700 split)** and **MSMT17** have fixed split so keeping ``split_id=0`` is fine.
+- **CUHK03 (classic split)** has 20 fixed splits, so do ``split_id=0~19``.
+- **VIPeR** contains 632 identities each with 2 images under two camera views. Evaluation should be done for 10 random splits. Each split randomly divides 632 identities to 316 train ids (632 images) and the other 316 test ids (632 images). Note that, in each random split, there are two sub-splits, one using camera-A as query and camera-B as gallery while the other one using camera-B as query and camera-A as gallery. Thus, there are totally 20 splits generated with ``split_id`` starting from 0 to 19. Models can be trained on ``split_id=[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]`` (because ``split_id=0`` and ``split_id=1`` share the same train set, and so on and so forth.). At test time, models trained on ``split_id=0`` can be directly evaluated on ``split_id=1``, models trained on ``split_id=2`` can be directly evaluated on ``split_id=3``, and so on and so forth.
+- **CUHK01** is similar to VIPeR in the split generation.
+- **GRID** , **iLIDS** and **PRID** have 10 random splits, so evaluation should be done by varying ``split_id`` from 0 to 9.
+- **SenseReID** has no training images and is used for evaluation only.
+
+
+.. note::
+    The ``split_id`` argument is defined in ``ImageDataManager`` and ``VideoDataManager``. Please refer to :ref:`torchreid_data`.
+
+
+Video ReID
+-----------
+- **MARS** and **DukeMTMC-VideoReID** have fixed single split so using ``split_id=0`` is ok.
+- **iLIDS-VID** and **PRID2011** have 10 predefined splits so evaluation should be done by varying ``split_id`` from 0 to 9.

strong_sort/deep/reid/docs/index.rst

@@ -0,0 +1,35 @@
+.. include:: ../README.rst
+
+
+.. toctree::
+    :hidden:
+
+    user_guide
+    datasets
+    evaluation
+
+.. toctree::
+    :caption: Package Reference
+    :hidden:
+
+    pkg/data
+    pkg/engine
+    pkg/losses
+    pkg/metrics
+    pkg/models
+    pkg/optim
+    pkg/utils
+
+.. toctree::
+    :caption: Resources
+    :hidden:
+
+    AWESOME_REID.md
+    MODEL_ZOO.md
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`

strong_sort/deep/reid/docs/pkg/data.rst

@@ -0,0 +1,86 @@
+.. _torchreid_data:
+
+torchreid.data
+==============
+
+
+Data Manager
+---------------------------
+
+.. automodule:: torchreid.data.datamanager
+    :members:
+
+
+Sampler
+-----------------------
+
+.. automodule:: torchreid.data.sampler
+    :members:
+
+
+Transforms
+---------------------------
+
+.. automodule:: torchreid.data.transforms
+    :members:
+
+
+Dataset
+---------------------------
+
+.. automodule:: torchreid.data.datasets.dataset
+    :members:
+
+
+.. automodule:: torchreid.data.datasets.__init__
+    :members:
+
+
+Image Datasets
+------------------------------
+
+.. automodule:: torchreid.data.datasets.image.market1501
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.cuhk03
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.dukemtmcreid
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.msmt17
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.viper
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.grid
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.cuhk01
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.ilids
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.sensereid
+    :members:
+
+.. automodule:: torchreid.data.datasets.image.prid
+    :members:
+
+
+Video Datasets
+------------------------------
+
+.. automodule:: torchreid.data.datasets.video.mars
+    :members:
+
+.. automodule:: torchreid.data.datasets.video.ilidsvid
+    :members:
+
+.. automodule:: torchreid.data.datasets.video.prid2011
+    :members:
+
+.. automodule:: torchreid.data.datasets.video.dukemtmcvidreid
+    :members:

strong_sort/deep/reid/docs/pkg/engine.rst

@@ -0,0 +1,31 @@
+.. _torchreid_engine:
+
+torchreid.engine
+==================
+
+
+Base Engine
+------------
+
+.. autoclass:: torchreid.engine.engine.Engine
+    :members:
+
+
+Image Engines
+-------------
+
+.. autoclass:: torchreid.engine.image.softmax.ImageSoftmaxEngine
+    :members:
+
+
+.. autoclass:: torchreid.engine.image.triplet.ImageTripletEngine
+    :members:
+
+
+Video Engines
+-------------
+
+.. autoclass:: torchreid.engine.video.softmax.VideoSoftmaxEngine
+
+
+.. autoclass:: torchreid.engine.video.triplet.VideoTripletEngine

strong_sort/deep/reid/docs/pkg/losses.rst

@@ -0,0 +1,18 @@
+.. _torchreid_losses:
+
+torchreid.losses
+=================
+
+
+Softmax
+--------
+
+.. automodule:: torchreid.losses.cross_entropy_loss
+    :members:
+
+
+Triplet
+-------
+
+.. automodule:: torchreid.losses.hard_mine_triplet_loss
+    :members:

strong_sort/deep/reid/docs/pkg/metrics.rst

@@ -0,0 +1,25 @@
+.. _torchreid_metrics:
+
+torchreid.metrics
+=================
+
+
+Distance
+---------
+
+.. automodule:: torchreid.metrics.distance
+    :members:
+
+
+Accuracy
+--------
+
+.. automodule:: torchreid.metrics.accuracy
+    :members:
+
+
+Rank
+-----
+
+.. automodule:: torchreid.metrics.rank
+    :members: evaluate_rank

strong_sort/deep/reid/docs/pkg/models.rst

@@ -0,0 +1,43 @@
+.. _torchreid_models:
+
+torchreid.models
+=================
+
+Interface
+---------
+
+.. automodule:: torchreid.models.__init__
+    :members:
+
+
+ImageNet Classification Models
+-------------------------------
+
+.. autoclass:: torchreid.models.resnet.ResNet
+.. autoclass:: torchreid.models.senet.SENet
+.. autoclass:: torchreid.models.densenet.DenseNet
+.. autoclass:: torchreid.models.inceptionresnetv2.InceptionResNetV2
+.. autoclass:: torchreid.models.inceptionv4.InceptionV4
+.. autoclass:: torchreid.models.xception.Xception
+
+
+Lightweight Models
+------------------
+
+.. autoclass:: torchreid.models.nasnet.NASNetAMobile
+.. autoclass:: torchreid.models.mobilenetv2.MobileNetV2
+.. autoclass:: torchreid.models.shufflenet.ShuffleNet
+.. autoclass:: torchreid.models.squeezenet.SqueezeNet
+.. autoclass:: torchreid.models.shufflenetv2.ShuffleNetV2
+
+
+ReID-specific Models
+--------------------
+
+.. autoclass:: torchreid.models.mudeep.MuDeep
+.. autoclass:: torchreid.models.resnetmid.ResNetMid
+.. autoclass:: torchreid.models.hacnn.HACNN
+.. autoclass:: torchreid.models.pcb.PCB
+.. autoclass:: torchreid.models.mlfn.MLFN
+.. autoclass:: torchreid.models.osnet.OSNet
+.. autoclass:: torchreid.models.osnet_ain.OSNet

strong_sort/deep/reid/docs/pkg/optim.rst

@@ -0,0 +1,18 @@
+.. _torchreid_optim:
+
+torchreid.optim
+=================
+
+
+Optimizer
+----------
+
+.. automodule:: torchreid.optim.optimizer
+    :members: build_optimizer
+
+
+LR Scheduler
+-------------
+
+.. automodule:: torchreid.optim.lr_scheduler
+    :members: build_lr_scheduler

strong_sort/deep/reid/docs/pkg/utils.rst

@@ -0,0 +1,41 @@
+.. _torchreid_utils:
+
+torchreid.utils
+=================
+
+Average Meter
+--------------
+
+.. automodule:: torchreid.utils.avgmeter
+    :members:
+
+
+Loggers
+-------
+
+.. automodule:: torchreid.utils.loggers
+    :members:
+
+
+Generic Tools
+---------------
+.. automodule:: torchreid.utils.tools
+    :members:
+
+
+ReID Tools
+----------
+
+.. automodule:: torchreid.utils.reidtools
+    :members:
+
+
+Torch Tools
+------------
+
+.. automodule:: torchreid.utils.torchtools
+    :members:
+
+
+.. automodule:: torchreid.utils.model_complexity
+    :members:

strong_sort/deep/reid/docs/user_guide.rst

@@ -0,0 +1,351 @@
+How-to
+============
+
+.. contents::
+   :local:
+
+
+Prepare datasets
+-----------------
+See :ref:`datasets`.
+
+
+Find model keys
+-----------------
+Keys are listed under the *Public keys* section within each model class in :ref:`torchreid_models`.
+
+
+Show available models
+----------------------
+
+.. code-block:: python
+    
+    import torchreid
+    torchreid.models.show_avai_models()
+
+
+Change the training sampler
+-----------------------------
+The default ``train_sampler`` is "RandomSampler". You can give the specific sampler name as input to ``train_sampler``, e.g. ``train_sampler='RandomIdentitySampler'`` for triplet loss.
+
+
+Choose an optimizer/lr_scheduler
+----------------------------------
+Please refer to the source code of ``build_optimizer``/``build_lr_scheduler`` in :ref:`torchreid_optim` for details.
+
+
+Resume training
+----------------
+Suppose the checkpoint is saved in "log/resnet50/model.pth.tar-30", you can do
+
+.. code-block:: python
+    
+    start_epoch = torchreid.utils.resume_from_checkpoint(
+        'log/resnet50/model.pth.tar-30',
+        model,
+        optimizer
+    )
+
+    engine.run(
+        save_dir='log/resnet50',
+        max_epoch=60,
+        start_epoch=start_epoch
+    )
+
+
+Compute model complexity
+--------------------------
+We provide a tool in ``torchreid.utils.model_complexity.py`` to automatically compute the model complexity, i.e. number of parameters and FLOPs.
+
+.. code-block:: python
+    
+    from torchreid import models, utils
+    
+    model = models.build_model(name='resnet50', num_classes=1000)
+    num_params, flops = utils.compute_model_complexity(model, (1, 3, 256, 128))
+
+    # show detailed complexity for each module
+    utils.compute_model_complexity(model, (1, 3, 256, 128), verbose=True)
+
+    # count flops for all layers including ReLU and BatchNorm
+    utils.compute_model_complexity(model, (1, 3, 256, 128), verbose=True, only_conv_linear=False)
+
+Note that (1) this function only provides an estimate of the theoretical time complexity rather than the actual running time which depends on implementations and hardware; (2) the FLOPs is only counted for layers that are used at test time. This means that redundant layers such as person ID classification layer will be ignored. The inference graph depends on how you define the computations in ``forward()``.
+
+
+Combine multiple datasets
+---------------------------
+Easy. Just give whatever datasets (keys) you want to the ``sources`` argument when instantiating a data manager. For example,
+
+.. code-block:: python
+    
+    datamanager = torchreid.data.ImageDataManager(
+        root='reid-data',
+        sources=['market1501', 'dukemtmcreid', 'cuhk03', 'msmt17'],
+        height=256,
+        width=128,
+        batch_size=32
+    )
+
+In this example, the target datasets are Market1501, DukeMTMC-reID, CUHK03 and MSMT17 as the ``targets`` argument is not specified. Please refer to ``Engine.test()`` in :ref:`torchreid_engine` for details regarding how evaluation is performed.
+
+
+Do cross-dataset evaluation
+-----------------------------
+Easy. Just give whatever datasets (keys) you want to the argument ``targets``, like
+
+.. code-block:: python
+    
+    datamanager = torchreid.data.ImageDataManager(
+        root='reid-data',
+        sources='market1501',
+        targets='dukemtmcreid', # or targets='cuhk03' or targets=['dukemtmcreid', 'cuhk03']
+        height=256,
+        width=128,
+        batch_size=32
+    )
+
+
+Combine train, query and gallery
+---------------------------------
+This can be easily done by setting ``combineall=True`` when instantiating a data manager. Below is an example of using Market1501,
+
+.. code-block:: python
+    
+    datamanager = torchreid.data.ImageDataManager(
+        root='reid-data',
+        sources='market1501',
+        height=256,
+        width=128,
+        batch_size=32,
+        market1501_500k=False,
+        combineall=True # it's me, here
+    )
+
+More specifically, with ``combineall=False``, you will get
+
+.. code-block:: none
+    
+    => Loaded Market1501
+      ----------------------------------------
+      subset   | # ids | # images | # cameras
+      ----------------------------------------
+      train    |   751 |    12936 |         6
+      query    |   750 |     3368 |         6
+      gallery  |   751 |    15913 |         6
+      ---------------------------------------
+
+with ``combineall=True``, you will get
+
+.. code-block:: none
+    
+    => Loaded Market1501
+      ----------------------------------------
+      subset   | # ids | # images | # cameras
+      ----------------------------------------
+      train    |  1501 |    29419 |         6
+      query    |   750 |     3368 |         6
+      gallery  |   751 |    15913 |         6
+      ---------------------------------------
+
+
+Optimize layers with different learning rates
+-----------------------------------------------
+A common practice for fine-tuning pretrained models is to use a smaller learning rate for base layers and a large learning rate for randomly initialized layers (referred to as ``new_layers``). ``torchreid.optim.optimizer`` has implemented such feature. What you need to do is to set ``staged_lr=True`` and give the names of ``new_layers`` such as "classifier".
+
+Below is an example of setting different learning rates for base layers and new layers in ResNet50,
+
+.. code-block:: python
+    
+    # New layer "classifier" has a learning rate of 0.01
+    # The base layers have a learning rate of 0.001
+    optimizer = torchreid.optim.build_optimizer(
+        model,
+        optim='sgd',
+        lr=0.01,
+        staged_lr=True,
+        new_layers='classifier',
+        base_lr_mult=0.1
+    )
+
+Please refer to :ref:`torchreid_optim` for more details.
+
+
+Do two-stepped transfer learning
+-------------------------------------
+To prevent the pretrained layers from being damaged by harmful gradients back-propagated from randomly initialized layers, one can adopt the *two-stepped transfer learning strategy* presented in `Deep Transfer Learning for Person Re-identification <https://arxiv.org/abs/1611.05244>`_. The basic idea is to pretrain the randomly initialized layers for few epochs while keeping the base layers frozen before training all layers end-to-end.
+
+This has been implemented in ``Engine.train()`` (see :ref:`torchreid_engine`). The arguments related to this feature are ``fixbase_epoch`` and ``open_layers``. Intuitively, ``fixbase_epoch`` denotes the number of epochs to keep the base layers frozen; ``open_layers`` means which layer is open for training.
+
+For example, say you want to pretrain the classification layer named "classifier" in ResNet50 for 5 epochs before training all layers, you can do
+
+.. code-block:: python
+    
+    engine.run(
+        save_dir='log/resnet50',
+        max_epoch=60,
+        eval_freq=10,
+        print_freq=10,
+        test_only=False,
+        fixbase_epoch=5,
+        open_layers='classifier'
+    )
+    # or open_layers=['fc', 'classifier'] if there is another fc layer that
+    # is randomly initialized, like resnet50_fc512
+
+Note that ``fixbase_epoch`` is counted into ``max_epoch``. In the above example, the base network will be fixed for 5 epochs and then open for training for 55 epochs. Thus, if you want to freeze some layers throughout the training, what you can do is to set ``fixbase_epoch`` equal to ``max_epoch`` and put the layer names in ``open_layers`` which you want to train.
+
+
+Test a trained model
+----------------------
+You can load a trained model using :code:`torchreid.utils.load_pretrained_weights(model, weight_path)` and set ``test_only=True`` in ``engine.run()``.
+
+
+Fine-tune a model pre-trained on reid datasets
+-----------------------------------------------
+Use :code:`torchreid.utils.load_pretrained_weights(model, weight_path)` to load the pre-trained weights and then fine-tune on the dataset you want.
+
+
+Visualize learning curves with tensorboard
+--------------------------------------------
+The ``SummaryWriter()`` for tensorboard will be automatically initialized in ``engine.run()`` when you are training your model. Therefore, you do not need to do extra jobs. After the training is done, the ``*tf.events*`` file will be saved in ``save_dir``. Then, you just call ``tensorboard --logdir=your_save_dir`` in your terminal and visit ``http://localhost:6006/`` in a web browser. See `pytorch tensorboard <https://pytorch.org/docs/stable/tensorboard.html>`_ for further information.
+
+
+Visualize ranking results
+---------------------------
+This can be achieved by setting ``visrank`` to true in ``engine.run()``. ``visrank_topk`` determines the top-k images to be visualized (Default is ``visrank_topk=10``). Note that ``visrank`` can only be used in test mode, i.e. ``test_only=True`` in ``engine.run()``. The output will be saved under ``save_dir/visrank_DATASETNAME`` where each plot contains the top-k similar gallery images given a query. An example is shown below where red and green denote incorrect and correct matches respectively.
+
+.. image:: figures/ranking_results.jpg
+    :width: 800px
+    :align: center
+
+
+Visualize activation maps
+--------------------------
+To understand where the CNN focuses on to extract features for ReID, you can visualize the activation maps as in `OSNet <https://arxiv.org/abs/1905.00953>`_. This is implemented in ``tools/visualize_actmap.py`` (check the code for more details). An example running command is
+
+.. code-block:: shell
+    
+    python tools/visualize_actmap.py \
+    --root $DATA/reid \
+    -d market1501 \
+    -m osnet_x1_0 \
+    --weights PATH_TO_PRETRAINED_WEIGHTS \
+    --save-dir log/visactmap_osnet_x1_0_market1501
+
+The output will look like (from left to right: image, activation map, overlapped image)
+
+.. image:: figures/actmap.jpg
+    :width: 300px
+    :align: center
+
+
+.. note::
+    In order to visualize activation maps, the CNN needs to output the last convolutional feature maps at eval mode. See ``torchreid/models/osnet.py`` for example.
+
+
+Use your own dataset
+----------------------
+1. Write your own dataset class. Below is a template for image dataset. However, it can also be applied to a video dataset class, for which you simply change ``ImageDataset`` to ``VideoDataset``.
+
+.. code-block:: python
+    
+    from __future__ import absolute_import
+    from __future__ import print_function
+    from __future__ import division
+
+    import sys
+    import os
+    import os.path as osp
+
+    from torchreid.data import ImageDataset
+
+
+    class NewDataset(ImageDataset):
+        dataset_dir = 'new_dataset'
+
+        def __init__(self, root='', **kwargs):
+            self.root = osp.abspath(osp.expanduser(root))
+            self.dataset_dir = osp.join(self.root, self.dataset_dir)
+
+            # All you need to do here is to generate three lists,
+            # which are train, query and gallery.
+            # Each list contains tuples of (img_path, pid, camid),
+            # where
+            # - img_path (str): absolute path to an image.
+            # - pid (int): person ID, e.g. 0, 1.
+            # - camid (int): camera ID, e.g. 0, 1.
+            # Note that
+            # - pid and camid should be 0-based.
+            # - query and gallery should share the same pid scope (e.g.
+            #   pid=0 in query refers to the same person as pid=0 in gallery).
+            # - train, query and gallery share the same camid scope (e.g.
+            #   camid=0 in train refers to the same camera as camid=0
+            #   in query/gallery).
+            train = ...
+            query = ...
+            gallery = ...
+
+            super(NewDataset, self).__init__(train, query, gallery, **kwargs)
+
+
+2. Register your dataset.
+
+.. code-block:: python
+    
+    import torchreid
+    torchreid.data.register_image_dataset('new_dataset', NewDataset)
+
+
+3. Initialize a data manager with your dataset.
+
+.. code-block:: python
+    
+    # use your own dataset only
+    datamanager = torchreid.data.ImageDataManager(
+        root='reid-data',
+        sources='new_dataset'
+    )
+    # combine with other datasets
+    datamanager = torchreid.data.ImageDataManager(
+        root='reid-data',
+        sources=['new_dataset', 'dukemtmcreid']
+    )
+    # cross-dataset evaluation
+    datamanager = torchreid.data.ImageDataManager(
+        root='reid-data',
+        sources=['new_dataset', 'dukemtmcreid'],
+        targets='market1501' # or targets=['market1501', 'cuhk03']
+    )
+
+
+
+Design your own Engine
+------------------------
+A new Engine should be designed if you have your own loss function. The base Engine class ``torchreid.engine.Engine`` has implemented some generic methods which you can inherit to avoid re-writing. Please refer to the source code for more details. You are suggested to see how ``ImageSoftmaxEngine`` and ``ImageTripletEngine`` are constructed (also ``VideoSoftmaxEngine`` and ``VideoTripletEngine``). All you need to implement might be just a ``forward_backward()`` function.
+
+
+Use Torchreid as a feature extractor in your projects
+-------------------------------------------------------
+We have provided a simple API for feature extraction, which accepts input of various types such as a list of image paths or numpy arrays. More details can be found in the code at ``torchreid/utils/feature_extractor.py``. Here we show a simple example of how to extract features given a list of image paths.
+
+.. code-block:: python
+
+    from torchreid.utils import FeatureExtractor
+
+    extractor = FeatureExtractor(
+        model_name='osnet_x1_0',
+        model_path='a/b/c/model.pth.tar',
+        device='cuda'
+    )
+
+    image_list = [
+        'a/b/c/image001.jpg',
+        'a/b/c/image002.jpg',
+        'a/b/c/image003.jpg',
+        'a/b/c/image004.jpg',
+        'a/b/c/image005.jpg'
+    ]
+
+    features = extractor(image_list)
+    print(features.shape) # output (5, 512)

strong_sort/deep/reid/linter.sh

@@ -0,0 +1,11 @@
+echo "Running isort"
+isort -y -sp .
+echo "Done"
+
+echo "Running yapf"
+yapf -i -r -vv -e build .
+echo "Done"
+
+echo "Running flake8"
+flake8 .
+echo "Done"

strong_sort/deep/reid/projects/DML/README.md

@@ -0,0 +1,16 @@
+# Deep mutual learning
+
+This repo implements [Deep Mutual Learning (CVPR'18)](https://zpascal.net/cvpr2018/Zhang_Deep_Mutual_Learning_CVPR_2018_paper.pdf) (DML) for person re-id.
+
+We used this code in our [OSNet](https://arxiv.org/pdf/1905.00953.pdf) paper (see Supp. B). The training command to reproduce the result of "triplet + DML" (Table 12f in the paper) is
+```bash
+python main.py \
+--config-file im_osnet_x1_0_dml_256x128_amsgrad_cosine.yaml \
+--root $DATA
+```
+
+`$DATA` corresponds to the path to your dataset folder.
+
+Change `model.deploy` to `both` if you wanna enable model ensembling.
+
+If you have any questions, please raise an issue in the Issues area.

strong_sort/deep/reid/projects/DML/default_config.py

@@ -0,0 +1,207 @@
+from yacs.config import CfgNode as CN
+
+
+def get_default_config():
+    cfg = CN()
+
+    # model
+    cfg.model = CN()
+    cfg.model.name = 'resnet50'
+    cfg.model.pretrained = True # automatically load pretrained model weights if available
+    cfg.model.load_weights1 = '' # path to model-1 weights
+    cfg.model.load_weights2 = '' # path to model-2 weights
+    cfg.model.resume1 = '' # path to checkpoint for resume training
+    cfg.model.resume2 = '' # path to checkpoint for resume training
+    cfg.model.deploy = 'model1' # model1, model2 or both
+
+    # data
+    cfg.data = CN()
+    cfg.data.type = 'image'
+    cfg.data.root = 'reid-data'
+    cfg.data.sources = ['market1501']
+    cfg.data.targets = ['market1501']
+    cfg.data.workers = 4 # number of data loading workers
+    cfg.data.split_id = 0 # split index
+    cfg.data.height = 256 # image height
+    cfg.data.width = 128 # image width
+    cfg.data.combineall = False # combine train, query and gallery for training
+    cfg.data.transforms = ['random_flip'] # data augmentation
+    cfg.data.norm_mean = [0.485, 0.456, 0.406] # default is imagenet mean
+    cfg.data.norm_std = [0.229, 0.224, 0.225] # default is imagenet std
+    cfg.data.save_dir = 'log' # path to save log
+    cfg.data.load_train_targets = False
+
+    # specific datasets
+    cfg.market1501 = CN()
+    cfg.market1501.use_500k_distractors = False # add 500k distractors to the gallery set for market1501
+    cfg.cuhk03 = CN()
+    cfg.cuhk03.labeled_images = False # use labeled images, if False, use detected images
+    cfg.cuhk03.classic_split = False # use classic split by Li et al. CVPR14
+    cfg.cuhk03.use_metric_cuhk03 = False # use cuhk03's metric for evaluation
+
+    # sampler
+    cfg.sampler = CN()
+    cfg.sampler.train_sampler = 'RandomSampler'
+    cfg.sampler.num_instances = 4 # number of instances per identity for RandomIdentitySampler
+
+    # video reid setting
+    cfg.video = CN()
+    cfg.video.seq_len = 15 # number of images to sample in a tracklet
+    cfg.video.sample_method = 'evenly' # how to sample images from a tracklet
+    cfg.video.pooling_method = 'avg' # how to pool features over a tracklet
+
+    # train
+    cfg.train = CN()
+    cfg.train.optim = 'adam'
+    cfg.train.lr = 0.0003
+    cfg.train.weight_decay = 5e-4
+    cfg.train.max_epoch = 60
+    cfg.train.start_epoch = 0
+    cfg.train.batch_size = 32
+    cfg.train.fixbase_epoch = 0 # number of epochs to fix base layers
+    cfg.train.open_layers = [
+        'classifier'
+    ] # layers for training while keeping others frozen
+    cfg.train.staged_lr = False # set different lr to different layers
+    cfg.train.new_layers = ['classifier'] # newly added layers with default lr
+    cfg.train.base_lr_mult = 0.1 # learning rate multiplier for base layers
+    cfg.train.lr_scheduler = 'single_step'
+    cfg.train.stepsize = [20] # stepsize to decay learning rate
+    cfg.train.gamma = 0.1 # learning rate decay multiplier
+    cfg.train.print_freq = 20 # print frequency
+    cfg.train.seed = 1 # random seed
+
+    # optimizer
+    cfg.sgd = CN()
+    cfg.sgd.momentum = 0.9 # momentum factor for sgd and rmsprop
+    cfg.sgd.dampening = 0. # dampening for momentum
+    cfg.sgd.nesterov = False # Nesterov momentum
+    cfg.rmsprop = CN()
+    cfg.rmsprop.alpha = 0.99 # smoothing constant
+    cfg.adam = CN()
+    cfg.adam.beta1 = 0.9 # exponential decay rate for first moment
+    cfg.adam.beta2 = 0.999 # exponential decay rate for second moment
+
+    # loss
+    cfg.loss = CN()
+    cfg.loss.name = 'triplet'
+    cfg.loss.softmax = CN()
+    cfg.loss.softmax.label_smooth = True # use label smoothing regularizer
+    cfg.loss.triplet = CN()
+    cfg.loss.triplet.margin = 0.3 # distance margin
+    cfg.loss.triplet.weight_t = 1. # weight to balance hard triplet loss
+    cfg.loss.triplet.weight_x = 0. # weight to balance cross entropy loss
+    cfg.loss.dml = CN()
+    cfg.loss.dml.weight_ml = 1. # weight for mutual learning loss
+
+    # test
+    cfg.test = CN()
+    cfg.test.batch_size = 100
+    cfg.test.dist_metric = 'euclidean' # distance metric, ['euclidean', 'cosine']
+    cfg.test.normalize_feature = False # normalize feature vectors before computing distance
+    cfg.test.ranks = [1, 5, 10, 20] # cmc ranks
+    cfg.test.evaluate = False # test only
+    cfg.test.eval_freq = -1 # evaluation frequency (-1 means to only test after training)
+    cfg.test.start_eval = 0 # start to evaluate after a specific epoch
+    cfg.test.rerank = False # use person re-ranking
+    cfg.test.visrank = False # visualize ranked results (only available when cfg.test.evaluate=True)
+    cfg.test.visrank_topk = 10 # top-k ranks to visualize
+
+    return cfg
+
+
+def imagedata_kwargs(cfg):
+    return {
+        'root': cfg.data.root,
+        'sources': cfg.data.sources,
+        'targets': cfg.data.targets,
+        'height': cfg.data.height,
+        'width': cfg.data.width,
+        'transforms': cfg.data.transforms,
+        'norm_mean': cfg.data.norm_mean,
+        'norm_std': cfg.data.norm_std,
+        'use_gpu': cfg.use_gpu,
+        'split_id': cfg.data.split_id,
+        'combineall': cfg.data.combineall,
+        'load_train_targets': cfg.data.load_train_targets,
+        'batch_size_train': cfg.train.batch_size,
+        'batch_size_test': cfg.test.batch_size,
+        'workers': cfg.data.workers,
+        'num_instances': cfg.sampler.num_instances,
+        'train_sampler': cfg.sampler.train_sampler,
+        # image
+        'cuhk03_labeled': cfg.cuhk03.labeled_images,
+        'cuhk03_classic_split': cfg.cuhk03.classic_split,
+        'market1501_500k': cfg.market1501.use_500k_distractors,
+    }
+
+
+def videodata_kwargs(cfg):
+    return {
+        'root': cfg.data.root,
+        'sources': cfg.data.sources,
+        'targets': cfg.data.targets,
+        'height': cfg.data.height,
+        'width': cfg.data.width,
+        'transforms': cfg.data.transforms,
+        'norm_mean': cfg.data.norm_mean,
+        'norm_std': cfg.data.norm_std,
+        'use_gpu': cfg.use_gpu,
+        'split_id': cfg.data.split_id,
+        'combineall': cfg.data.combineall,
+        'batch_size_train': cfg.train.batch_size,
+        'batch_size_test': cfg.test.batch_size,
+        'workers': cfg.data.workers,
+        'num_instances': cfg.sampler.num_instances,
+        'train_sampler': cfg.sampler.train_sampler,
+        # video
+        'seq_len': cfg.video.seq_len,
+        'sample_method': cfg.video.sample_method
+    }
+
+
+def optimizer_kwargs(cfg):
+    return {
+        'optim': cfg.train.optim,
+        'lr': cfg.train.lr,
+        'weight_decay': cfg.train.weight_decay,
+        'momentum': cfg.sgd.momentum,
+        'sgd_dampening': cfg.sgd.dampening,
+        'sgd_nesterov': cfg.sgd.nesterov,
+        'rmsprop_alpha': cfg.rmsprop.alpha,
+        'adam_beta1': cfg.adam.beta1,
+        'adam_beta2': cfg.adam.beta2,
+        'staged_lr': cfg.train.staged_lr,
+        'new_layers': cfg.train.new_layers,
+        'base_lr_mult': cfg.train.base_lr_mult
+    }
+
+
+def lr_scheduler_kwargs(cfg):
+    return {
+        'lr_scheduler': cfg.train.lr_scheduler,
+        'stepsize': cfg.train.stepsize,
+        'gamma': cfg.train.gamma,
+        'max_epoch': cfg.train.max_epoch
+    }
+
+
+def engine_run_kwargs(cfg):
+    return {
+        'save_dir': cfg.data.save_dir,
+        'max_epoch': cfg.train.max_epoch,
+        'start_epoch': cfg.train.start_epoch,
+        'fixbase_epoch': cfg.train.fixbase_epoch,
+        'open_layers': cfg.train.open_layers,
+        'start_eval': cfg.test.start_eval,
+        'eval_freq': cfg.test.eval_freq,
+        'test_only': cfg.test.evaluate,
+        'print_freq': cfg.train.print_freq,
+        'dist_metric': cfg.test.dist_metric,
+        'normalize_feature': cfg.test.normalize_feature,
+        'visrank': cfg.test.visrank,
+        'visrank_topk': cfg.test.visrank_topk,
+        'use_metric_cuhk03': cfg.cuhk03.use_metric_cuhk03,
+        'ranks': cfg.test.ranks,
+        'rerank': cfg.test.rerank
+    }

strong_sort/deep/reid/projects/DML/dml.py

@@ -0,0 +1,149 @@
+from __future__ import division, print_function, absolute_import
+import torch
+from torch.nn import functional as F
+
+from torchreid.utils import open_all_layers, open_specified_layers
+from torchreid.engine import Engine
+from torchreid.losses import TripletLoss, CrossEntropyLoss
+
+
+class ImageDMLEngine(Engine):
+
+    def __init__(
+        self,
+        datamanager,
+        model1,
+        optimizer1,
+        scheduler1,
+        model2,
+        optimizer2,
+        scheduler2,
+        margin=0.3,
+        weight_t=0.5,
+        weight_x=1.,
+        weight_ml=1.,
+        use_gpu=True,
+        label_smooth=True,
+        deploy='model1'
+    ):
+        super(ImageDMLEngine, self).__init__(datamanager, use_gpu)
+
+        self.model1 = model1
+        self.optimizer1 = optimizer1
+        self.scheduler1 = scheduler1
+        self.register_model('model1', model1, optimizer1, scheduler1)
+
+        self.model2 = model2
+        self.optimizer2 = optimizer2
+        self.scheduler2 = scheduler2
+        self.register_model('model2', model2, optimizer2, scheduler2)
+
+        self.weight_t = weight_t
+        self.weight_x = weight_x
+        self.weight_ml = weight_ml
+
+        assert deploy in ['model1', 'model2', 'both']
+        self.deploy = deploy
+
+        self.criterion_t = TripletLoss(margin=margin)
+        self.criterion_x = CrossEntropyLoss(
+            num_classes=self.datamanager.num_train_pids,
+            use_gpu=self.use_gpu,
+            label_smooth=label_smooth
+        )
+
+    def forward_backward(self, data):
+        imgs, pids = self.parse_data_for_train(data)
+
+        if self.use_gpu:
+            imgs = imgs.cuda()
+            pids = pids.cuda()
+
+        outputs1, features1 = self.model1(imgs)
+        loss1_x = self.compute_loss(self.criterion_x, outputs1, pids)
+        loss1_t = self.compute_loss(self.criterion_t, features1, pids)
+
+        outputs2, features2 = self.model2(imgs)
+        loss2_x = self.compute_loss(self.criterion_x, outputs2, pids)
+        loss2_t = self.compute_loss(self.criterion_t, features2, pids)
+
+        loss1_ml = self.compute_kl_div(
+            outputs2.detach(), outputs1, is_logit=True
+        )
+        loss2_ml = self.compute_kl_div(
+            outputs1.detach(), outputs2, is_logit=True
+        )
+
+        loss1 = 0
+        loss1 += loss1_x * self.weight_x
+        loss1 += loss1_t * self.weight_t
+        loss1 += loss1_ml * self.weight_ml
+
+        loss2 = 0
+        loss2 += loss2_x * self.weight_x
+        loss2 += loss2_t * self.weight_t
+        loss2 += loss2_ml * self.weight_ml
+
+        self.optimizer1.zero_grad()
+        loss1.backward()
+        self.optimizer1.step()
+
+        self.optimizer2.zero_grad()
+        loss2.backward()
+        self.optimizer2.step()
+
+        loss_dict = {
+            'loss1_x': loss1_x.item(),
+            'loss1_t': loss1_t.item(),
+            'loss1_ml': loss1_ml.item(),
+            'loss2_x': loss1_x.item(),
+            'loss2_t': loss1_t.item(),
+            'loss2_ml': loss1_ml.item()
+        }
+
+        return loss_dict
+
+    @staticmethod
+    def compute_kl_div(p, q, is_logit=True):
+        if is_logit:
+            p = F.softmax(p, dim=1)
+            q = F.softmax(q, dim=1)
+        return -(p * torch.log(q + 1e-8)).sum(1).mean()
+
+    def two_stepped_transfer_learning(
+        self, epoch, fixbase_epoch, open_layers, model=None
+    ):
+        """Two stepped transfer learning.
+
+        The idea is to freeze base layers for a certain number of epochs
+        and then open all layers for training.
+
+        Reference: https://arxiv.org/abs/1611.05244
+        """
+        model1 = self.model1
+        model2 = self.model2
+
+        if (epoch + 1) <= fixbase_epoch and open_layers is not None:
+            print(
+                '* Only train {} (epoch: {}/{})'.format(
+                    open_layers, epoch + 1, fixbase_epoch
+                )
+            )
+            open_specified_layers(model1, open_layers)
+            open_specified_layers(model2, open_layers)
+        else:
+            open_all_layers(model1)
+            open_all_layers(model2)
+
+    def extract_features(self, input):
+        if self.deploy == 'model1':
+            return self.model1(input)
+
+        elif self.deploy == 'model2':
+            return self.model2(input)
+
+        else:
+            features = []
+            features.append(self.model1(input))
+            features.append(self.model2(input))
+            return torch.cat(features, 1)

strong_sort/deep/reid/projects/DML/im_osnet_x1_0_dml_256x128_amsgrad_cosine.yaml

@@ -0,0 +1,42 @@
+model:
+  name: 'osnet_x1_0'
+  pretrained: True
+  deploy: 'model1'
+
+data:
+  type: 'image'
+  sources: ['market1501']
+  targets: ['market1501']
+  height: 256
+  width: 128
+  combineall: False
+  transforms: ['random_flip', 'random_erase']
+  save_dir: 'log/osnet_x1_0_market1501_dml_cosinelr'
+
+loss:
+  name: 'triplet'
+  softmax:
+    label_smooth: True
+  triplet:
+    margin: 0.3
+    weight_t: 0.5
+    weight_x: 1.
+  dml:
+    weight_ml: 1.
+
+train:
+  optim: 'amsgrad'
+  lr: 0.0015
+  max_epoch: 250
+  batch_size: 64
+  fixbase_epoch: 10
+  open_layers: ['classifier']
+  lr_scheduler: 'cosine'
+
+test:
+  batch_size: 300
+  dist_metric: 'cosine'
+  normalize_feature: False
+  evaluate: False
+  eval_freq: -1
+  rerank: False

strong_sort/deep/reid/projects/DML/main.py

@@ -0,0 +1,166 @@
+import sys
+import copy
+import time
+import os.path as osp
+import argparse
+import torch
+import torch.nn as nn
+
+import torchreid
+from torchreid.utils import (
+    Logger, check_isfile, set_random_seed, collect_env_info,
+    resume_from_checkpoint, load_pretrained_weights, compute_model_complexity
+)
+
+from dml import ImageDMLEngine
+from default_config import (
+    imagedata_kwargs, optimizer_kwargs, engine_run_kwargs, get_default_config,
+    lr_scheduler_kwargs
+)
+
+
+def reset_config(cfg, args):
+    if args.root:
+        cfg.data.root = args.root
+    if args.sources:
+        cfg.data.sources = args.sources
+    if args.targets:
+        cfg.data.targets = args.targets
+    if args.transforms:
+        cfg.data.transforms = args.transforms
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        '--config-file', type=str, default='', help='path to config file'
+    )
+    parser.add_argument(
+        '-s',
+        '--sources',
+        type=str,
+        nargs='+',
+        help='source datasets (delimited by space)'
+    )
+    parser.add_argument(
+        '-t',
+        '--targets',
+        type=str,
+        nargs='+',
+        help='target datasets (delimited by space)'
+    )
+    parser.add_argument(
+        '--transforms', type=str, nargs='+', help='data augmentation'
+    )
+    parser.add_argument(
+        '--root', type=str, default='', help='path to data root'
+    )
+    parser.add_argument(
+        'opts',
+        default=None,
+        nargs=argparse.REMAINDER,
+        help='Modify config options using the command-line'
+    )
+    args = parser.parse_args()
+
+    cfg = get_default_config()
+    cfg.use_gpu = torch.cuda.is_available()
+    if args.config_file:
+        cfg.merge_from_file(args.config_file)
+    reset_config(cfg, args)
+    cfg.merge_from_list(args.opts)
+    set_random_seed(cfg.train.seed)
+
+    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
+    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
+    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
+
+    print('Show configuration\n{}\n'.format(cfg))
+    print('Collecting env info ...')
+    print('** System info **\n{}\n'.format(collect_env_info()))
+
+    if cfg.use_gpu:
+        torch.backends.cudnn.benchmark = True
+
+    datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))
+
+    print('Building model-1: {}'.format(cfg.model.name))
+    model1 = torchreid.models.build_model(
+        name=cfg.model.name,
+        num_classes=datamanager.num_train_pids,
+        loss=cfg.loss.name,
+        pretrained=cfg.model.pretrained,
+        use_gpu=cfg.use_gpu
+    )
+    num_params, flops = compute_model_complexity(
+        model1, (1, 3, cfg.data.height, cfg.data.width)
+    )
+    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
+
+    print('Copying model-1 to model-2')
+    model2 = copy.deepcopy(model1)
+
+    if cfg.model.load_weights1 and check_isfile(cfg.model.load_weights1):
+        load_pretrained_weights(model1, cfg.model.load_weights1)
+
+    if cfg.model.load_weights2 and check_isfile(cfg.model.load_weights2):
+        load_pretrained_weights(model2, cfg.model.load_weights2)
+
+    if cfg.use_gpu:
+        model1 = nn.DataParallel(model1).cuda()
+        model2 = nn.DataParallel(model2).cuda()
+
+    optimizer1 = torchreid.optim.build_optimizer(
+        model1, **optimizer_kwargs(cfg)
+    )
+    scheduler1 = torchreid.optim.build_lr_scheduler(
+        optimizer1, **lr_scheduler_kwargs(cfg)
+    )
+
+    optimizer2 = torchreid.optim.build_optimizer(
+        model2, **optimizer_kwargs(cfg)
+    )
+    scheduler2 = torchreid.optim.build_lr_scheduler(
+        optimizer2, **lr_scheduler_kwargs(cfg)
+    )
+
+    if cfg.model.resume1 and check_isfile(cfg.model.resume1):
+        cfg.train.start_epoch = resume_from_checkpoint(
+            cfg.model.resume1,
+            model1,
+            optimizer=optimizer1,
+            scheduler=scheduler1
+        )
+
+    if cfg.model.resume2 and check_isfile(cfg.model.resume2):
+        resume_from_checkpoint(
+            cfg.model.resume2,
+            model2,
+            optimizer=optimizer2,
+            scheduler=scheduler2
+        )
+
+    print('Building DML-engine for image-reid')
+    engine = ImageDMLEngine(
+        datamanager,
+        model1,
+        optimizer1,
+        scheduler1,
+        model2,
+        optimizer2,
+        scheduler2,
+        margin=cfg.loss.triplet.margin,
+        weight_t=cfg.loss.triplet.weight_t,
+        weight_x=cfg.loss.triplet.weight_x,
+        weight_ml=cfg.loss.dml.weight_ml,
+        use_gpu=cfg.use_gpu,
+        label_smooth=cfg.loss.softmax.label_smooth,
+        deploy=cfg.model.deploy
+    )
+    engine.run(**engine_run_kwargs(cfg))
+
+
+if __name__ == '__main__':
+    main()

strong_sort/deep/reid/projects/OSNet_AIN/README.md

@@ -0,0 +1,38 @@
+# Differentiable NAS for OSNet-AIN
+
+## Introduction
+This repository contains the neural architecture search (NAS) code (based on [Torchreid](https://arxiv.org/abs/1910.10093)) for [OSNet-AIN](https://arxiv.org/abs/1910.06827), an extension of [OSNet](https://arxiv.org/abs/1905.00953) that achieves strong performance on cross-domain person re-identification (re-ID) benchmarks (*without using any target data*). OSNet-AIN builds on the idea of using [instance normalisation](https://arxiv.org/abs/1607.08022) (IN) layers to eliminate instance-specific contrast in images for domain-generalisable representation learning. This is inspired by the [neural style transfer](https://arxiv.org/abs/1703.06868) works that use IN to remove image styles. Though IN naturally suits the cross-domain person re-ID task, it still remains unclear that where to insert IN to a re-ID CNN can maximise the performance gain. To avoid exhaustively evaluating all possible designs, OSNet-AIN learns to search for the optimal OSNet+IN design from data using a differentiable NAS algorithm. For technical details, please refer to our paper at https://arxiv.org/abs/1910.06827.
+
+<div align="center">
+  <img src="https://drive.google.com/uc?export=view&id=1yvVIi2Ml7WBe85Uhaa54qyG4g8z-MGEB" width="500px" />
+</div>
+
+## Training
+Assume the reid data is stored at `$DATA`. Run
+```
+python main.py --config-file nas.yaml --root $DATA
+```
+
+The structure of the found architecture will be shown at the end of training.
+
+The default config was designed for 8 Tesla V100 32GB GPUs. You can modify the batch size based on your device memory.
+
+**Note** that the test result obtained at the end of architecture search is not meaningful (due to the stochastic sampling layers). Therefore, do not rely on the result to judge the model performance. Instead, you should construct the found architecture in `osnet_child.py` and re-train and evaluate the model on the reid datasets.
+
+## Citation
+If you find this code useful to your research, please consider citing the following papers.
+```
+@article{zhou2021osnet,
+  title={Learning Generalisable Omni-Scale Representations for Person Re-Identification},
+  author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
+  journal={TPAMI},
+  year={2021}
+}
+
+@inproceedings{zhou2019osnet,
+  title={Omni-Scale Feature Learning for Person Re-Identification},
+  author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},
+  booktitle={ICCV},
+  year={2019}
+}
+```

strong_sort/deep/reid/projects/OSNet_AIN/default_config.py

@@ -0,0 +1,210 @@
+from yacs.config import CfgNode as CN
+
+
+def get_default_config():
+    cfg = CN()
+
+    # model
+    cfg.model = CN()
+    cfg.model.name = 'resnet50'
+    cfg.model.pretrained = True # automatically load pretrained model weights if available
+    cfg.model.load_weights = '' # path to model weights
+    cfg.model.resume = '' # path to checkpoint for resume training
+
+    # NAS
+    cfg.nas = CN()
+    cfg.nas.mc_iter = 1 # Monte Carlo sampling
+    cfg.nas.init_lmda = 10. # initial lambda value
+    cfg.nas.min_lmda = 1. # minimum lambda value
+    cfg.nas.lmda_decay_step = 20 # decay step for lambda
+    cfg.nas.lmda_decay_rate = 0.5 # decay rate for lambda
+    cfg.nas.fixed_lmda = False # keep lambda unchanged
+
+    # data
+    cfg.data = CN()
+    cfg.data.type = 'image'
+    cfg.data.root = 'reid-data'
+    cfg.data.sources = ['market1501']
+    cfg.data.targets = ['market1501']
+    cfg.data.workers = 4 # number of data loading workers
+    cfg.data.split_id = 0 # split index
+    cfg.data.height = 256 # image height
+    cfg.data.width = 128 # image width
+    cfg.data.combineall = False # combine train, query and gallery for training
+    cfg.data.transforms = ['random_flip'] # data augmentation
+    cfg.data.norm_mean = [0.485, 0.456, 0.406] # default is imagenet mean
+    cfg.data.norm_std = [0.229, 0.224, 0.225] # default is imagenet std
+    cfg.data.save_dir = 'log' # path to save log
+
+    # specific datasets
+    cfg.market1501 = CN()
+    cfg.market1501.use_500k_distractors = False # add 500k distractors to the gallery set for market1501
+    cfg.cuhk03 = CN()
+    cfg.cuhk03.labeled_images = False # use labeled images, if False, use detected images
+    cfg.cuhk03.classic_split = False # use classic split by Li et al. CVPR14
+    cfg.cuhk03.use_metric_cuhk03 = False # use cuhk03's metric for evaluation
+
+    # sampler
+    cfg.sampler = CN()
+    cfg.sampler.train_sampler = 'RandomSampler'
+    cfg.sampler.num_instances = 4 # number of instances per identity for RandomIdentitySampler
+
+    # video reid setting
+    cfg.video = CN()
+    cfg.video.seq_len = 15 # number of images to sample in a tracklet
+    cfg.video.sample_method = 'evenly' # how to sample images from a tracklet
+    cfg.video.pooling_method = 'avg' # how to pool features over a tracklet
+
+    # train
+    cfg.train = CN()
+    cfg.train.optim = 'adam'
+    cfg.train.lr = 0.0003
+    cfg.train.weight_decay = 5e-4
+    cfg.train.max_epoch = 60
+    cfg.train.start_epoch = 0
+    cfg.train.batch_size = 32
+    cfg.train.fixbase_epoch = 0 # number of epochs to fix base layers
+    cfg.train.open_layers = [
+        'classifier'
+    ] # layers for training while keeping others frozen
+    cfg.train.staged_lr = False # set different lr to different layers
+    cfg.train.new_layers = ['classifier'] # newly added layers with default lr
+    cfg.train.base_lr_mult = 0.1 # learning rate multiplier for base layers
+    cfg.train.lr_scheduler = 'single_step'
+    cfg.train.stepsize = [20] # stepsize to decay learning rate
+    cfg.train.gamma = 0.1 # learning rate decay multiplier
+    cfg.train.print_freq = 20 # print frequency
+    cfg.train.seed = 1 # random seed
+
+    # optimizer
+    cfg.sgd = CN()
+    cfg.sgd.momentum = 0.9 # momentum factor for sgd and rmsprop
+    cfg.sgd.dampening = 0. # dampening for momentum
+    cfg.sgd.nesterov = False # Nesterov momentum
+    cfg.rmsprop = CN()
+    cfg.rmsprop.alpha = 0.99 # smoothing constant
+    cfg.adam = CN()
+    cfg.adam.beta1 = 0.9 # exponential decay rate for first moment
+    cfg.adam.beta2 = 0.999 # exponential decay rate for second moment
+
+    # loss
+    cfg.loss = CN()
+    cfg.loss.name = 'softmax'
+    cfg.loss.softmax = CN()
+    cfg.loss.softmax.label_smooth = True # use label smoothing regularizer
+    cfg.loss.triplet = CN()
+    cfg.loss.triplet.margin = 0.3 # distance margin
+    cfg.loss.triplet.weight_t = 1. # weight to balance hard triplet loss
+    cfg.loss.triplet.weight_x = 0. # weight to balance cross entropy loss
+
+    # test
+    cfg.test = CN()
+    cfg.test.batch_size = 100
+    cfg.test.dist_metric = 'euclidean' # distance metric, ['euclidean', 'cosine']
+    cfg.test.normalize_feature = False # normalize feature vectors before computing distance
+    cfg.test.ranks = [1, 5, 10, 20] # cmc ranks
+    cfg.test.evaluate = False # test only
+    cfg.test.eval_freq = -1 # evaluation frequency (-1 means to only test after training)
+    cfg.test.start_eval = 0 # start to evaluate after a specific epoch
+    cfg.test.rerank = False # use person re-ranking
+    cfg.test.visrank = False # visualize ranked results (only available when cfg.test.evaluate=True)
+    cfg.test.visrank_topk = 10 # top-k ranks to visualize
+    cfg.test.visactmap = False # visualize CNN activation maps
+
+    return cfg
+
+
+def imagedata_kwargs(cfg):
+    return {
+        'root': cfg.data.root,
+        'sources': cfg.data.sources,
+        'targets': cfg.data.targets,
+        'height': cfg.data.height,
+        'width': cfg.data.width,
+        'transforms': cfg.data.transforms,
+        'norm_mean': cfg.data.norm_mean,
+        'norm_std': cfg.data.norm_std,
+        'use_gpu': cfg.use_gpu,
+        'split_id': cfg.data.split_id,
+        'combineall': cfg.data.combineall,
+        'batch_size_train': cfg.train.batch_size,
+        'batch_size_test': cfg.test.batch_size,
+        'workers': cfg.data.workers,
+        'num_instances': cfg.sampler.num_instances,
+        'train_sampler': cfg.sampler.train_sampler,
+        # image
+        'cuhk03_labeled': cfg.cuhk03.labeled_images,
+        'cuhk03_classic_split': cfg.cuhk03.classic_split,
+        'market1501_500k': cfg.market1501.use_500k_distractors,
+    }
+
+
+def videodata_kwargs(cfg):
+    return {
+        'root': cfg.data.root,
+        'sources': cfg.data.sources,
+        'targets': cfg.data.targets,
+        'height': cfg.data.height,
+        'width': cfg.data.width,
+        'transforms': cfg.data.transforms,
+        'norm_mean': cfg.data.norm_mean,
+        'norm_std': cfg.data.norm_std,
+        'use_gpu': cfg.use_gpu,
+        'split_id': cfg.data.split_id,
+        'combineall': cfg.data.combineall,
+        'batch_size_train': cfg.train.batch_size,
+        'batch_size_test': cfg.test.batch_size,
+        'workers': cfg.data.workers,
+        'num_instances': cfg.sampler.num_instances,
+        'train_sampler': cfg.sampler.train_sampler,
+        # video
+        'seq_len': cfg.video.seq_len,
+        'sample_method': cfg.video.sample_method
+    }
+
+
+def optimizer_kwargs(cfg):
+    return {
+        'optim': cfg.train.optim,
+        'lr': cfg.train.lr,
+        'weight_decay': cfg.train.weight_decay,
+        'momentum': cfg.sgd.momentum,
+        'sgd_dampening': cfg.sgd.dampening,
+        'sgd_nesterov': cfg.sgd.nesterov,
+        'rmsprop_alpha': cfg.rmsprop.alpha,
+        'adam_beta1': cfg.adam.beta1,
+        'adam_beta2': cfg.adam.beta2,
+        'staged_lr': cfg.train.staged_lr,
+        'new_layers': cfg.train.new_layers,
+        'base_lr_mult': cfg.train.base_lr_mult
+    }
+
+
+def lr_scheduler_kwargs(cfg):
+    return {
+        'lr_scheduler': cfg.train.lr_scheduler,
+        'stepsize': cfg.train.stepsize,
+        'gamma': cfg.train.gamma,
+        'max_epoch': cfg.train.max_epoch
+    }
+
+
+def engine_run_kwargs(cfg):
+    return {
+        'save_dir': cfg.data.save_dir,
+        'max_epoch': cfg.train.max_epoch,
+        'start_epoch': cfg.train.start_epoch,
+        'fixbase_epoch': cfg.train.fixbase_epoch,
+        'open_layers': cfg.train.open_layers,
+        'start_eval': cfg.test.start_eval,
+        'eval_freq': cfg.test.eval_freq,
+        'test_only': cfg.test.evaluate,
+        'print_freq': cfg.train.print_freq,
+        'dist_metric': cfg.test.dist_metric,
+        'normalize_feature': cfg.test.normalize_feature,
+        'visrank': cfg.test.visrank,
+        'visrank_topk': cfg.test.visrank_topk,
+        'use_metric_cuhk03': cfg.cuhk03.use_metric_cuhk03,
+        'ranks': cfg.test.ranks,
+        'rerank': cfg.test.rerank
+    }

strong_sort/deep/reid/projects/OSNet_AIN/main.py

@@ -0,0 +1,145 @@
+import os
+import sys
+import time
+import os.path as osp
+import argparse
+import torch
+import torch.nn as nn
+
+import torchreid
+from torchreid.utils import (
+    Logger, check_isfile, set_random_seed, collect_env_info,
+    resume_from_checkpoint, compute_model_complexity
+)
+
+import osnet_search as osnet_models
+from softmax_nas import ImageSoftmaxNASEngine
+from default_config import (
+    imagedata_kwargs, optimizer_kwargs, engine_run_kwargs, get_default_config,
+    lr_scheduler_kwargs
+)
+
+
+def reset_config(cfg, args):
+    if args.root:
+        cfg.data.root = args.root
+    if args.sources:
+        cfg.data.sources = args.sources
+    if args.targets:
+        cfg.data.targets = args.targets
+    if args.transforms:
+        cfg.data.transforms = args.transforms
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        '--config-file', type=str, default='', help='path to config file'
+    )
+    parser.add_argument(
+        '-s',
+        '--sources',
+        type=str,
+        nargs='+',
+        help='source datasets (delimited by space)'
+    )
+    parser.add_argument(
+        '-t',
+        '--targets',
+        type=str,
+        nargs='+',
+        help='target datasets (delimited by space)'
+    )
+    parser.add_argument(
+        '--transforms', type=str, nargs='+', help='data augmentation'
+    )
+    parser.add_argument(
+        '--root', type=str, default='', help='path to data root'
+    )
+    parser.add_argument(
+        '--gpu-devices',
+        type=str,
+        default='',
+    )
+    parser.add_argument(
+        'opts',
+        default=None,
+        nargs=argparse.REMAINDER,
+        help='Modify config options using the command-line'
+    )
+    args = parser.parse_args()
+
+    cfg = get_default_config()
+    cfg.use_gpu = torch.cuda.is_available()
+    if args.config_file:
+        cfg.merge_from_file(args.config_file)
+    reset_config(cfg, args)
+    cfg.merge_from_list(args.opts)
+    set_random_seed(cfg.train.seed)
+
+    if cfg.use_gpu and args.gpu_devices:
+        # if gpu_devices is not specified, all available gpus will be used
+        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
+    log_name = 'test.log' if cfg.test.evaluate else 'train.log'
+    log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
+    sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
+
+    print('Show configuration\n{}\n'.format(cfg))
+    print('Collecting env info ...')
+    print('** System info **\n{}\n'.format(collect_env_info()))
+
+    if cfg.use_gpu:
+        torch.backends.cudnn.benchmark = True
+
+    datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))
+
+    print('Building model: {}'.format(cfg.model.name))
+    model = osnet_models.build_model(
+        cfg.model.name, num_classes=datamanager.num_train_pids
+    )
+    num_params, flops = compute_model_complexity(
+        model, (1, 3, cfg.data.height, cfg.data.width)
+    )
+    print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
+
+    if cfg.use_gpu:
+        model = nn.DataParallel(model).cuda()
+
+    optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
+    scheduler = torchreid.optim.build_lr_scheduler(
+        optimizer, **lr_scheduler_kwargs(cfg)
+    )
+
+    if cfg.model.resume and check_isfile(cfg.model.resume):
+        cfg.train.start_epoch = resume_from_checkpoint(
+            cfg.model.resume, model, optimizer=optimizer
+        )
+
+    print('Building NAS engine')
+    engine = ImageSoftmaxNASEngine(
+        datamanager,
+        model,
+        optimizer,
+        scheduler=scheduler,
+        use_gpu=cfg.use_gpu,
+        label_smooth=cfg.loss.softmax.label_smooth,
+        mc_iter=cfg.nas.mc_iter,
+        init_lmda=cfg.nas.init_lmda,
+        min_lmda=cfg.nas.min_lmda,
+        lmda_decay_step=cfg.nas.lmda_decay_step,
+        lmda_decay_rate=cfg.nas.lmda_decay_rate,
+        fixed_lmda=cfg.nas.fixed_lmda
+    )
+    engine.run(**engine_run_kwargs(cfg))
+
+    print('*** Display the found architecture ***')
+    if cfg.use_gpu:
+        model.module.build_child_graph()
+    else:
+        model.build_child_graph()
+
+
+if __name__ == '__main__':
+    main()