Init Upload

.dockerignore

@@ -0,0 +1,79 @@
+.git
+*.Dockerfile
+.DS_Store
+.gitignore
+.dockerignore
+
+/credentials
+/cache
+/store
+
+
+# https://raw.githubusercontent.com/github/gitignore/master/Python.gitignore
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# 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
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+
+# ignore all markdown files (md) beside all README*.md other than README-secret.md
+*.md 
+*.egg-info/
+thunder-collection*.json
+*.code-workspace

.gitignore

@@ -0,0 +1,183 @@
+traffic_monitoring/detection/vis_img/
+traffic_monitoring/detection/vis_video/
+web/output/
+
+# 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/
+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
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# 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/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+*.db
+*.avi
+vis/
+cache/
+uploads/
+web/static
+.flaskenv
+mprofile_*
+migrations/
+core.*
+ckpts/
+.devcontainer/
+.vscode/
+*.pth
+*.zip
+*.trt 
+*.engine
+timing.cache
+timing.cache.lock

Dockerfile

@@ -0,0 +1,131 @@
+### This Dockerfile is modified from MMDeploy to build MMDeploy for GPU devices
+### We update the tensorrt version and cuda version to 11.8
+FROM nvcr.io/nvidia/tensorrt:23.08-py3
+
+ARG CUDA=11.8
+ARG PYTHON_VERSION=3.10
+ARG TORCH_VERSION=2.0.0
+ARG TORCHVISION_VERSION=0.15.1
+ARG ONNXRUNTIME_VERSION=1.15.1
+ARG PPLCV_VERSION=0.7.0
+ENV FORCE_CUDA="1"
+ARG MMCV_VERSION="==2.0.0"
+ARG MMENGINE_VERSION="==0.8.4"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+### change the system source for installing libs
+ARG USE_SRC_INSIDE=false
+RUN if [ ${USE_SRC_INSIDE} == true ] ; \
+    then \
+        sed -i s/archive.ubuntu.com/mirrors.aliyun.com/g /etc/apt/sources.list ; \
+        sed -i s/security.ubuntu.com/mirrors.aliyun.com/g /etc/apt/sources.list ; \
+        echo "Use aliyun source for installing libs" ; \
+    else \
+        echo "Keep the download source unchanged" ; \
+    fi
+
+### update apt and install libs
+RUN apt-get update &&\
+    apt-get install -y vim libsm6 libxext6 libxrender-dev libgl1-mesa-glx git wget libssl-dev libopencv-dev libspdlog-dev --no-install-recommends &&\
+    rm -rf /var/lib/apt/lists/*
+
+RUN curl -fsSL -v -o ~/miniconda.sh -O  https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh  && \
+    chmod +x ~/miniconda.sh && \
+    bash ~/miniconda.sh -b -p /opt/conda && \
+    rm ~/miniconda.sh && \
+    /opt/conda/bin/conda install -y python=${PYTHON_VERSION} conda-build pyyaml numpy ipython cython typing typing_extensions mkl mkl-include ninja && \
+    /opt/conda/bin/conda clean -ya
+
+### change the pip source for installing packages
+RUN if [ ${USE_SRC_INSIDE} == true ] ; \
+    then \
+        /opt/conda/bin/pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple; \
+        echo "pip using tsinghua source" ; \
+    else \
+        echo "Keep pip the download source unchanged" ; \
+    fi
+
+### install pytorch openmim
+RUN /opt/conda/bin/conda install pytorch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} cudatoolkit=${CUDA} -c pytorch -c conda-forge -y \
+    && /opt/conda/bin/pip install --no-cache-dir openmim
+
+### pytorch mmcv onnxruntime
+RUN /opt/conda/bin/mim install --no-cache-dir "mmcv"${MMCV_VERSION} onnxruntime-gpu==${ONNXRUNTIME_VERSION} mmengine${MMENGINE_VERSION}
+
+ENV PATH /opt/conda/bin:$PATH
+WORKDIR /root/workspace
+
+### get onnxruntime
+RUN wget https://github.com/microsoft/onnxruntime/releases/download/v${ONNXRUNTIME_VERSION}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz \
+    && tar -zxvf onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz
+
+### cp trt from pip to conda
+RUN cp -r /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorrt* /opt/conda/lib/python${PYTHON_VERSION}/site-packages/
+
+### install mmdeploy
+ENV ONNXRUNTIME_DIR=/root/workspace/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}
+ENV TENSORRT_DIR=/workspace/tensorrt
+ARG VERSION
+RUN git clone -b main https://github.com/open-mmlab/mmdeploy &&\
+    cd mmdeploy &&\
+    if [ -z ${VERSION} ] ; then echo "No MMDeploy version passed in, building on main" ; else git checkout tags/v${VERSION} -b tag_v${VERSION} ; fi &&\
+    git submodule update --init --recursive &&\
+    mkdir -p build &&\
+    cd build &&\
+    cmake -DMMDEPLOY_TARGET_BACKENDS="ort;trt" .. &&\
+    make -j$(nproc) &&\
+    cd .. &&\
+    /opt/conda/bin/mim install -e .
+
+### build sdk
+# RUN git clone https://github.com/openppl-public/ppl.cv.git &&\
+#     cd ppl.cv &&\
+#     git checkout tags/v${PPLCV_VERSION} -b v${PPLCV_VERSION} &&\
+#     ./build.sh cuda
+
+ENV BACKUP_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
+ENV LD_LIBRARY_PATH=/usr/local/cuda/compat/lib.real/:$LD_LIBRARY_PATH
+
+RUN cd /root/workspace/mmdeploy &&\
+    rm -rf build/CM* build/cmake-install.cmake build/Makefile build/csrc &&\
+    mkdir -p build && cd build &&\
+    cmake .. \
+        -DMMDEPLOY_BUILD_EXAMPLES=ON \
+        -DCMAKE_CXX_COMPILER=g++ \
+        -DTENSORRT_DIR=${TENSORRT_DIR} \
+        -DONNXRUNTIME_DIR=${ONNXRUNTIME_DIR} \
+        -DMMDEPLOY_BUILD_SDK_PYTHON_API=ON \
+        -DMMDEPLOY_TARGET_DEVICES="cuda;cpu" \
+        -DMMDEPLOY_TARGET_BACKENDS="ort;trt" \
+        -DMMDEPLOY_CODEBASES=all &&\
+    make -j$(nproc) && make install &&\
+    export SPDLOG_LEVEL=warn &&\
+    if [ -z ${VERSION} ] ; then echo "Built MMDeploy for GPU devices successfully!" ; else echo "Built MMDeploy version v${VERSION} for GPU devices successfully!" ; fi
+        # -DMMDEPLOY_BUILD_SDK=ON \
+        # -Dpplcv_DIR=/root/workspace/ppl.cv/cuda-build/install/lib/cmake/ppl \
+
+ENV LD_LIBRARY_PATH="/root/workspace/mmdeploy/build/lib:${BACKUP_LD_LIBRARY_PATH}"
+ENV CUDA_HOME /usr/local/cuda-11.8/
+WORKDIR /root/workspace/cc-demo
+
+COPY . /root/workspace/cc-demo
+RUN mkdir -p /data && mv ./data /data/human_detection
+RUN rm -rf /var/lib/apt/lists/*
+RUN apt-get update -y
+RUN apt-get install ffmpeg libsm6 libxext6  -y
+RUN apt-get clean -y
+RUN pip install --upgrade pip
+
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+RUN mim install mmdet==3.1.0
+RUN mim install mmpose==1.1.0
+RUN mim install mmyolo==0.6.0
+RUN pip install cython-bbox==0.1.3
+RUN pip install lap
+RUN pip install gradio
+RUN python setup.py develop 
+
+RUN echo 'export PYTHONPATH=$PYTHONPATH:./' >> ~/.bashrc
+
+CMD /bin/bash -c "cd projects && gradio human_detection/demo_app.py"

models/base/__init__.py

@@ -0,0 +1,2 @@
+from .trt_base import TRT_Base
+from .onnx_base import ONNX_Base

models/base/onnx_base.py

@@ -0,0 +1,58 @@
+# Inference for onnx model (.onnx)
+from typing import List
+import numpy as np
+import onnxruntime as ort
+import os, torch
+
+class ONNX_Base():
+    def __init__(self,
+                 model_path: str, 
+                 device: str='0'):
+        self.model_path = model_path
+        self.device = self.select_device(device)
+        self.session = self.create_session(model_path)
+        
+    def create_session(self, model_path):
+
+        providers = ['CPUExecutionProvider']
+        if torch.cuda.is_available():
+            providers.insert(0, 'CUDAExecutionProvider')
+        ort_session = ort.InferenceSession(model_path, providers=providers)
+        return ort_session
+    
+    def select_device(self, device: str)->torch.device:
+        """ Select device to be used for inference.
+        Args:
+            param device: 'cpu' or '0' or '0,1,2,3'
+        Return: 
+            torch.device
+        """
+        cpu = device.lower() == "cpu"
+        if cpu:
+            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+            return torch.device("cpu")
+        else:
+            assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'
+            os.environ['CUDA_VISBILE_DEVICES'] = device
+            torch.cuda.set_device(int(device))
+            return torch.device(f"cuda:{device}")
+
+    def infer_batch(self, image_batch: np.ndarray) -> List[np.ndarray]:
+        """ Inference for onnx model.
+        Args:
+            param image_batch: (batch_size, height, width, channels)
+        Return:
+            results: List[np.ndarray]
+        """
+        input_name = self.session.get_inputs()[0].name
+        results = self.session.run(None, {input_name: image_batch})
+        return results
+
+    
+
+    
+
+
+
+
+    

models/base/trt_base.py

@@ -0,0 +1,97 @@
+
+from typing import Tuple, Dict,List
+import tensorrt as trt
+import numpy as np
+import torch
+import time
+import os
+from collections import OrderedDict, namedtuple
+
+class TRT_Base():
+    def __init__(self, 
+                 input_shape: Tuple[int, int, int], 
+                 model_path: str,
+                 device: str='0'):
+        """ Tensor RT base class for inference.
+
+        Args:
+            input_shape (Tuple[int, int, int]): image size (3, H, W)
+            model_path (str): path to the model.trt
+            device (str, optional): CUDA device. Defaults to '0'.
+        """
+        self.input_shape = input_shape
+        self.model_path = model_path
+        self.device = self.select_device(device)
+        self.init_model()
+
+    def select_device(self, device: str)->torch.device:
+        """ Select device to be used for inference.
+        Args:
+            param device: 'cpu' or '0' or '0,1,2,3'
+        Return: 
+            torch.device
+        """
+        cpu = device.lower() == "cpu"
+        if cpu:
+            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+            return torch.device("cpu")
+        else:
+            assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'
+            torch.cuda.set_device(int(device))
+            return torch.device("cuda")
+
+    def init_model(self):
+        """ Initialize TensorRT engine and context."""
+        logger = trt.Logger(trt.Logger.INFO)
+        trt.init_libnvinfer_plugins(logger, namespace="")
+        with open(self.model_path, 'rb') as f, trt.Runtime(logger) as runtime:
+            engine = runtime.deserialize_cuda_engine(f.read())
+        context = engine.create_execution_context()
+        self.model = {
+            "engine": engine,
+            "context": context
+        }
+        bindings, binding_addrs = self.get_bindings(input_shape=self.input_shape)
+        input_names = [binding_name for binding_name in binding_addrs.keys() if (self.model["engine"].binding_is_input(binding_name))]
+        
+        for _ in range(10):
+            for name in input_names:
+                binding_addrs[name] = int(torch.randn(bindings[name].shape).to(self.device).data_ptr())
+            context.execute_v2(list(binding_addrs.values()))
+        self.model.update({
+            'binding_addrs': binding_addrs,
+            'bindings': bindings,
+            'rt_shapes': self.input_shape
+        })
+
+    def get_bindings(self, input_shape: Tuple[int, int, int]):
+        """ Get bindings and binding addresses for TensorRT engine.
+        Args: 
+            input_shape (Tuple[int, int, int]): image size (3, H, W)
+        """
+        
+        self.model["context"].set_binding_shape(0, input_shape)
+        bindings = OrderedDict()
+        Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
+        for index in range(self.model["engine"].num_bindings):
+            name = self.model["engine"].get_binding_name(index)
+            dtype = trt.nptype(self.model["engine"].get_binding_dtype(index))
+            shape = tuple(self.model["context"].get_binding_shape(index))
+            data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(self.device)
+            bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
+        binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
+        return bindings, binding_addrs
+
+    def change_runtime_dimension(self, input_shape: Tuple[int, int, int]):
+        """ Support inference with Dynamic shape.
+
+        Args:
+            input_shape (Tuple[int, int, int]): image size (3, H, W)
+        """
+        if (input_shape == self.model["rt_shapes"]): return
+        bindings, binding_addrs = self.get_bindings(input_shape)
+        self.model['binding_addrs'] = binding_addrs
+        self.model['bindings'] = bindings
+        self.model['rt_shapes'] = input_shape
+
+    

models/detectors/yolov7.py

@@ -0,0 +1,228 @@
+from typing import List, Dict, Tuple, Optional, Union
+import torch
+import time
+from models.base.trt_base import TRT_Base
+from models.base.onnx_base import ONNX_Base
+import cv2 
+import numpy as np 
+from mmcv.ops.nms import batched_nms
+
+class YOLOv7Base():
+    def __init__(self,
+                    class_map_ids: Optional[Dict]=None,
+                    preprocess_cfg: Dict=dict(                      
+                        border_color=(114, 114, 114), 
+                        auto=False, 
+                        scaleFill=True, 
+                        scaleup=True, 
+                        stride=32),
+                    nms_agnostic_cfg: Dict=dict(
+                        type='nms',
+                        iou_threshold=0.9,
+                        class_agnostic=True),   
+                    use_torch: bool=False,
+                    ):
+        """ YOLOv7 class for inference.
+
+        Args:
+            class_map_ids (Dict): class mapping dictionary to map model's class to original class. Eg {0: 1, 1: 0, 2: 2, 3: 3} mean we swap class ID between 0 and 1.
+            preprocess_cfg (Dict):
+                - border_color (Tuple[int, int, int]): padding value.
+                - auto (bool): resize input to minimum rectangle.
+                - scaleFill (bool): stretching the input image.
+                - scaleUp (bool): allow scale up the input image.
+                - stride (int): stride of the model.
+            nms_agnostic_cfg (Dict):
+                - type (str): nms type (nms, softnms).
+                - iou_threshold (float): IoU threshold of NMS.
+                - class_agnostic (bool): enable class-agnostic NMS instead of NMS for each class.
+            use_torch (bool): use torch tensor or numpy array in preprocess and postprocess function.
+        """
+        self.preprocess_cfg=preprocess_cfg
+        self.nms_agnostic_cfg=nms_agnostic_cfg
+        self.class_map_ids = class_map_ids
+        self.use_torch= use_torch
+        
+    def letterbox(self, 
+                  img: np.ndarray, 
+                  new_shape: Tuple[int, int]=(640, 640), 
+                  border_color: Tuple[int, int, int]=(114, 114, 114), 
+                  auto: bool=True, 
+                  scaleFill: bool=False, 
+                  scaleup: bool=True, 
+                  stride: int=32):
+        """ Resize input image.
+
+        Args:
+            img (np.ndarray): input image.
+            new_shape (Tuple[int, int]): the shape of output image.
+            border_color: Tuple[int, int, int]: padding value.
+            auto (bool): resize input to minimum rectangle.
+            scaleFill (bool): stretching the input image.
+            scaleUp (bool): allow scale up the input image.
+            stride (int): stride of the model.
+        """
+
+        # Resize and pad image while meeting stride-multiple constraints
+        shape = img.shape[:2]  # current shape [height, width]
+        if isinstance(new_shape, int):
+            new_shape = (new_shape, new_shape)
+
+        # Scale ratio (new / old)
+        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+        if not scaleup:  # only scale down, do not scale up (for better test mAP)
+            r = min(r, 1.0)
+
+        # Compute padding
+        ratio = r, r  # width, height ratios
+        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
+        if auto:  # minimum rectangle
+            dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
+        elif scaleFill:  # stretch
+            dw, dh = 0.0, 0.0
+            new_unpad = (new_shape[1], new_shape[0])
+            ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
+
+        dw /= 2  # divide padding into 2 sides
+        dh /= 2
+
+        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+        
+        if shape[::-1] != new_unpad:  # resize
+            new_img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
+        else:
+            new_img = img
+        new_img = cv2.copyMakeBorder(new_img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=border_color)  # add border
+        return new_img, ratio, (dw, dh)
+    
+    def preprocess(self, input_data: np.ndarray):
+        """ Preprocess function for input data.
+
+        Args:
+            input_data (np.ndarray): batch input image.
+        """
+        tensor_data = []
+        if ((isinstance(input_data, np.ndarray)) and (len(input_data.shape) == 3)):
+            input_data = [input_data]
+        for i in range(len(input_data)):
+            img = input_data[i]
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            preprocessed_img, ratio, dwdh = self.letterbox(img, new_shape=self.input_shape[2:], **self.preprocess_cfg)
+            height, width = preprocessed_img.shape[0], preprocessed_img.shape[1]
+            if self.use_torch:
+                tensor_data.append(torch.from_numpy(preprocessed_img).to(self.device))
+            else:
+                tensor_data.append(preprocessed_img)
+        if self.use_torch:
+            tensor_data = torch.stack(tensor_data, dim=0)[:, :, :, [2, 1, 0]].permute(0, 3, 1, 2).float().contiguous()/255.0
+        else:
+            tensor_data = np.stack(tensor_data, axis=0)/255.0
+            
+        return tensor_data, height, width, ratio[0], dwdh
+
+    def postprocess(self, 
+                    boxes: Union[torch.tensor,np.ndarray], 
+                    r: float, 
+                    dwdh: Tuple[float, float]):
+        """ Postprocess function for input data.
+
+        Args:
+            boxes (torch.tensor or np.ndarray): output boxes.
+            r (float): ratio between model shape and input shape.
+            dwdh (Tuple[float, float]): top-left padding position.
+        """
+        dwdh = dwdh * 2
+        if self.use_torch:
+            dwdh = torch.tensor(dwdh)
+        boxes -= dwdh
+        boxes /= r
+        return boxes
+    
+class YOLOv7TRT(TRT_Base, YOLOv7Base):
+    def __init__(self, 
+                    class_map_ids: Optional[Dict],
+                    preprocess_cfg: Dict,
+                    nms_agnostic_cfg: Dict,
+                    img_shape: Tuple[int, int]=(640, 640),
+                    batch_size: int=32,
+                    model_path: str="",
+                    device: str='0',):
+        """ YOLOv7 TRT class for inference, which is based on TRT_Base and YOLOv7Base.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        YOLOv7Base.__init__(self, class_map_ids=class_map_ids, 
+                                         preprocess_cfg=preprocess_cfg, 
+                                         nms_agnostic_cfg=nms_agnostic_cfg, 
+                                         use_torch=True)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> List[np.ndarray]:
+        """ Batch inference function for batch input image.
+
+        Args:
+            image_batch (np.ndarray): batch of input image.
+        """
+
+        tensor_data, height, width, ratio, dwdh = self.preprocess(image_batch)
+        self.change_runtime_dimension(input_shape=(len(tensor_data), 3, height, width))
+        self.model['binding_addrs']['images'] = int(tensor_data.data_ptr())
+        self.model['context'].execute_v2(list(self.model['binding_addrs'].values()))
+        nums = self.model['bindings']['num_dets'].data.cpu()
+        boxes = self.model['bindings']['det_boxes'].data.cpu()
+        scores = self.model['bindings']['det_scores'].data.cpu()
+        classes = self.model['bindings']['det_classes'].data.cpu()
+        
+        # Rearrange the classes idx.
+        new_classes = classes.clone()
+        if self.class_map_ids is not None:
+            for idx_s, idx_t in self.class_map_ids.items():
+                new_classes[classes == idx_s] = idx_t
+        boxes = self.postprocess(boxes, ratio, dwdh)
+        det_outputs = []
+        for idx in range(len(nums)):
+            num = nums[idx, 0]
+            frame_boxes = boxes[idx, :num]
+            frame_scores = scores[idx, :num]
+            frame_labels = new_classes[idx, :num]
+            result_boxes, keep = batched_nms(frame_boxes.float(), frame_scores.float(), frame_labels, nms_cfg=self.nms_agnostic_cfg)
+            det_outputs.append({"boxes": result_boxes,"labels": frame_labels[keep]})
+        return det_outputs 
+    
+class YOLOv7ONNX(ONNX_Base, YOLOv7Base):
+    def __init__(self, 
+                    class_map_ids,
+                    preprocess_cfg,
+                    nms_agnostic_cfg,
+                    model_path: str="",
+                    device: str='0',):
+        """ YOLOv7 ONNX class for inference, which is based on ONNX_Base and YOLOv7Base.
+        """
+        super().__init__(model_path,device)
+        YOLOv7Base.__init__(self, class_map_ids=class_map_ids, 
+                            preprocess_cfg=preprocess_cfg, 
+                            nms_agnostic_cfg=nms_agnostic_cfg, 
+                            use_torch=False)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> List[np.ndarray]:
+        numpy_array_data, height, width, ratio, dwdh = self.preprocess(image_batch)
+        results = super().infer_batch(numpy_array_data)
+        
+        det_outputs = []
+        batch_size = len(results[0])
+        for idx in range(batch_size):
+            boxes = results[0][idx][1:5]
+            classes = results[0][idx][5]
+            scores = results[0][idx][-1]
+            boxes = self.postprocess(boxes, ratio, dwdh)
+
+            result_boxes, keep = batched_nms(boxes.float(), scores.float(), classes, nms_cfg=self.nms_agnostic_cfg)
+            det_outputs.append({"boxes": result_boxes,"labels": classes[keep]})
+        return det_outputs 
+    
+
+    
+    

models/models/base/onnx_base.py

@@ -0,0 +1,68 @@
+# Inference for onnx model (.onnx)
+from typing import List
+import numpy as np
+import onnxruntime as ort
+import os, torch
+
+class ONNX_Base():
+    def __init__(self,
+                 input_shape,
+                 model_path: str, 
+                 device: str='0'):
+        self.input_shape = input_shape
+        self.model_path = model_path
+        self.device = self.select_device(device)
+        self.session = self.create_session(model_path)
+        
+    def create_session(self, model_path: str) -> ort.InferenceSession:
+        """_summary_
+
+        Args:
+            model_path (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        
+        providers = ['CPUExecutionProvider']
+        if torch.cuda.is_available():
+            providers.insert(0, 'CUDAExecutionProvider')
+        ort_session = ort.InferenceSession(model_path, providers=providers)
+        return ort_session
+    
+    def select_device(self, device: str)->torch.device:
+        """ Select device to be used for inference.
+        Args:
+            param device: 'cpu' or '0' or '0,1,2,3'
+        Return: 
+            torch.device
+        """
+        cpu = device.lower() == "cpu"
+        if cpu:
+            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+            return torch.device("cpu")
+        else:
+            assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'
+            os.environ['CUDA_VISBILE_DEVICES'] = device
+            torch.cuda.set_device(int(device))
+            return torch.device(f"cuda:{device}")
+
+    def infer_batch(self, image_batch: np.ndarray) -> List[np.ndarray]:
+        """ Inference for onnx model.
+        Args:
+            param image_batch: (batch_size, height, width, channels)
+        Return:
+            results: List[np.ndarray]
+        """
+        input_name = self.session.get_inputs()[0].name
+        results = self.session.run(None, {input_name: image_batch})
+        return results
+
+    
+
+    
+
+
+
+
+    

models/models/base/trt_base.py

@@ -0,0 +1,96 @@
+
+from typing import Tuple, Dict,List
+import tensorrt as trt
+import numpy as np
+import torch
+import time
+import os
+from collections import OrderedDict, namedtuple
+
+class TRT_Base():
+    def __init__(self, 
+                 input_shape: Tuple[int, int, int], 
+                 model_path: str,
+                 device: str='0'):
+        """ Tensor RT base class for inference.
+
+        Args:
+            input_shape (Tuple[int, int, int]): image size (3, H, W)
+            model_path (str): path to the model.trt
+            device (str, optional): CUDA device. Defaults to '0'.
+        """
+        self.input_shape = input_shape
+        self.model_path = model_path
+        self.device = self.select_device(device)
+        self.init_model()
+
+    def select_device(self, device: str)->torch.device:
+        """ Select device to be used for inference.
+        Args:
+            param device: 'cpu' or '0' or '0,1,2,3'
+        Return: 
+            torch.device
+        """
+        cpu = device.lower() == "cpu"
+        if cpu:
+            os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+            return torch.device("cpu")
+        else:
+            assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'
+            torch.cuda.set_device(int(device))
+            return torch.device("cuda")
+
+    def init_model(self):
+        """ Initialize TensorRT engine and context."""
+        logger = trt.Logger(trt.Logger.INFO)
+        trt.init_libnvinfer_plugins(logger, namespace="")
+        with open(self.model_path, 'rb') as f, trt.Runtime(logger) as runtime:
+            engine = runtime.deserialize_cuda_engine(f.read())
+        context = engine.create_execution_context()
+        self.model = {
+            "engine": engine,
+            "context": context
+        }
+        bindings, binding_addrs = self.get_bindings(input_shape=self.input_shape)
+        input_names = [binding_name for binding_name in binding_addrs.keys() if (self.model["engine"].binding_is_input(binding_name))]
+        
+        for _ in range(10):
+            for name in input_names:
+                binding_addrs[name] = int(torch.randn(bindings[name].shape).to(self.device).data_ptr())
+            context.execute_v2(list(binding_addrs.values()))
+        self.model.update({
+            'binding_addrs': binding_addrs,
+            'bindings': bindings,
+            'rt_shapes': self.input_shape
+        })
+
+    def get_bindings(self, input_shape: Tuple[int, int, int]):
+        """ Get bindings and binding addresses for TensorRT engine.
+        Args: 
+            input_shape (Tuple[int, int, int]): image size (3, H, W)
+        """
+        self.model["context"].set_binding_shape(0, input_shape)
+        bindings = OrderedDict()
+        Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
+        for index in range(self.model["engine"].num_bindings):
+            name = self.model["engine"].get_binding_name(index)
+            dtype = trt.nptype(self.model["engine"].get_binding_dtype(index))
+            shape = tuple(self.model["context"].get_binding_shape(index))
+            data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(self.device)
+            bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
+        binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
+        return bindings, binding_addrs
+
+    def change_runtime_dimension(self, input_shape: Tuple[int, int, int]):
+        """ Support inference with Dynamic shape.
+
+        Args:
+            input_shape (Tuple[int, int, int]): image size (3, H, W)
+        """
+        if (input_shape == self.model["rt_shapes"]): return
+        bindings, binding_addrs = self.get_bindings(input_shape)
+        self.model['binding_addrs'] = binding_addrs
+        self.model['bindings'] = bindings
+        self.model['rt_shapes'] = input_shape
+
+    

models/models/detectors/mmyolov8.py

@@ -0,0 +1,86 @@
+from typing import List, Dict, Tuple
+import numpy as np 
+import torch
+from .yolov7 import YOLOBase
+from models.base.trt_base import TRT_Base
+from models.base.onnx_base import ONNX_Base
+
+class MMYOLOv8TRT(TRT_Base, YOLOBase):
+    def __init__(self, 
+                    preprocess_cfg: Dict=dict(                      
+                        border_color=(114, 114, 114), 
+                        auto=False, 
+                        scaleFill=True, 
+                        scaleup=True, 
+                        stride=32),
+                    nms_agnostic_cfg: Dict=dict(
+                        type='nms',
+                        iou_threshold=0.9,
+                        class_agnostic=True),   
+                    score_thr=0.1,
+                    img_shape: Tuple[int, int]=(640, 640),
+                    batch_size: int=32,
+                    model_path: str="",
+                    device: str='0',):
+        """ YOLOv8 TRT class for inference, which is based on TRT_Base and YOLOBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        YOLOBase.__init__(self, preprocess_cfg=preprocess_cfg, 
+                                  nms_agnostic_cfg=nms_agnostic_cfg, 
+                                  score_thr=score_thr,
+                                  use_torch=True)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> List[Dict]:
+        """ Batch inference function for batch input image.
+
+        Args:
+            image_batch (np.ndarray): batch of input image.
+        """
+
+        tensor_data, height, width, ratio, dwdh = self.preprocess(image_batch)
+        self.change_runtime_dimension(input_shape=(len(tensor_data), 3, height, width))
+        self.model['binding_addrs']['input'] = int(tensor_data.data_ptr())
+        self.model['context'].execute_v2(list(self.model['binding_addrs'].values()))
+        dets = self.model['bindings']['dets'].data.cpu()
+        classes = self.model['bindings']['labels'].data.cpu()
+        boxes = dets[:,:,:4]
+        scores = dets[:,:,4]
+        
+        return self.post_process(boxes, scores, classes, ratio, dwdh)
+
+class MMYOLOv8ONNX(ONNX_Base, YOLOBase):
+    def __init__(self, 
+                    preprocess_cfg,
+                    nms_agnostic_cfg,
+                    score_thr=0.1,
+                    img_shape: Tuple[int, int]=(640, 640),
+                    batch_size: int=32,
+                    model_path: str="",
+                    device: str='0',):
+        """ YOLOv7 ONNX class for inference, which is based on ONNX_Base and YOLOBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        YOLOBase.__init__(self,  
+                            preprocess_cfg=preprocess_cfg, 
+                            nms_agnostic_cfg=nms_agnostic_cfg, 
+                            score_thr=score_thr,
+                            use_torch=False)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> List[Dict]:
+
+        numpy_array_data, height, width, ratio, dwdh = self.preprocess(image_batch)
+        numpy_array_data = numpy_array_data.astype(np.float32)
+        results = super().infer_batch(numpy_array_data)
+        dets, classes = results
+        dets = torch.from_numpy(dets)
+        classes = torch.from_numpy(classes)
+        boxes = dets[:,:,:4]
+        scores = dets[:,:,4]
+        
+        return self.post_process(boxes, scores, classes, ratio, dwdh)

models/models/detectors/yolov7.py

@@ -0,0 +1,251 @@
+from typing import List, Dict, Tuple, Union
+import torch, cv2
+import numpy as np 
+
+from models.base.trt_base import TRT_Base
+from models.base.onnx_base import ONNX_Base
+from mmcv.ops.nms import batched_nms, nms
+
+class YOLOBase():
+    def __init__(self,
+                    preprocess_cfg: Dict=dict(                      
+                        border_color=(114, 114, 114), 
+                        auto=False, 
+                        scaleFill=True, 
+                        scaleup=True, 
+                        stride=32),
+                    nms_agnostic_cfg: Dict=dict(
+                        type='nms',
+                        iou_threshold=0.9,
+                        class_agnostic=True),   
+                    score_thr=0.1,
+                    use_torch: bool=False,
+                    ):
+        """ This base-class has preprocess and postprocess function for YOLO model.
+
+        Args:
+            preprocess_cfg (Dict):
+                - border_color (Tuple[int, int, int]): padding value.
+                - auto (bool): resize input to minimum rectangle.
+                - scaleFill (bool): stretching the input image.
+                - scaleUp (bool): allow scale up the input image.
+                - stride (int): stride of the model.
+            nms_agnostic_cfg (Dict):
+                - type (str): nms type (nms, softnms).
+                - iou_threshold (float): IoU threshold of NMS.
+                - class_agnostic (bool): enable class-agnostic NMS instead of NMS for each class.
+            use_torch (bool): use torch tensor or numpy array in preprocess and postprocess function.
+        """
+        self.preprocess_cfg=preprocess_cfg
+        self.nms_agnostic_cfg=nms_agnostic_cfg
+        self.use_torch= use_torch
+        self.score_thr = score_thr
+        
+    def letterbox(self, 
+                  img: np.ndarray, 
+                  new_shape: Tuple[int, int]=(640, 640), 
+                  border_color: Tuple[int, int, int]=(114, 114, 114), 
+                  auto: bool=True, 
+                  scaleFill: bool=False, 
+                  scaleup: bool=True, 
+                  stride: int=32):
+        """ Resize input image.
+
+        Args:
+            img (np.ndarray): input image.
+            new_shape (Tuple[int, int]): the shape of output image.
+            border_color: Tuple[int, int, int]: padding value.
+            auto (bool): resize input to minimum rectangle.
+            scaleFill (bool): stretching the input image.
+            scaleUp (bool): allow scale up the input image.
+            stride (int): stride of the model.
+        """
+
+        # Resize and pad image while meeting stride-multiple constraints
+        shape = img.shape[:2]  # current shape [height, width]
+        if isinstance(new_shape, int):
+            new_shape = (new_shape, new_shape)
+
+        # Scale ratio (new / old)
+        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+        if not scaleup:  # only scale down, do not scale up (for better test mAP)
+            r = min(r, 1.0)
+
+        # Compute padding
+        ratio = r, r  # width, height ratios
+        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
+        if auto:  # minimum rectangle
+            dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
+        elif scaleFill:  # stretch
+            dw, dh = 0.0, 0.0
+            new_unpad = (new_shape[1], new_shape[0])
+            ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
+
+        dw /= 2  # divide padding into 2 sides
+        dh /= 2
+
+        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+        
+        if shape[::-1] != new_unpad:  # resize
+            new_img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
+        else:
+            new_img = img
+        new_img = cv2.copyMakeBorder(new_img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=border_color)  # add border
+        return new_img, ratio, (dw, dh)
+    
+    def preprocess(self, input_data: np.ndarray):
+        """ Preprocess function for input data.
+
+        Args:
+            input_data (np.ndarray): batch input image.
+        """
+        tensor_data = []
+        if ((isinstance(input_data, np.ndarray)) and (len(input_data.shape) == 3)):
+            input_data = [input_data]
+        for i in range(len(input_data)):
+            img = input_data[i]
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            preprocessed_img, ratio, dwdh = self.letterbox(img, new_shape=self.input_shape[2:], **self.preprocess_cfg)
+            height, width = preprocessed_img.shape[0], preprocessed_img.shape[1]
+            if self.use_torch:
+                tensor_data.append(torch.from_numpy(preprocessed_img).to(self.device))
+            else:
+                tensor_data.append(preprocessed_img)
+        if self.use_torch:
+            tensor_data = torch.stack(tensor_data, dim=0).permute(0, 3, 1, 2).float().contiguous()/255.0
+        else:
+            tensor_data = np.stack(tensor_data, axis=0).transpose(0, 3, 1, 2)/255.0
+            
+        return tensor_data, height, width, ratio[0], dwdh
+
+    def scale_boxes(self, 
+                    boxes: Union[torch.Tensor,np.ndarray], 
+                    r: float, 
+                    dwdh: Tuple[float, float]) -> torch.Tensor:
+        """ Scale predicted boxes to the original image shape.
+
+        Args:
+            boxes (torch.tensor or np.ndarray): output boxes.
+            r (float): ratio between model shape and input shape.
+            dwdh (Tuple[float, float]): top-left padding position.
+        """
+        dwdh = dwdh * 2
+        if isinstance(boxes, torch.Tensor):
+            dwdh = torch.tensor(dwdh)
+        boxes -= dwdh
+        boxes /= r
+        return boxes
+    
+    def post_process(self,
+                    boxes: torch.Tensor, 
+                    scores: torch.Tensor, 
+                    classes: torch.Tensor, 
+                    ratio: float, 
+                    dwdh: Tuple[int,int]) -> List[Dict]:
+        """ Postprocess function for input data.
+
+        Args:
+            boxes (torch.Tensor): output boxes of shape [N_box, 4].
+            scores (torch.Tensor): confidence scores of shape [N_box,].
+            classes (torch.Tensor): class ids of shape [N_box,].
+            ratio (float): ratio between original image shape and model's input shape.
+            dwdh (Tuple[int,int]): padding margin output from letterbox function.
+
+        Returns:
+            List[Dict]: _description_
+        """
+        boxes = self.scale_boxes(boxes, ratio, dwdh)
+        det_outputs = []
+        batch_size = len(boxes)
+        for idx in range(batch_size):
+            frame_boxes = boxes[idx]
+            frame_scores = scores[idx]
+            frame_labels = classes[idx]
+            # Filter out confidence scores below threshold
+            index = frame_scores > self.score_thr
+            frame_boxes = frame_boxes[index]
+            frame_scores = frame_scores[index]
+            frame_labels = frame_labels[index]
+            # Use NMS to suppress class-agnostic boxes
+            result_boxes, keep = batched_nms(frame_boxes.float(), frame_scores.float(), frame_labels, nms_cfg=self.nms_agnostic_cfg)
+            det_outputs.append({"boxes": result_boxes,"labels": frame_labels[keep]})
+        return det_outputs 
+    
+class YOLOv7TRT(TRT_Base, YOLOBase):
+    def __init__(self, 
+                    preprocess_cfg: Dict,
+                    nms_agnostic_cfg: Dict,
+                    score_thr=0.1,
+                    img_shape: Tuple[int, int]=(640, 640),
+                    batch_size: int=32,
+                    model_path: str="",
+                    device: str='0',):
+        """ YOLOv7 TRT class for inference, which is based on TRT_Base and YOLOBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        YOLOBase.__init__(self, preprocess_cfg=preprocess_cfg, 
+                                  nms_agnostic_cfg=nms_agnostic_cfg, 
+                                  score_thr=score_thr,
+                                  use_torch=True)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> List[Dict]:
+        """ Batch inference function for batch input image.
+
+        Args:
+            image_batch (np.ndarray): batch of input image.
+        """
+
+        tensor_data, height, width, ratio, dwdh = self.preprocess(image_batch)
+        self.change_runtime_dimension(input_shape=(len(tensor_data), 3, height, width))
+        self.model['binding_addrs']['images'] = int(tensor_data.data_ptr())
+        self.model['context'].execute_v2(list(self.model['binding_addrs'].values()))
+        nums = self.model['bindings']['num_dets'].data.cpu()
+        boxes = self.model['bindings']['det_boxes'].data.cpu()
+        scores = self.model['bindings']['det_scores'].data.cpu()
+        classes = self.model['bindings']['det_classes'].data.cpu()
+        
+        return self.post_process(boxes, scores, classes, ratio, dwdh)
+    
+class YOLOv7ONNX(ONNX_Base, YOLOBase):
+    def __init__(self, 
+                    preprocess_cfg,
+                    nms_agnostic_cfg,
+                    score_thr=0.1,
+                    img_shape: Tuple[int, int]=(640, 640),
+                    batch_size: int=32,
+                    model_path: str="",
+                    device: str='0',):
+        """ YOLOv7 ONNX class for inference, which is based on ONNX_Base and YOLOBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        YOLOBase.__init__(self,  
+                            preprocess_cfg=preprocess_cfg, 
+                            nms_agnostic_cfg=nms_agnostic_cfg, 
+                            score_thr=score_thr,
+                            use_torch=False)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> List[Dict]:
+
+        numpy_array_data, height, width, ratio, dwdh = self.preprocess(image_batch)
+        numpy_array_data = numpy_array_data.astype(np.float32)
+        results = super().infer_batch(numpy_array_data)
+        num_dets, boxes, scores, classes = results
+        boxes = torch.from_numpy(boxes)
+        scores = torch.from_numpy(scores)
+        classes = torch.from_numpy(classes)
+        return self.post_process(boxes, scores, classes, ratio, dwdh)        
+        
+
+
+    
+
+    
+    

models/models/engine/threading_func.py

@@ -0,0 +1,221 @@
+from queue import Queue, Full, Empty
+from threading import Event
+from mmcv import VideoReader
+import logging
+from gradio import Progress
+from models.trackers.byte_track import BYTETracker
+import torch
+import numpy as np
+
+def queue_clear(q: Queue):
+    """ Clear all items in the queue.
+
+    Args:
+        q (Queue): input queue.
+    """
+    with q.mutex: q.queue.clear()
+
+def queue_get(q: Queue, eStop: Event, retry_interval=1, item_idx=None, default_item=None):
+    """wrapper for queue.get() with timeout, retry and event stop.
+
+    Args:
+        q (Queue): input queue.
+        eStop (Event): event to stop the thread.
+        retry_interval (int, optional): time to wait before retry to get the item. Defaults to 1 second.
+        item_idx (_type_, optional): index of item to get. This is used for logging information. Defaults to None.
+        default_item (_type_, optional): default item to return if error or early stop. Defaults to None.
+
+    Returns:
+        any: item in the queue.
+    """
+    if not q.empty():
+        return q.get()
+
+    while not eStop.is_set():
+        try:
+            item = q.get(timeout=retry_interval)
+            return item
+        except Empty:
+            if item_idx is not None:
+                logging.info(f"Waiting to get item {item_idx}")
+    if item_idx is not None:
+        logging.info(f"Early Stop. Return Default item at iter {item_idx}")
+    return default_item
+
+def queue_put(q: Queue, item, eStop: Event, retry_interval=1, item_idx=None):
+    """ wrapper for queue.put() with timeout, retry and event stop.
+
+    Args:
+        q (Queue): input queue.
+        item (_type_): item to put in the queue.
+        eStop (Event): event to stop the thread.
+        retry_interval (int, optional): time to wait before retry to put the item. Defaults to 1 second.
+        item_idx (_type_, optional): index of item to put. This is used for logging information. Defaults to None.
+    """
+    if not q.full():
+        q.put(item)
+        return 
+
+    while not eStop.is_set():
+        try:
+            q.put(item, timeout=retry_interval)
+            return
+        except Full:
+            if item_idx is not None:
+                logging.info(f"Waiting to put item at {item_idx}")
+    if item_idx is not None:
+        logging.info(f"Early Stop. No item is put at iter {item_idx}")  
+
+def batch_extract_thread(video_path: str, 
+                         img_batch_queue: Queue, 
+                         vis_img_batch_queue: Queue, 
+                         eStop: Event, 
+                         batch_size=32):
+    """Thread function to extract a batch of frames from video and put it to img_batch_queue and vis_img_batch_queue.
+
+    Args:
+        video_path (str): input video path.
+        img_batch_queue (Queue): output queue for batch of frames, used for processing.
+        vis_img_batch_queue (Queue): output queue for batch of frames, used for visualization.
+        eStop (Event): event to stop the thread.
+        batch_size (int, optional): number of images in a batch. Defaults to 32.
+    """
+    logging.info("Start Batch Extract Thread")
+    vidcap = VideoReader(video_path)
+    vis_img_batch_queue.put([vidcap.fps, vidcap.width, vidcap.height, len(vidcap)])
+    start_frame_idx = 0
+    last_frame_idx = len(vidcap)
+    end_frame_idx = start_frame_idx
+
+    while (start_frame_idx < last_frame_idx):
+        if eStop.is_set(): break
+        end_frame_idx = min(start_frame_idx + batch_size, last_frame_idx)
+        img_batch = []
+        for frame_idx in range(start_frame_idx, end_frame_idx):
+            img = vidcap[frame_idx]
+            if (img is None):
+                break
+            img_batch.append(img)
+        if (len(img_batch) == 0):
+            break
+        item_data = [start_frame_idx, img_batch]
+        queue_put(img_batch_queue, item_data, eStop)
+        queue_put(vis_img_batch_queue, item_data , eStop)
+        start_frame_idx = end_frame_idx
+    
+    if eStop.is_set():
+        queue_clear(img_batch_queue)
+        queue_clear(vis_img_batch_queue)
+    else:
+        logging.info(f"Finish batch_extract_thread for video_file {video_path} at end_frame_idx {end_frame_idx}.")
+    img_batch_queue.put(None)
+    vis_img_batch_queue.put(None)
+
+def detect_thread(obj_detector, 
+                  img_batch_queue: Queue, 
+                  det_queue: Queue, 
+                  eStop: Event,
+                  put_img_batch: bool=False):
+    """ detect_thread function to run detection on a batch of frames.
+
+    Args:
+        obj_detector (_type_): object detector, for example YOLOV7TRT/-ONXX.
+        img_batch_queue (Queue): input queue for batch of frames, which is the output from batch_extract_thread.
+        det_queue (Queue): output queue for detection results.
+        eStop (Event): event to stop the thread.
+        put_img_batch (bool, optional): If True, the input image batch will also put tp det_queue. 
+                This is often used for later step that require images of detected objects, such Human Pose or ReID.
+                Defaults to False.
+    """
+    logging.info("Start Detection Thread")
+    item = img_batch_queue.get()
+    start_frame_idx = -1
+    while item is not None:
+        if eStop.is_set(): break
+        start_frame_idx, img_batch = item
+        logging.info(f"Run detection at frame idx: {start_frame_idx}")
+        try:
+            det_result = obj_detector.infer_batch(img_batch)
+            item_data = [start_frame_idx, det_result]
+            if (put_img_batch):
+                item_data.append(img_batch)
+            queue_put(det_queue, item_data, eStop)
+        except Exception as e:
+            error_msg=[501, f"Error when running detection at frame idx: {start_frame_idx}]. "]
+            log_error_message = f"{error_msg[1]}. Error {e}"
+            logging.exception(log_error_message)
+            eStop.set()
+            break  
+        item = img_batch_queue.get()  
+
+    # Finish this thread. 
+    if eStop.is_set():
+        logging.warning(f"Early stop detect_thread at start_frame_idx {start_frame_idx}")
+        queue_clear(det_queue)
+    else:
+        logging.info(f"Finish detect_thread at start_frame_idx {start_frame_idx}.")    
+    det_queue.put(None)
+
+def bytetrack_thread(tracker_cfg, det_queue: Queue, track_queue: Queue, eStop: Event, conf_thres: float):
+    logging.info("Start Tracking Thread")
+    tracker = BYTETracker(
+        **tracker_cfg
+    )
+    item = det_queue.get()
+    start_frame_idx = -1
+    
+    while item is not None:
+        if eStop.is_set():break
+        start_frame_idx, det_result = item
+
+        if isinstance(det_result[0]['boxes'],np.ndarray):
+            det_result = [{key:torch.from_numpy(value) for key,value in dict_det.items()} for dict_det in det_result]
+            
+        try:
+            track_result = tracker.track_batch(start_frame_idx,det_result,conf_thres)
+        except Exception as e:
+            error_msg=[501,f"Error when running tracking at start_frame_idx {start_frame_idx}: {e}"]
+            log_error_message = f"Error {error_msg[0]}: {error_msg[1]}"
+            logging.error(log_error_message)
+            eStop.set()
+            break
+        queue_put(track_queue, [start_frame_idx, track_result], eStop)
+        item = det_queue.get()
+
+    # Finish this thread 
+    if eStop.is_set():
+        logging.warning(f"Early stop at start_frame_idx {start_frame_idx}.")
+        queue_clear(track_queue)
+    else:
+        logging.info(f"Finish track_thread.")
+    track_queue.put(None)
+    
+def update_progress_thread(visualize_queue: Queue, progress: Progress, eStop: Event):
+    """Show the progress of the video processing on Gradio, measured by the number of frames visualized.
+
+    Args:
+        visualize_queue (Queue): input queue for batch of frames, which is the output from batch_extract_thread.
+        progress (Progress): Gradio progress bar.
+        eStop (Event): event to stop the thread.
+    """
+    
+    fps, width, height, total_num_frames = visualize_queue.get()
+    progress(0, desc="Starting...")
+    start_frame_idx = -1
+    for frame_idx in progress.tqdm(range(total_num_frames), total=total_num_frames):
+        item = visualize_queue.get()
+        if (item is None):
+            break
+        start_frame_idx = item
+        if (start_frame_idx != frame_idx):
+            error_msg=[501, f"Error when runing update progress at start_frame_idx {start_frame_idx}. "]
+            log_error_message = f"Error {error_msg[0]}: {error_msg[1]}"
+            logging.error(log_error_message)
+            eStop.set()
+            break
+    
+    # Finish this thread    
+    if eStop.is_set():
+        logging.warning(f"Early stop at start_frame_idx {start_frame_idx}")
+    else:
+        logging.info(f"Finish update_progress_thread.")

models/models/engine/utils.py

@@ -0,0 +1,248 @@
+import numpy as np
+import cv2, torch
+
+from typing import Tuple, Dict, List
+
+def bbox_xyxy2cs(bbox: np.ndarray,
+                 padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
+    """Transform the bbox format from (x,y,w,h) into (center, scale)
+
+    Args:
+        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
+            as (left, top, right, bottom)
+        padding (float): BBox padding factor that will be multilied to scale.
+            Default: 1.0
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
+            (n, 2)
+        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
+            (n, 2)
+    """
+    # convert single bbox from (4, ) to (1, 4)
+    dim = bbox.ndim
+    if dim == 1:
+        bbox = bbox[None, :]
+
+    # get bbox center and scale
+    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
+    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
+    scale = np.hstack([x2 - x1, y2 - y1]) * padding
+
+    if dim == 1:
+        center = center[0]
+        scale = scale[0]
+
+    return center, scale
+
+def decode(simcc_x: np.ndarray, 
+            simcc_y: np.ndarray,
+            simcc_split_ratio, 
+            use_torch=False) -> Tuple[np.ndarray, np.ndarray]:
+    """Modulate simcc distribution with Gaussian.
+
+    Args:
+        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
+        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
+        simcc_split_ratio (int): The split ratio of simcc.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
+        - np.ndarray[float32]: scores in shape (K,) or (n, K)
+    """
+    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y, use_torch=use_torch)
+    keypoints /= simcc_split_ratio
+
+    return keypoints, scores
+
+
+def _fix_aspect_ratio(bbox_scale: np.ndarray,
+                        aspect_ratio: float) -> np.ndarray:
+    """Extend the scale to match the given aspect ratio.
+
+    Args:
+        scale (np.ndarray): The image scale (w, h) in shape (2, )
+        aspect_ratio (float): The ratio of ``w/h``
+
+    Returns:
+        np.ndarray: The reshaped image scale in (2, )
+    """
+    w, h = np.hsplit(bbox_scale, [1])
+    bbox_scale = np.where(w > h * aspect_ratio,
+                        np.hstack([w, w / aspect_ratio]),
+                        np.hstack([h * aspect_ratio, h]))
+    return bbox_scale
+
+
+def _rotate_point(pt: np.ndarray, 
+                    angle_rad: float) -> np.ndarray:
+    """Rotate a point by an angle.
+
+    Args:
+        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+        angle_rad (float): rotation angle in radian
+
+    Returns:
+        np.ndarray: Rotated point in shape (2, )
+    """
+    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+    rot_mat = np.array([[cs, -sn], [sn, cs]])
+    return rot_mat @ pt
+
+
+def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """To calculate the affine matrix, three pairs of points are required. This
+    function is used to get the 3rd point, given 2D points a & b.
+
+    The 3rd point is defined by rotating vector `a - b` by 90 degrees
+    anticlockwise, using b as the rotation center.
+
+    Args:
+        a (np.ndarray): The 1st point (x,y) in shape (2, )
+        b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+    Returns:
+        np.ndarray: The 3rd point.
+    """
+    direction = a - b
+    c = b + np.r_[-direction[1], direction[0]]
+    return c
+
+def get_warp_matrix(center: np.ndarray,
+                    scale: np.ndarray,
+                    rot: float,
+                    output_size: Tuple[int, int],
+                    shift: Tuple[float, float] = (0., 0.),
+                    inv: bool = False) -> np.ndarray:
+    """Calculate the affine transformation matrix that can warp the bbox area
+    in the input image to the output size.
+
+    Args:
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        rot (float): Rotation angle (degree).
+        output_size (np.ndarray[2, ] | list(2,)): Size of the
+            destination heatmaps.
+        shift (0-100%): Shift translation ratio wrt the width/height.
+            Default (0., 0.).
+        inv (bool): Option to inverse the affine transform direction.
+            (inv=False: src->dst or inv=True: dst->src)
+
+    Returns:
+        np.ndarray: A 2x3 transformation matrix
+    """
+    shift = np.array(shift)
+    src_w = scale[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    # compute transformation matrix
+    rot_rad = np.deg2rad(rot)
+    src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
+    dst_dir = np.array([0., dst_w * -0.5])
+
+    # get four corners of the src rectangle in the original image
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale * shift
+    src[1, :] = center + src_dir + scale * shift
+    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
+
+    # get four corners of the dst rectangle in the input image
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return warp_mat
+
+
+def top_down_affine(input_size: dict, 
+                    bbox_scale: dict, 
+                    bbox_center: dict,
+                    img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Get the bbox image as the model input by affine transform.
+
+        Args:
+            input_size (dict): The input size of the model.
+            bbox_scale (dict): The bbox scale of the img.
+            bbox_center (dict): The bbox center of the img.
+            img (np.ndarray): The original image.
+
+        Returns:
+            tuple: A tuple containing center and scale.
+            - np.ndarray[float32]: img after affine transform.
+            - np.ndarray[float32]: bbox scale after affine transform.
+        """
+        w, h = input_size
+        warp_size = (int(w), int(h))
+
+        # reshape bbox to fixed aspect ratio
+        bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
+
+        # get the affine matrix
+        center = bbox_center
+        scale = bbox_scale
+        rot = 0
+        warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
+
+        # do affine transform
+        img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
+
+        return img, bbox_scale
+
+def get_simcc_maximum(simcc_x: np.ndarray,
+                        simcc_y: np.ndarray,
+                        use_torch=False) -> Tuple[np.ndarray, np.ndarray]:
+    """Get maximum response location and value from simcc representations.
+
+    Note:
+        instance number: N
+        num_keypoints: K
+        heatmap height: H
+        heatmap width: W
+
+    Args:
+        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
+        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
+
+    Returns:
+        tuple:
+        - locs (np.ndarray): locations of maximum heatmap responses in shape
+            (K, 2) or (N, K, 2)
+        - vals (np.ndarray): values of maximum heatmap responses in shape
+            (K,) or (N, K)
+    """
+    N, K, Wx = simcc_x.shape
+    simcc_x = simcc_x.reshape(N * K, -1)
+    simcc_y = simcc_y.reshape(N * K, -1)
+
+    # get maximum value locations
+    x_locs = np.argmax(simcc_x, axis=1)
+    y_locs = np.argmax(simcc_y, axis=1)
+    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
+    if use_torch:
+        max_val_x = torch.max(simcc_x, dim=1)[0]
+        max_val_y = torch.max(simcc_y, dim=1)[0]
+    else:
+        max_val_x = np.amax(simcc_x, axis=1)
+        max_val_y = np.amax(simcc_y, axis=1)
+
+    # get maximum value across x and y axis
+    mask = max_val_x > max_val_y
+    max_val_x[mask] = max_val_y[mask]
+    vals = max_val_x
+    locs[vals <= 0.] = -1
+
+    # reshape
+    locs = locs.reshape(N, K, 2)
+    vals = vals.reshape(N, K)
+
+    return locs, vals

models/models/engine/visualizer.py

@@ -0,0 +1,143 @@
+from abc import abstractmethod
+from typing import List, Optional
+import cv2
+import subprocess
+import numpy as np 
+
+def putText(img, text: str, position, 
+               text_font: int=0, text_scale: int=1,
+               bg_color=(255,255,255),
+               text_color=(255,0,255),
+               bg_thickness=8, 
+               text_thickness=1, 
+               lineType=cv2.LINE_AA):
+    """ Function to put text on image.
+    
+    Args:
+        img (_type_): 
+        text (str): _description_
+        position (_type_): Top-left position of text.
+        text_font (int, optional): font size of text. Defaults to 0.
+        text_scale (int, optional): text scale. Defaults to 1.
+        bg_color (tuple, optional): text background color. Defaults to (255,255,255).
+        text_color (tuple, optional): text foreground color. Defaults to (255,0,255).
+        bg_thickness (int, optional): text background thickness. Defaults to 8.
+        text_thickness (int, optional): text foreground thickness. Defaults to 1.
+        lineType (_type_, optional): line type. Defaults to cv2.LINE_AA.
+
+    Returns:
+        _type_: _description_
+    """
+    img = cv2.putText(img, text, position, text_font, text_scale, bg_color, thickness=bg_thickness, lineType=lineType)
+    img = cv2.putText(img, text, position, text_font, text_scale, text_color, thickness=text_thickness, lineType=lineType)
+    return img 
+
+
+class BaseVisualizer():
+
+    def __init__(self, class_names: Optional[List[str]], fps: int=-1, min_width: int=-1):
+        """ Visualizer class for visualization (track_results + count_results).
+
+        Args:
+            class_map_ids (Dict): class mapping dictionary to map model's class to original class. Eg {0: 1, 1: 0, 2: 2, 3: 3} mean we swap class ID between 0 and 1.
+            fps (int): FPS for output video. If fps = -1, it will have same fps as input video.
+            min_width (int): minimum width for output video (height will be scaled to keep aspect ratio as input video). If min_width = -1, it will have same resolution as input video.
+        """
+
+        self.fps = fps
+        self.min_width = min_width
+        self.class_names = class_names
+    
+    def init_writer(self, input_video_info: List[int], output_path: str):
+        """ Init video writer for write visualized frame to output video.
+
+        Args:
+            input_video_info (List[int]): It is a list that includes 4 elements of input video information (fps, width, height, num_frames).
+            output_path (str): Path to save output video. 
+        """
+        if (self.fps == -1):
+            self.fps = input_video_info[0]
+        self.width, self.height = input_video_info[1], input_video_info[2]
+        if (self.min_width > 0):
+            out_width = min(self.min_width, self.width)
+            self.height = (self.height * out_width)//self.width
+            self.width = out_width
+        self.output_path = output_path
+        self.writer = cv2.VideoWriter(self.output_path, cv2.VideoWriter_fourcc(*"mp4v"), int(self.fps), (self.width, self.height))
+        
+    @staticmethod
+    def get_color(idx):
+        idx = idx * 3
+        color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
+
+        return color
+
+    @staticmethod
+    def draw_dash_line(img,pt1,pt2,color,thickness=1,style='dotted',gap=20):
+        dist =((pt1[0]-pt2[0])**2+(pt1[1]-pt2[1])**2)**.5
+        pts= []
+        for i in  np.arange(0,dist,gap):
+            r=i/dist
+            x=int((pt1[0]*(1-r)+pt2[0]*r)+.5)
+            y=int((pt1[1]*(1-r)+pt2[1]*r)+.5)
+            p = (x,y)
+            pts.append(p)
+        if len(pts) ==0: 
+            return
+        if style=='dotted':
+            for p in pts:
+                cv2.circle(img,p,thickness,color,-1)
+        else:
+            s=pts[0]
+            e=pts[0]
+            i=0
+            for p in pts:
+                s=e
+                e=p
+                if i%2==1:
+                    cv2.line(img,s,e,color,thickness)
+                i+=1
+
+    @staticmethod
+    def draw_dash_poly(img,pts,color,thickness=1,style='dotted',gap=20):
+        """ draw a polygon with dash line.
+
+        Args:
+            img (_type_): input image.
+            pts (_type_): _description_
+            color (_type_): _description_
+            thickness (int, optional): _description_. Defaults to 1.
+            style (str, optional): _description_. Defaults to 'dotted'.
+            gap (int, optional): _description_. Defaults to 20.
+
+        Returns:
+            _type_: _description_
+        """
+        s=pts[0]
+        e=pts[0]
+        pts.append(pts.pop(0))
+        for p in pts:
+            s=e
+            e=p
+            BaseVisualizer.draw_dash_line(img,s,e,color,thickness,style,gap=gap)
+        return img
+
+    @staticmethod
+    def draw_dash_rect(img,pt1,pt2,color,thickness=1,style='dotted',gap=10):
+        pts = [pt1,(pt2[0],pt1[1]),pt2,(pt1[0],pt2[1])] 
+        return BaseVisualizer.draw_dash_poly(img,pts,color,thickness,style,gap=gap)   
+    
+    def close(self):
+        """ Function to release video writer. It should be called after finish visualization for all input frames.
+        """
+        self.writer.release()
+    
+    def convert(self):
+        subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {self.output_path} -c:v libx264 {self.output_path}".split())
+
+    @abstractmethod
+    def visualize(self, *args,**kwargs):
+        """ Each project should implement this function to visualize a frame.
+
+        """
+        raise NotImplementedError

models/models/pose/rtmpose.py

@@ -0,0 +1,193 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+
+from typing import List, Tuple, Dict
+
+import numpy as np
+
+
+from models.base.onnx_base import ONNX_Base
+from models.base.trt_base import TRT_Base
+from models.engine.utils import *
+
+class RTMPose():
+    def __init__(self, 
+                 use_torch: bool=False) -> None:
+        self.use_torch = use_torch
+
+    def preprocess(self,
+                   input_data: np.ndarray, 
+                   input_size: Tuple[int, int] = (192, 256)) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Do preprocessing for RTMPose model inference.
+
+        Args:
+            img (np.ndarray): Input image in shape.
+            input_size (tuple): Input image size in shape (w, h).
+
+        Returns:
+            tuple:
+            - resized_img (np.ndarray): Preprocessed image.
+            - center (np.ndarray): Center of image.
+            - scale (np.ndarray): Scale of image.
+        """
+        tensor_data = []
+        # get shape of image
+        scales =[]
+        centers = []
+        for i in range(len(input_data)):
+            img = input_data[i]
+            img_shape = img.shape[:2]
+            bbox = np.array([0, 0, img_shape[1], img_shape[0]])
+
+            # get center and scale
+            center, scale = bbox_xyxy2cs(bbox, padding=1.25)
+            # do affine transformation
+            resized_img, scale = top_down_affine(input_size, scale, center, img)
+
+            # normalize image
+            mean = np.array([123.675, 116.28, 103.53])
+            std = np.array([58.395, 57.12, 57.375])
+            resized_img = (resized_img - mean) / std
+
+            centers.append(center)
+            scales.append(scale)
+
+            if self.use_torch:
+                tensor_data.append(torch.from_numpy(resized_img).to(self.device))
+            else:
+                tensor_data.append(resized_img.transpose(2, 0, 1))
+
+            
+        if self.use_torch:
+            tensor_data = torch.stack(tensor_data, dim=0)[:, :, :, [2, 1, 0]].permute(0, 3, 1, 2).float().contiguous()
+        else:
+            tensor_data = np.stack(tensor_data, axis=0)
+
+        return tensor_data, centers, scales
+
+    def postprocess(self, outputs: List[np.ndarray],
+                model_input_size: Tuple[int, int],
+                centers: List[np.ndarray],
+                scales: List[np.ndarray],
+                simcc_split_ratio: float = 2.0,
+                use_torch=False
+                ) -> Tuple[np.ndarray, np.ndarray]:
+        """Postprocess for RTMPose model output.
+
+        Args:
+            outputs (np.ndarray): Output of RTMPose model.
+            model_input_size (tuple): RTMPose model Input image size.
+            center List[tuple(int,int)]: List of Center of bbox in shape (x, y).
+            scale List[tuple(int,int)]: List of Scales of bbox in shape (w, h).
+            simcc_split_ratio (float): Split ratio of simcc.
+
+        Returns:
+            tuple:
+            - keypoints (np.ndarray): Rescaled keypoints.
+            - scores (np.ndarray): Model predict scores.
+        """
+        # use simcc to decode
+        simcc_x, simcc_y = outputs
+        tensor_keypoints = []
+        tensor_scores = []
+        assert simcc_x.shape[0] == simcc_y.shape[0]
+        for i in range(simcc_x.shape[0]):
+            
+            simcc_x_3d = simcc_x[i][np.newaxis, :, :]
+            simcc_y_3d = simcc_y[i][np.newaxis, :, :]
+
+            keypoints, scores = decode(simcc_x_3d, simcc_y_3d, simcc_split_ratio, use_torch=use_torch)
+
+            # rescale keypoints
+            keypoints = keypoints / model_input_size * scales[i] + centers[i] - scales[i] / 2
+
+            tensor_keypoints.append(keypoints)
+            tensor_scores.append(scores)
+        
+        tensor_keypoints = np.vstack(tensor_keypoints)
+        tensor_scores = np.vstack(tensor_scores)
+        return tensor_keypoints, tensor_scores
+
+    def crop_objects(self, image: np.ndarray, bounding_boxes: np.ndarray):
+        """ Function to crop objects in input image.
+
+        Args:
+            image (np.ndarray): input image with shape (H, W, C).
+            bounding_boxes (np.ndarray): Array with shape Nx4 with N is the number of objects.
+        """
+        max_h, max_w = image.shape[:2]
+        cropped_images = []
+        for box in bounding_boxes:
+            x_top, y_top, x_bottom, y_bottom, _ = box.astype(int).tolist()
+            x_top = max(0, x_top)
+            y_top = max(0, y_top)
+            x_bottom = min(x_bottom, max_w)
+            y_bottom = min(y_bottom, max_h)
+            cropped_image = image[y_top:y_bottom, x_top:x_bottom]
+            cropped_images.append(cropped_image)
+        return cropped_images
+class RTMPoseONNX(ONNX_Base, RTMPose):
+    def __init__(self,
+                use_torch,
+                img_shape: Tuple[int, int, int]=(3, 256, 192),
+                batch_size: int=32,
+                model_path: str="",
+                device: str='0'):
+        #/home/ccvn/Workspace/haimd/CC-Demo-Collection/end2end.onnx
+        """_summary_
+        RTMPose ONNX class for inference, which is base on ONNX_BASE and RTMPose 
+        Args:
+            use_torch (_type_): use torch tensor or numpy array in preprocess and postprocess function.
+            img_shape (Tuple[int, int], optional): _description_. Defaults to (640, 640).
+            batch_size (int, optional): _description_. Defaults to 32.
+            model_path (str, optional): _description_. Defaults to "".
+            device (str, optional): _description_. Defaults to '0'.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        RTMPose.__init__(self,
+                        use_torch=use_torch)
+    
+    def infer_batch(self, image_batch: np.ndarray):
+
+        h, w = self.session.get_inputs()[0].shape[2:]
+        model_input_size = (w, h)
+        numpy_array_data, centers, scales = self.preprocess(image_batch, model_input_size)
+        numpy_array_data = numpy_array_data.astype(np.float32)
+        results = super().infer_batch(numpy_array_data)
+        keypoints, scores = self.postprocess(results, model_input_size, centers, scales)
+        return {'keypoints':keypoints, 'scores': scores}
+
+class RTMPoseTRT(TRT_Base, RTMPose):
+    def __init__(self, 
+                    use_torch,
+                    img_shape: Tuple[int, int, int]=(3, 256, 192),
+                    batch_size: int=1,
+                    model_path: str="",
+                    device: str='0',):
+        """ RTMPoseTRT class for inference, which is based on TRT_Base and RTMPose.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        RTMPose.__init__(self,
+                         use_torch=use_torch)
+    
+    def infer_batch(self, image_batch: np.ndarray):
+
+        model_input_size = (self.img_shape[-1], self.img_shape[1])
+        tensor_data, centers, scales = self.preprocess(image_batch, model_input_size)
+        self.change_runtime_dimension(input_shape=(len(tensor_data), 3, model_input_size[1], model_input_size[0]))
+        self.model['binding_addrs']['input'] = int(tensor_data.data_ptr())
+        self.model['context'].execute_v2(list(self.model['binding_addrs'].values()))
+        simcc_x = self.model['bindings']['simcc_x'].data.cpu()
+        simcc_y = self.model['bindings']['simcc_y'].data.cpu()
+
+        results = (simcc_x, simcc_y)
+        keypoints, scores = self.postprocess(results, model_input_size, centers, scales, use_torch=self.use_torch)
+        return {'keypoints':keypoints, 'scores': scores}
+
+    

models/models/reids/solider.py

@@ -0,0 +1,165 @@
+from typing import Tuple, Union
+from models.base.trt_base import TRT_Base
+from models.base.onnx_base import ONNX_Base
+import torch
+import cv2
+import numpy as np
+
+class SOLIDERBase(): 
+    def __init__(self, use_torch: bool=False):
+        """ SOLIDERBase class for inference.
+
+        Args:
+            preprocess_cfg (Dict):
+                - mean (List[float, float, float]): mean offset values for preprocessing.
+                - std (List[float, float, float]): standard deviation offset values for preprocessing.
+            use_torch (bool): use torch tensor or numpy array in preprocess and postprocess function.
+        """
+        self.use_torch = use_torch
+    
+    def crop_objects(self, image: np.ndarray, bounding_boxes: np.ndarray):
+        """ Function to crop objects in input image.
+
+        Args:
+            image (np.ndarray): input image with shape (H, W, C).
+            bounding_boxes (np.ndarray): Array with shape Nx4 with N is the number of objects.
+        """
+        max_h, max_w = image.shape[:2]
+        cropped_images = []
+        for box in bounding_boxes:
+            x_top, y_top, x_bottom, y_bottom, _ = box.astype(int).tolist()
+            x_top = max(0, x_top)
+            y_top = max(0, y_top)
+            x_bottom = min(x_bottom, max_w)
+            y_bottom = min(y_bottom, max_h)
+            cropped_image = image[y_top:y_bottom, x_top:x_bottom]
+            cropped_images.append(cropped_image)
+        return cropped_images
+    
+    def preprocess(self, input_data: np.ndarray):
+        """ Preprocess function for input data.
+
+        Args:
+            input_data (np.ndarray): batch input image.
+        """
+        tensor_data = []
+        if ((isinstance(input_data, np.ndarray)) and (len(input_data.shape) == 3)):
+            input_data = [input_data]
+        for i in range(len(input_data)):
+            img = input_data[i]
+            img = cv2.resize(img, self.input_shape[2:][::-1], interpolation=cv2.INTER_LINEAR)
+            if self.use_torch:
+                tensor_data.append(torch.from_numpy(img).to(self.device))
+            else:
+                tensor_data.append(img)
+        if self.use_torch:
+            tensor_data = torch.stack(tensor_data, dim=0)
+            tensor_data = tensor_data.permute(0, 3, 1, 2).float().contiguous()
+        else:
+            tensor_data = np.stack(tensor_data, axis=0)
+            tensor_data = tensor_data.transpose((0, 3, 1, 2)).astype(np.float32)
+        return tensor_data
+
+    def postprocess(self, embeddings: Union[torch.tensor,np.ndarray]):
+        """ Postprocess function for input data.
+
+        Args:
+            embeddings (torch.tensor or np.ndarray): output embeddingss.
+        """
+
+        if (self.use_torch):
+            norms = torch.unsqueeze(torch.norm(embeddings, dim=-1), dim=-1)
+        else:
+            norms = np.expand_dims(np.linalg.norm(embeddings, axis=-1), axis=-1)
+        
+        normalized_embeddings = embeddings/norms
+        return normalized_embeddings
+
+    def batch_padding(self, input_data: Union[torch.tensor, np.ndarray], batch_size: int) -> Union[torch.tensor, np.ndarray]:
+        """Since the current model does not support Dynamic batch size, we perform padding to the input data.
+
+        Args:
+            input_data (Union[torch.tensor, np.ndarray]): input data. 
+            batch_size (int): batch size for inference.
+
+        Returns:
+            input_data (Union[torch.tensor, np.ndarray]): input data. 
+        """
+        n_pad = batch_size - len(input_data)
+        if n_pad>0:
+            if self.use_torch:
+                input_data = torch.cat((input_data, input_data[:n_pad]), dim=0)
+            else:
+                input_data = np.concatenate((input_data, input_data[:n_pad]), axis=0)
+        return input_data
+
+class SOLIDERONNX(ONNX_Base, SOLIDERBase):
+    def __init__(self, 
+                    batch_size: int,
+                    model_path: str,
+                    img_shape: Tuple[int, int, int]=(3, 384, 128),
+                    device: str='0',):
+        """ SOLIDER ONNX class for inference, which is based on ONNX_Base and SOLIDERBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        SOLIDERBase.__init__(self, use_torch=False)
+    
+    def infer_batch(self, image_batch: np.ndarray) -> np.ndarray:
+        """ Batch inference function for batch input image.
+
+        Args:
+            image_batch (np.ndarray): batch of input image.
+        """
+        num_images = len(image_batch)
+        assert num_images <= self.batch_size, "the number of input images must be smaller or equal to the Batch size."
+        numpy_array_data = self.preprocess(image_batch)
+        
+        # Padding data to the batch size
+        padding_data = self.batch_padding(numpy_array_data, self.batch_size)
+        results = super().infer_batch(padding_data)
+        
+        # Crop the padding data and postprocess
+        feats = results[0][:num_images]
+        feats = self.postprocess(feats)
+        return feats
+
+class SOLIDERTRT(TRT_Base, SOLIDERBase):
+    def __init__(self, 
+                    batch_size: int,
+                    model_path: str,
+                    img_shape: Tuple[int, int, int]=(3, 384, 128),
+                    device: str='0',):
+        """ SOLIDER TRT class for inference, which is based on TRT_Base and SOLIDERBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        SOLIDERBase.__init__(self, use_torch=True)
+        
+    def infer_batch(self, image_batch: np.ndarray) -> np.ndarray:
+        """ Batch inference function for batch input image.
+
+        Args:
+            image_batch (np.ndarray): batch of input image.
+        """
+        num_images = len(image_batch)
+        assert num_images <= self.batch_size, "the number of input images must be smaller or equal to the Batch size."
+        tensor_data = self.preprocess(image_batch)
+        
+        # Padding data to the batch size
+        padding_data = self.batch_padding(tensor_data, self.batch_size)
+        self.model['binding_addrs']['input'] = int(padding_data.data_ptr())
+        self.model['context'].execute_v2(list(self.model['binding_addrs'].values()))
+        feats = self.model['bindings']['output'].data.cpu()
+        
+        # Crop the padding data and postprocess
+        feats = feats[:num_images]
+        feats = self.postprocess(feats)
+        feats = feats.float().numpy()
+        return feats
+
+

models/models/trackers/__init__.py

@@ -0,0 +1,2 @@
+from .byte_track import BYTETracker
+from .reid_parallel_tracker import ParallelTracker

models/models/trackers/byte_track.py

@@ -0,0 +1,74 @@
+from mmdet.models.trackers.byte_tracker import ByteTracker as MMByteTracker
+from mmdet.structures import DetDataSample
+from mmengine.structures import InstanceData
+from typing import List, Dict
+import numpy as np
+import torch
+
+class BYTETracker(MMByteTracker):
+    def __init__(self, *args, **kwargs):
+        """ ByteTracker class for tracking.
+
+        Args:
+            obj_score_thrs (Dict):
+                - high (float): if detection box > high -> high_score_detections for first association.
+                - low (float): if low < detection box < high -> low_score_detections for second association.
+            init_track_thr (float): Detection score threshold for initializing a new tracklet.
+            weight_iou_with_det_scores (bool): Whether using detection scores to weight IOU which is used for matching.
+            match_iou_thrs (Dict): IOU distance threshold for matching between two frames.
+                - high (float): Threshold of the first matching.
+                - low (float): Threshold of the second matching.
+                - tentative (float): Threshold of the matching for tentative tracklets.
+            num_frames_retain (int): If a track is disappeared more than num_frames_retain frames, it will be deleted in the memo.
+            motion (Dict): Config for motion.
+                - type (str): Motion type (KalmanFilter, LinearFilter).
+        """
+        super().__init__(*args, **kwargs)
+    
+    def prepare_det_data_sample(self, frame_idx: int, boxes: torch.tensor, labels: torch.tensor, scores: torch.tensor):
+        """ Function to prepare DetDataSample.
+
+        Args:
+            frame_idx (int): frame_idx of this frame.
+            boxes (torch.tensor(N, 4)): bounding boxes prediction of this frame.
+            labels (torch.tensor(N)): labels prediction of this frame.
+            scores (torch.tensor(N)): scores prediction of this frame.
+        """
+
+        det_data_sample = DetDataSample()
+        det_data_sample.meta_info = dict(frame_id=frame_idx)
+        det_data_sample.pred_instances = InstanceData()
+        det_data_sample.pred_instances.bboxes = boxes
+        det_data_sample.pred_instances.labels = labels
+        det_data_sample.pred_instances.scores = scores
+        return det_data_sample
+
+    def track_batch(self, start_frame_idx: int, det_results: List[Dict], conf_thres: float=0.0):
+        """ Batch inference function for batch det results.
+
+        Args:
+            start_frame_idx (int): start_frame_idx of this batch.
+            det_results (List[Dict]): detection results of this batch.
+            conf_thres (float): If a tracklet's confidence < confidence threshold, it will be removed.
+        """
+
+        track_outputs = []
+        for frame_id, frame_det_outputs in enumerate(det_results):
+            boxes, labels = frame_det_outputs.pop("boxes"), frame_det_outputs.pop("labels")
+            det_data_sample = self.prepare_det_data_sample(frame_id+start_frame_idx, boxes[:, :4], labels, boxes[:, 4])
+            track_instances = self.track(det_data_sample)
+            boxes = torch.cat([track_instances.bboxes, torch.unsqueeze(track_instances.scores, dim=-1)], dim=-1)
+            labels = track_instances.labels
+            track_ids = track_instances.instances_id
+            boxes = np.around(boxes.numpy(),decimals=3)
+            labels = labels.numpy().astype(np.uint)
+            track_ids = track_ids.numpy().astype(np.uint)
+            idxs = np.where(boxes[:, 4] > conf_thres)[0]
+            
+            track_outputs.append({
+                    "boxes": boxes[idxs],
+                    "labels": labels[idxs],
+                    "ids": track_ids[idxs]
+                })
+        return track_outputs
+        

models/models/trackers/reid_parallel_tracker/__init__.py

@@ -0,0 +1 @@
+from .parallel_tracker import ParallelTracker

models/models/trackers/reid_parallel_tracker/base_tracker.py

@@ -0,0 +1,315 @@
+
+from typing import List, Tuple, Dict
+import numpy as np
+from .core.tracklet import (Tracklet, TrackState, add_stracks, subtract_stracks, remove_duplicate_stracks)
+from .core.kalman_filter import KalmanFilter
+from .core.basetrack import BaseTrack
+from .core.matching import iou_scores
+
+class BaseTracker(object):
+    def __init__(self, 
+                det_thr=dict(high=0.3,low=0.1, min_height=10, min_width=10),
+                new_track_cfg = dict(active_thr=0.9, active_iou=0.7, thr=0.4, min_size=(10,5), feat_buffer=30),
+                lost_track_cfg = dict(max_length=32, min_size=(10,5)),
+                smooth_update = False,
+                ):
+        """ Base class for SORT tracker
+
+        Args:
+            det_thr (dict, optional): 
+                + high: threshold score to consider highly confident detection. 
+                Defaults to 0.3.
+                + low : threshold score to consider low confident detection. 
+                Detection with lower score than this threshold are ignored. 
+                Defaults to 0.1.
+                
+            new_track_cfg (dict, optional): Config for initializing new track.
+                + thr (float, optional): threshold to initialize new track.
+                    A detection with score higher than this threshold will be initialized as a new (unconfirmed) track if it does not match with any tracks. 
+                    Defaults to 0.4.
+                + active_thr (float, optional): threshold to activate a new track. 
+                + active_iou (float, optional): threshold to activate a new track. A new track (score > active)thr and iou < active_iou) is a high confident detection without significant overlap with other objects are activated immediatly without confirming in the next frames.
+                + min_size (tuple, optional): minimum (high,width) of the bounding box to be considered as new track.
+                Defaults to (80,40).
+                + feat_buffer (int, optional): number of frames to store the features of the new track.
+            lost_track_cfg (dict, optional): Config for lost track.
+                + max_length (int): number of frames that the lost tracks are keep before being removed. It is also the length of buffer to store the features. 
+                Defaults to 30.
+                + min_size (tuple, optional): If the lost object size smaller than this min_size(high,width) will be removed. 
+                Defaults to (40,20).
+                + tracking_region (x1,y1,x2,y2): Top-Left, Bottom-right coordinates of the tracking region. If objects move out of this region, they will be removed.
+            smooth_update (bool, optional): If True, when a lost object is refind, we interpolate its missing coordinate during lost, and use these interpolated bboxes to update Kalman Filter. Thus, avoid excessive gain when updating the Kalman filter (smoother).
+        """
+        self.tracked_stracks = []  # type: list[Tracklet]
+        self.lost_stracks = []  # type: list[Tracklet]
+        self.removed_stracks = []  # type: list[Tracklet]
+        BaseTrack.clear_count()
+
+        self.frame_id = 0
+        self.det_thr = det_thr
+        self.new_track_cfg = new_track_cfg
+        self.lost_track_cfg = lost_track_cfg
+        
+
+        self.kalman_filter = KalmanFilter()
+        self.smooth_update = smooth_update
+
+    def preprocess_det_result(self,det_results: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        boxes = det_results['boxes']
+        boxes = boxes.reshape(-1, 5)
+        h = boxes[:,3]-boxes[:,1]
+        w = boxes[:,2]-boxes[:,0]
+        valid_inds = np.logical_and(h>self.det_thr["min_height"], w>self.det_thr["min_width"])
+        for k,v in det_results.items():
+            if k in ['boxes', 'labels', 'angles', 'obj_imgs', 'embeddings']:
+                if k == 'obj_imgs':
+                    det_results[k] = [v[_i] for _i, _valid in enumerate(valid_inds) if _valid]
+                else:
+                    det_results[k] = v[valid_inds]
+        return det_results
+
+    def split_detections_by_scores(self, 
+                                   det_result: Dict[str, np.ndarray])-> Tuple[List[Tracklet], List[Tracklet]]:
+        """ Split the detections into high score/lower score group.
+            det_result is a dict of {'boxes': np.ndarray(x1,y1,x2,y2,score), 'labels': np.ndarray}
+            Return:
+                detections_high: list[Tracklet]
+                detections_low: list[Tracklet]
+        """ 
+        detections_high = []
+        detections_low = []
+        
+        feat_history = self.new_track_cfg["feat_buffer"]
+        if len(det_result['boxes']):
+            bboxes = det_result['boxes'][:, :4]
+            scores = det_result['boxes'][:, 4]
+            classes = det_result['labels']
+            angles = np.array(det_result.get('angles',[None]*len(scores)))
+            features = np.array(det_result.get('embeddings',[None]*len(scores)))
+            obj_imgs = np.array(det_result.get('obj_imgs',[None]*len(scores)))
+
+            # Find high threshold detections
+            inds_high = scores >= self.det_thr["high"]
+
+            enable_reid_buffer = False
+            if hasattr(self, 'enable_reid_buffer'):
+                enable_reid_buffer = self.enable_reid_buffer
+
+            if np.any(inds_high):
+                detections_high = [Tracklet(Tracklet.tlbr_to_tlwh(tlbr), s, c, a, feat,feat_history=feat_history,
+                                            obj_img=obj_img, enable_buffer=enable_reid_buffer) for
+                                    (tlbr, s, c, a ,feat, obj_img) in zip(bboxes[inds_high], scores[inds_high], classes[inds_high],
+                                                                        angles[inds_high], features[inds_high], obj_imgs[inds_high])]
+            # Find low threshold detections
+            inds_low = np.logical_and(scores > self.det_thr["low"],
+                                      scores < self.det_thr["high"])
+            if np.any(inds_low):
+                detections_low = [Tracklet(Tracklet.tlbr_to_tlwh(tlbr), s, c, a, feat, feat_history=feat_history,
+                                           obj_img=obj_img, enable_buffer=enable_reid_buffer) for
+                                    (tlbr, s, c, a, feat, obj_img) in zip(bboxes[inds_low], scores[inds_low], classes[inds_low],
+                                                                          angles[inds_low], features[inds_low], obj_imgs[inds_low])]
+
+        return detections_high, detections_low
+
+    def split_tracks_by_activation(self) -> Tuple[List[Tracklet], List[Tracklet]]:
+        """ Split the tracks into trackpool=(tracked_tracks + lost_tracks) and unconfirmed (just initialize)
+        Returns:
+            strack_pool: List[Tracklet]
+            unconfirmed: List[Tracklet]
+        """
+        unconfirmed = []
+        tracked_stracks = []  # type: list[Tracklet]
+        for track in self.tracked_stracks:
+            if not track.is_activated:
+                unconfirmed.append(track)
+            else:
+                tracked_stracks.append(track)
+        strack_pool = add_stracks(tracked_stracks, self.lost_stracks)
+        return strack_pool, unconfirmed
+
+    def predict_with_gmc(self, 
+                         strack_pool: List[Tracklet], 
+                         unconfirmed: List[Tracklet], 
+                         Hmat: np.array=None) -> None:
+        """ Predict the current location with KF, and compensate for Camera Motion
+
+        Args:
+            strack_pool (List[Tracklet]): list of tracked tracks 
+            unconfirmed (List[Tracklet]): list of unconfirmed tracks
+            Hmat (np.array): Homography transformation matrix
+        """
+        Tracklet.multi_predict(strack_pool)
+        if Hmat is not None:
+            Tracklet.multi_gmc(strack_pool,Hmat)
+            Tracklet.multi_gmc(unconfirmed,Hmat)
+
+    def update_matched_tracks(self, 
+                              matches: np.ndarray, 
+                              strack_pool: List[Tracklet], 
+                              detections: List[Tracklet])-> Tuple[List[Tracklet], List[Tracklet]]:
+        """Update the matched tracks with Kalman Filter
+
+        Args:
+            matches (np.ndarray): [Nx2] index of the matched tracks and detections
+            strack_pool (List[Tracklet]): List of tracked tracks
+            detections (List[Tracklet]): List of detections
+
+        Returns:
+            activated_stracks (List[Tracklet]): List of tracked tracks that continue to be tracked (activated)
+            refind_stracks (List[Tracklet]): List of lost tracks that are refound in this frame (refind)
+        """
+        activated_stracks,refind_stracks=[],[]
+        for itracked, idet in matches:
+            track = strack_pool[itracked]
+            det = detections[idet]
+            if track.state == TrackState.Tracked:
+                #Perform Kalman Update/Feature Update
+                if self.smooth_update:
+                    track.smooth_update(det, self.frame_id)
+                else:
+                    track.update(det, self.frame_id)
+                activated_stracks.append(track)
+            else:
+                track.re_activate(det, self.frame_id, new_id=False)
+                refind_stracks.append(track)
+        
+        return activated_stracks,refind_stracks
+
+    def init_new_tracks(self, 
+                        detections: List[Tracklet], 
+                        u_detection: np.ndarray)-> List[Tracklet]:
+        """Initialize new tracks
+
+        Args:
+            detections (List[Tracklet]): List of detection objects
+            u_detection (np.ndarray): indices of the detections that are not matched with any tracks, 
+                                        and are considerd as new detections if its score is high enough
+        Returns:
+            List[Tracklet]: List of new tracks
+        """
+        new_tracks= []
+        for inew in u_detection:
+            det_= detections[inew] 
+            if det_.score >=self.new_track_cfg["thr"] and (not det_.is_too_small(self.new_track_cfg["min_size"])):
+                new_tracks.append(det_)
+        # The activate function will initialize the new tracks with a new id 
+        for track in new_tracks: 
+            # By default, new_track status is Unconfirmed (is_activated = False), except for the first frame
+            track.activate(self.kalman_filter, self.frame_id)        
+        return new_tracks 
+
+    def activate_new_tracks(self,new_tracks, current_tracks):
+        ious =iou_scores(new_tracks, current_tracks)
+        iou_max = ious.max(axis=1) if ious.shape[1]>0 else np.zeros(len(new_tracks))
+        active_thr = self.new_track_cfg.get('active_thr',0.7)
+        active_iou = self.new_track_cfg.get('active_iou',0.5)
+        for track, iou in zip(new_tracks, iou_max):
+            # For very high confident detection and non-overlap objects, we can activate it directly
+            if track.score >= active_thr and iou < active_iou:
+                track.mark_activated()    
+
+    def remove_lost_tracks(self):
+        removed_stracks=[]
+        """ Remove lost tracks if they are already lost for a certain conditions"""
+        for track in self.lost_stracks:
+            is_expired = track.is_expired(self.frame_id, self.lost_track_cfg["max_length"])
+            is_out_border, is_too_small = False, False
+            if self.lost_track_cfg.get('tracking_region',None) is not None:
+                is_out_border = track.is_out_border(self.lost_track_cfg["tracking_region"])
+            if self.lost_track_cfg.get('min_size',None) is not None:
+                is_too_small = track.is_too_small(self.lost_track_cfg["min_size"])
+                    
+            if is_expired or is_out_border or is_too_small:                
+                track.mark_removed()
+                removed_stracks.append(track)
+        return removed_stracks
+    
+    def merge_results(self, 
+                      activated_stracks: List[Tracklet],
+                      refind_stracks: List[Tracklet],
+                      new_stracks: List[Tracklet],
+                      removed_stracks: List[Tracklet]) -> Tuple[Dict, Dict]:
+        """ Merge the results from different types of tracks into the final results
+
+        Args:
+            activated_stracks (List[Tracklet]): activated tracks
+            refind_stracks (List[Tracklet]): refind tracks
+            removed_stracks (List[Tracklet]): removed tracks
+
+        Returns:
+            active_tracks (dict): dict of active tracks in the current frame. See format_track_results for the format.
+            lost_tracks (dict): dict of lost tracks in the current frame. See format_track_results for the format.
+        """
+        self.tracked_stracks = [t for t in self.tracked_stracks if t.state == TrackState.Tracked]
+        self.tracked_stracks = add_stracks(self.tracked_stracks, activated_stracks)
+        self.tracked_stracks = add_stracks(self.tracked_stracks, refind_stracks)
+        self.tracked_stracks = add_stracks(self.tracked_stracks, new_stracks)
+        self.lost_stracks = subtract_stracks(self.lost_stracks, self.tracked_stracks)
+        self.lost_stracks = subtract_stracks(self.lost_stracks, self.removed_stracks)
+        self.removed_stracks.extend(removed_stracks)
+        self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)
+
+        active_tracks = [track for track in self.tracked_stracks if track.is_activated]
+        active_tracks = self.format_track_results(active_tracks) if len(active_tracks)>0 else None
+        lost_tracks = self.format_track_results(self.lost_stracks) if len(self.lost_stracks)>0 else None
+        return active_tracks, lost_tracks
+
+    def format_track_results(self, 
+                             tracklets: List[Tracklet]) -> Dict[str, np.ndarray]:
+        """Format the tracking results to the required format
+
+        Args:
+            tracklets (Lost): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        tlbrs = []
+        ids = []
+        scores = []
+        cls = []
+        vel = []
+        angles = []
+        for t in tracklets:
+            tlbrs.append(t.tlbr)
+            ids.append(t.track_id)
+            scores.append(t.score)
+            cls.append(t.cls)
+            vel.append(t.vel_dir)
+            angles.append(t.angle)
+
+        track_outputs={
+            "boxes": np.concatenate([np.array(tlbrs), np.expand_dims(np.array(scores), axis=1)], axis=1),
+            "labels": np.array(cls),
+            "ids": np.array(ids),
+            "velocity": np.array(vel),  # motion velocity
+            "angles": np.array(angles), # body orientation
+        }
+        return track_outputs
+
+    def update(self, 
+               det_result: Dict, 
+               Hmat: np.array=None,
+               meta_data: Dict=None) -> Tuple[Dict, Dict]:
+        """ The main function to perform tracking, which may includes the follow steps:
+            1. Split the detections into high score/lower score group: 
+                - split_detections_by_scores
+            2. Split the tracks into trackpool=(tracked_tracks + lost_tracks) and unconfirmed (just initialize). 
+                - split_tracks_by_activation
+                - predict_with_gmc: predict the current location of these tracklets with KF, and compensate for Camera Motion
+            3. First association with high score detection boxes: 
+                - matcher_high 
+                - update_matched_tracks
+            4. Second association with low score detection boxes
+                - matcher_low
+                - update_matched_tracks if they are activated or refind
+                - mark new lost tracks
+            5. Third association, between new detections and unconfirmed tracks (usually tracks with only one beginning frame)
+                - matcher_unconfirmed
+                - remove unconfirmed tracks that does not match any detections
+                - init new track if the unconfirmed track is matched with a detection
+            6. Remove lost tracks if they are already lost for a certain frames
+            7. Update status for these trackes: active, lost, removed, uncofirmed.
+               Merge results and format the results
+        """
+        raise NotImplementedError

models/models/trackers/reid_parallel_tracker/core/basetrack.py

@@ -0,0 +1,165 @@
+import numpy as np
+from collections import OrderedDict
+
+
+class TrackState(object):
+    New = 0
+    Tracked = 1
+    Lost = 2
+    LongLost = 3
+    Removed = 4
+    Merged = 5
+
+
+class BaseTrack(object):
+    _count = 0
+
+    track_id = 0
+    is_activated = False
+    state = TrackState.New
+
+    history = OrderedDict()
+    feat_buffer = []
+    curr_feature = None
+    score = 0
+    start_frame = 0
+    frame_id = 0
+    time_since_update = 0
+        
+    # multi-camera
+    location = (np.inf, np.inf)
+    
+    is_merge = False
+
+    @property
+    def end_frame(self):
+        return self.frame_id
+
+    @staticmethod
+    def next_id():
+        BaseTrack._count += 1
+        return BaseTrack._count
+
+    def activate(self, *args):
+        raise NotImplementedError
+
+    def predict(self):
+        raise NotImplementedError
+
+    def update(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def mark_lost(self):
+        self.state = TrackState.Lost
+
+    def mark_long_lost(self):
+        self.state = TrackState.LongLost
+
+    def mark_removed(self):
+        self.state = TrackState.Removed
+
+    def mark_merged(self):
+        self.state = TrackState.Merged
+
+    def mark_activated(self):
+        self.is_activated = True
+        
+    def is_expired(self, frame_id, max_time_lost):
+        is_expired = (frame_id - self.end_frame) >= max_time_lost
+        return is_expired
+        
+    def is_out_border(self, tracking_region):
+        tlbr = self.tlbr
+        is_out= (tlbr[0] <= tracking_region[0]) or (tlbr[1] <= tracking_region[1]) or \
+                    (tlbr[2] >= tracking_region[2]) or (tlbr[3] >= tracking_region[3])
+        return is_out 
+
+    def is_too_small(self, min_size):
+        min_h, min_w = min_size
+        x,y,w,h = tuple(self.xywh)
+        return  (w < min_w)  or (h < min_h) or (w*h < min_h*min_w)
+    
+    def is_active(self):
+        return self.state == TrackState.Tracked or self.state == TrackState.New
+    
+    def is_lost(self):
+        return self.state == TrackState.Lost or self.state == TrackState.LongLost
+    
+    @staticmethod
+    def clear_count():
+        BaseTrack._count = 0
+   
+    @property
+    def tlwh(self):
+        """Get current position in bounding box format `(top left x, top left y,
+                width, height)`.
+            mean is the (x_center, y_top, w, h)
+        """
+
+        if self.mean is None:
+            return self._tlwh.copy()
+        ret = self.mean[:4].copy()
+        # ret[:2] -= ret[2:] / 2
+        ret[0] -= ret[2] / 2
+        return ret
+
+    @property
+    def tlbr(self):
+        """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
+        `(top left, bottom right)`.
+        """
+        ret = self.tlwh.copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    @property
+    def xywh(self):
+        """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
+        `(top left, bottom right)`.
+        """
+        ret = self.tlwh.copy()
+        ret[:2] += ret[2:] / 2.0
+        return ret
+
+    @staticmethod
+    def tlwh_to_xywh(tlwh):
+        """Convert bounding box to format `(center x, y top, width,
+        height)`.
+        """
+        ret = np.asarray(tlwh).copy()
+        # ret[:2] += ret[2:] / 2
+        ret[0] += ret[2] / 2
+        return ret
+
+    def to_xywh(self):
+        return self.tlwh_to_xywh(self.tlwh)
+
+    @staticmethod
+    def tlbr_to_tlwh(tlbr):
+        ret = np.asarray(tlbr).copy()
+        ret[2:] -= ret[:2]
+        return ret
+
+    @staticmethod
+    def tlwh_to_tlbr(tlwh):
+        ret = np.asarray(tlwh).copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    @staticmethod
+    def area(tlbr):
+        w=tlbr[:,2]-tlbr[:,0]
+        h=tlbr[:,3]-tlbr[:,1]
+        return w*h
+
+    @staticmethod
+    def height(tlbr):
+        h=tlbr[:,3]-tlbr[:,1]
+        return h
+
+    @staticmethod
+    def width(tlbr):
+        w=tlbr[:,2]-tlbr[:,0]
+        return w    
+    def __repr__(self):
+        return 'OT_{}_({}-{})'.format(self.track_id, self.start_frame, self.end_frame)

models/models/trackers/reid_parallel_tracker/core/homography.py

@@ -0,0 +1,89 @@
+import cv2
+import numpy as np
+
+def compute_Hmat_bev(landmarks):
+    assert len(landmarks)==4
+    '''
+        p1-----p2
+        |      |
+        p3-----p4
+    '''
+    x1,y1=landmarks[0]
+    x2,y2=landmarks[1]
+    x3,y3=landmarks[2]
+    x4,y4=landmarks[3]
+    h13=np.sqrt((x4-x2)**2 +(y4-y2)**2)
+    h24=np.sqrt((x3-x1)**2 +(y3-y1)**2)
+    new_p1=x1,(y3+y1)/2-h13
+    new_p2=x2,(y3+y1)/2-h24
+    new_p3=x1,(y3+y1)/2+h13
+    new_p4=x2,(y3+y1)/2+h24
+    # new_p1=(x3+x1)/2,(y3+y1)/2-h13
+    # new_p2=(x2+x4)/2,(y3+y1)/2-h24
+    # new_p3=(x3+x1)/2,(y3+y1)/2+h13
+    # new_p4=(x2+x4)/2,(y3+y1)/2+h24
+    src=np.float32(landmarks)
+    dst=np.float32([new_p1,new_p2,new_p3,new_p4])
+    #bird's eye view transformation matrix
+    Hmat = cv2.getPerspectiveTransform(src, dst) # The transformation matrix
+    return Hmat, np.rint(src), np.rint(dst)
+
+def compute_tang_alpha(landmarks_bev):
+    '''
+          p1---p2
+          /      \
+        p3--------p4
+    '''
+    x1,y1=landmarks_bev[0]
+    x2,y2=landmarks_bev[1]
+    x3,y3=landmarks_bev[2]
+    x4,y4=landmarks_bev[3]
+    ta_left = (x1-x3)/(y3-y1)
+    ta_right = (x4-x2)/(y4-y2)
+    return ta_left,ta_right 
+
+def compute_Hmat_left(landmarks):
+    assert len(landmarks)==4
+    '''
+        p1-----p2
+        |      |
+        p3-----p4
+    '''
+    x1,y1=landmarks[0]
+    x2,y2=landmarks[1]
+    x3,y3=landmarks[2]
+    x4,y4=landmarks[3]
+    l12=np.sqrt((x2-x1)**2 +(y2-y1)**2)
+    l34=np.sqrt((x4-x3)**2 +(y4-y3)**2)
+    h13=np.sqrt((x3-x1)**2 +(y3-y1)**2)
+    new_p1=(x1-l12,y3-h13)
+    new_p2=(x2,y3-h13)
+    new_p3=(x3-l34,y3)
+    new_p4=(x4,y3)
+    src=np.float32(landmarks)
+    dst=np.float32([new_p1,new_p2,new_p3,new_p4])
+    Hmat = cv2.getPerspectiveTransform(src, dst) # The transformation matrix
+    return Hmat, np.rint(src), np.rint(dst)
+
+def compute_Hmat_right(landmarks):
+    assert len(landmarks)==4
+    '''
+        p1-----p2
+        |      |
+        p3-----p4
+    '''
+    x1,y1=landmarks[0]
+    x2,y2=landmarks[1]
+    x3,y3=landmarks[2]
+    x4,y4=landmarks[3]
+    l12=np.sqrt((x2-x1)**2 +(y2-y1)**2)
+    l34=np.sqrt((x4-x3)**2 +(y4-y3)**2)
+
+    new_p1=(x2+x1)/2-l12,y2
+    new_p2=(x2+x1)/2+l12,y2
+    new_p3=(x3+x4)/2-l34,y4
+    new_p4=(x3+x4)/2+l34,y4
+    src=np.float32(landmarks)
+    dst=np.float32([new_p1,new_p2,new_p3,new_p4])
+    Hmat = cv2.getPerspectiveTransform(src, dst) # The transformation matrix
+    return Hmat, np.rint(src), np.rint(dst)

models/models/trackers/reid_parallel_tracker/core/kalman_filter.py

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

models/models/trackers/reid_parallel_tracker/core/matching.py

@@ -0,0 +1,405 @@
+import numpy as np
+import scipy
+import lap
+import torch
+from scipy.spatial.distance import cdist
+from numpy.linalg import norm
+from copy import deepcopy
+from mmdet.evaluation.functional import bbox_overlaps
+# from cython_bbox import bbox_overlaps as bbox_ious
+
+DEG2RAD= np.pi/180
+
+def merge_matches(m1, m2, shape):
+    O,P,Q = shape
+    m1 = np.asarray(m1)
+    m2 = np.asarray(m2)
+
+    M1 = scipy.sparse.coo_matrix((np.ones(len(m1)), (m1[:, 0], m1[:, 1])), shape=(O, P))
+    M2 = scipy.sparse.coo_matrix((np.ones(len(m2)), (m2[:, 0], m2[:, 1])), shape=(P, Q))
+
+    mask = M1*M2
+    match = mask.nonzero()
+    match = list(zip(match[0], match[1]))
+    unmatched_O = tuple(set(range(O)) - set([i for i, j in match]))
+    unmatched_Q = tuple(set(range(Q)) - set([j for i, j in match]))
+
+    return match, unmatched_O, unmatched_Q
+
+def linear_assignment(cost_matrix, thresh):
+    if cost_matrix.size == 0:
+        return np.empty((0, 2), dtype=int), np.array(range(cost_matrix.shape[0]), dtype=int), np.array(range(cost_matrix.shape[1]), dtype=int)
+    matches, unmatched_a, unmatched_b = [], [], []
+    cost, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
+    for ix, mx in enumerate(x):
+        if mx >= 0:
+            matches.append([ix, mx])
+    unmatched_a = np.where(x < 0)[0]
+    unmatched_b = np.where(y < 0)[0]
+    return np.asarray(matches, dtype=int).reshape(-1, 2), np.array(unmatched_a, dtype=int), np.array(unmatched_b, dtype=int)
+
+def topk_assignment(cost_matrix:np.ndarray, 
+                    thresh: float, 
+                    tolerance: float=None, 
+                    tolerance_ratio: float=0.15,
+                    topk: int=1):
+    """_summary_
+
+    Args:
+        cost_matrix (np.ndarray): [num_tracks x num_dets]
+        thresh (float): maximum distance threshold
+        tolerance (float, optional): the maximum distance between the best and second matches. Defaults to 0.05.
+        topk (int, optional):topk matching. Defaults to 1.
+
+    Returns:
+        _type_: _description_
+    """
+    if cost_matrix.size == 0:
+        return np.empty((0, 2), dtype=int), np.array(range(cost_matrix.shape[0]), dtype=int), np.array(range(cost_matrix.shape[1]), dtype=int)
+    
+    topk=min(topk,cost_matrix.shape[1])
+    if topk==1:
+        return  linear_assignment(cost_matrix, thresh)
+    
+    tolerance = tolerance or tolerance_ratio*thresh
+    # select topk cols has lowest cost value for each row
+    match_idx = np.argsort(cost_matrix, axis=1)[:, :topk]
+    match_scores = cost_matrix[np.arange(cost_matrix.shape[0])[:, None], match_idx]
+    valid_match = (match_scores < thresh) & ((match_scores-match_scores[:,0:1]) <= tolerance)
+    
+    match_cols = match_idx[valid_match]
+    unmatched_a = np.where(np.sum(valid_match, axis=1) == 0)[0]
+    unmatched_b = np.array([i for i in range(cost_matrix.shape[1]) if i not in match_cols])
+    
+    # add index of rows 
+    match_rows = np.repeat(np.arange(cost_matrix.shape[0])[:,None],topk,1)
+    match_rows = match_rows[valid_match]
+
+    matches = np.array([match_rows,match_cols]).T
+    return np.asarray(matches, dtype=int), np.array(unmatched_a, dtype=int), np.array(unmatched_b, dtype=int)
+    
+def center_distance(track, det):
+    diff = track.xywh[:2] - det.xywh[:2] 
+    # return L2 norm
+    return norm(diff, ord=2)
+    
+def ious(atlbrs, btlbrs):
+    """
+    Compute cost based on IoU
+    :type atlbrs: list[tlbr] | np.ndarray
+    :type atlbrs: list[tlbr] | np.ndarray
+
+    :rtype ious np.ndarray
+    """
+    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+    if ious.size == 0:
+        return ious
+
+    # ious = bbox_ious(
+    ious = bbox_overlaps(
+        np.ascontiguousarray(atlbrs, dtype=np.float32),
+        np.ascontiguousarray(btlbrs, dtype=np.float32)
+    )
+
+    return ious
+
+def ious_adaptive_height(atlbrs, btlbrs):
+    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+    if ious.size == 0:
+        return ious
+    b_tlbrs =  deepcopy(btlbrs)
+    _ious=[]
+    for a_tlbr in atlbrs:
+        ha_ = a_tlbr[3]-a_tlbr[1]
+        # Adjust the height in b_
+        for b_tlbr in b_tlbrs: b_tlbr[3] = b_tlbr[1]+ha_ 
+        a_ious = bbox_overlaps(
+            np.ascontiguousarray([a_tlbr], dtype=np.float32),
+            np.ascontiguousarray(b_tlbrs, dtype=np.float32)
+        )
+        _ious.append(a_ious)
+    _ious = np.concatenate(_ious,axis=0)
+    return _ious 
+
+def tlbr_expand(tlbr, scale=1.2):
+    w = tlbr[2] - tlbr[0]
+    h = tlbr[3] - tlbr[1]
+
+    half_scale = 0.5 * (scale-1)
+    tlbr[0] -= half_scale * w
+    tlbr[1] -= half_scale * h
+    tlbr[2] += half_scale * w
+    tlbr[3] += half_scale * h
+
+    return tlbr
+
+def iou_scores(atracks, btracks, adaptive_height=False):
+    """
+    Compute cost based on IoU
+    :type atracks: list[Tracklet]
+    :type btracks: list[Tracklet]
+
+    :rtype cost_matrix np.ndarray
+    """
+
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlbr for track in atracks]
+        btlbrs = [track.tlbr for track in btracks]
+    _ious = ious_adaptive_height(atlbrs, btlbrs) if adaptive_height else ious(atlbrs, btlbrs) 
+    return _ious
+
+def vel_consistent_scores(tracks,detections):
+    if (len(tracks)>0 and isinstance(tracks[0], np.ndarray)) or (len(detections) > 0 and isinstance(detections[0], np.ndarray)):
+        atlbrs = tracks
+        btlbrs = detections
+    else:
+        atlbrs = np.array([track.tlbr for track in tracks])
+        btlbrs = np.array([track.tlbr for track in detections])
+    if len(atlbrs) == 0 or len(btlbrs) == 0 :
+        return np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+
+    # Current velocity direction
+    vel_dir_cur = np.array([s.vel_dir for s in tracks])
+    vel_dir_cur = np.expand_dims(vel_dir_cur,axis=1)
+    # Compute expected velocity direction
+    xa,ya = atlbrs[:,0] + atlbrs[:,2]/2,atlbrs[:,1]
+    xb,yb = btlbrs[:,0] + btlbrs[:,2]/2,btlbrs[:,1]
+    dx = xb[None,:]-xa[:,None]
+    dy = yb[None,:]-ya[:,None]
+    norm = np.sqrt(dx*dx + dy*dy) +1e-6
+    vel_dir = np.stack([dx/norm,dy/norm],axis=2)
+
+    vel_scores = np.sum(vel_dir_cur*vel_dir,axis=2)
+    return vel_scores 
+
+def ort_consistent_scores(tracks,detections):
+    a_angles = np.array([track.angle for track in tracks])
+    b_angles = np.array([track.angle for track in detections])
+    if len(a_angles) == 0 or len(b_angles) == 0 :
+        return np.zeros((len(a_angles), len(b_angles)), dtype=np.float32)
+    diff = np.absolute(a_angles[:,None] - b_angles[None,:])
+    diff = np.minimum(diff,360-diff)
+    return np.cos(diff*DEG2RAD)
+ 
+def expand_iou_scores(atracks, btracks,expand_scale=1.2):
+    """
+    Compute cost based on IoU
+    :type atracks: list[Tracklet]
+    :type btracks: list[Tracklet]
+
+    :rtype cost_matrix np.ndarray
+    """
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [tlbr_expand(track.tlbr, scale=expand_scale) for track in atracks]
+        btlbrs = [tlbr_expand(track.tlbr, scale=expand_scale) for track in btracks]
+    _ious = ious(atlbrs, btlbrs)
+    return _ious
+        
+def embedding_distance(tracks, detections, metric='cosine'):
+    """
+    :param tracks: list[Tracklet]
+    :param detections: list[BaseTrack]
+    :param metric:
+    :return: cost_matrix np.ndarray
+    """
+
+    cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float32)
+    if cost_matrix.size == 0:
+        return cost_matrix
+    det_features = np.asarray([track.curr_feat for track in detections], dtype=np.float32)
+    track_features = np.asarray([track.smooth_feat for track in tracks], dtype=np.float32)
+
+    cost_matrix = np.maximum(0.0, cdist(track_features, det_features, metric))  # / 2.0  # Nomalized features
+    return cost_matrix
+
+def embedding_cosine_distance(tracklets, dets, feature_type='curr_feat', norm=False):
+    # embedding distance base on cosine similarity
+
+    # if tracklets and dets are already embeddings
+    if ((len(tracklets)>0 and isinstance(tracklets[0], np.ndarray)) 
+        or (len(dets) > 0 and isinstance(dets[0], np.ndarray))):
+        track_feats = tracklets
+        det_feats = dets
+    else:
+        # else, take the embedding from tracklet class
+        track_feats = [getattr(track, feature_type) for track in tracklets]
+        det_feats = [getattr(det, 'curr_feat') for det in dets]
+    # cal culate distance base on cosine
+    if len(track_feats) > 0 and len(det_feats) > 0:
+        track_feats = np.stack(track_feats)
+        det_feats = np.stack(det_feats)
+        if norm:
+            track_feats = track_feats / np.linalg.norm(track_feats, axis=1)[:, None]
+            det_feats = det_feats / np.linalg.norm(det_feats, axis=1)[:, None]
+        cosine = np.matmul(track_feats, det_feats.T)
+        return 1 - cosine
+    else:
+        return np.empty((len(track_feats), len(det_feats)), dtype=np.float64)
+
+def euclidean_distance_matrix(atracks, btracks):
+    """
+    Calculates the Euclidean distance matrix between middle top left points of atracks and btracks
+
+    :param atracks: A list of tracklets or atracks representing the first set of track data.
+    :type atracks: list[Tracklet] size N
+    :param btracks: A list of tracklets or atracks representing the second set of track data.
+    :type btracks: list[Tracklet] size M
+
+    :return: a dictionary
+            - distance_matrix size NxM: A numpy.ndarray representing the Euclidean distance between each pair of tracklets.
+            - width_atracks size Nx1: A numpy.ndarray indicating widths of boxes in atracks
+    :rtype: dict
+    """
+    if len(atracks) == 0 or len(btracks) == 0:
+        emtyp_matrix = np.empty((len(atracks), len(btracks)), dtype=np.float32)
+        return dict(distance_matrix=emtyp_matrix, width_atracks=emtyp_matrix)
+    
+    if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlbr for track in atracks]
+        btlbrs = [track.tlbr for track in btracks]
+    
+    # Extract the middle top points from atlbrs and btlbrs
+    atlbrs_midtop = np.array([((bbox[0] + bbox[2]) / 2, bbox[1]) for bbox in atlbrs])
+    btlbrs_midtop = np.array([((bbox[0] + bbox[2]) / 2, bbox[1]) for bbox in btlbrs])
+
+    # Calculate the Euclidean distance between each pair of middle top points
+    distance_matrix = np.sqrt(np.sum((atlbrs_midtop[:, np.newaxis] - btlbrs_midtop) ** 2, axis=-1))
+
+    # evaluate distance_matrix valid or not
+    width_atracks = np.array([(bbox[2] - bbox[0]) / 2 for bbox in atlbrs])[:, None]
+
+    return dict(distance_matrix=distance_matrix, width_atracks=width_atracks)
+
+# def gate_cost_matrix(kf, cost_matrix, tracks, detections, only_position=False):
+#     if cost_matrix.size == 0:
+#         return cost_matrix
+#     gating_dim = 2 if only_position else 4
+#     gating_threshold = kalman_filter.chi2inv95[gating_dim]
+#     # measurements = np.asarray([det.to_xyah() for det in detections])
+#     measurements = np.asarray([det.to_xywh() for det in detections])
+#     for row, track in enumerate(tracks):
+#         gating_distance = kf.gating_distance(
+#             track.mean, track.covariance, measurements, only_position)
+#         cost_matrix[row, gating_distance > gating_threshold] = np.inf
+#     return cost_matrix
+
+
+# def fuse_motion(kf, cost_matrix, tracks, detections, only_position=False, lambda_=0.98):
+#     if cost_matrix.size == 0:
+#         return cost_matrix
+#     gating_dim = 2 if only_position else 4
+#     gating_threshold = kalman_filter.chi2inv95[gating_dim]
+#     # measurements = np.asarray([det.to_xyah() for det in detections])
+#     measurements = np.asarray([det.to_xywh() for det in detections])
+#     for row, track in enumerate(tracks):
+#         gating_distance = kf.gating_distance(
+#             track.mean, track.covariance, measurements, only_position, metric='maha')
+#         cost_matrix[row, gating_distance > gating_threshold] = np.inf
+#         cost_matrix[row] = lambda_ * cost_matrix[row] + (1 - lambda_) * gating_distance
+#     return cost_matrix
+
+
+def fuse_iou(cost_matrix, tracks, detections):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    reid_sim = 1 - cost_matrix
+    iou_dist = 1 -iou_scores(tracks, detections)
+    iou_sim = 1 - iou_dist
+    fuse_sim = reid_sim * (1 + iou_sim) / 2
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    #fuse_sim = fuse_sim * (1 + det_scores) / 2
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost
+
+def fuse_score(cost_matrix, detections):
+    if cost_matrix.size == 0:
+        return cost_matrix
+    iou_sim = 1 - cost_matrix
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    fuse_sim = iou_sim * det_scores
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost
+
+def solider_distance(tracklets, dets, feature_type='curr_feat', norm=True):
+    # if tracklets and dets are already embeddings
+    if ((len(tracklets)>0 and isinstance(tracklets[0], np.ndarray)) 
+        or (len(dets) > 0 and isinstance(dets[0], np.ndarray))):
+        track_feats = tracklets
+        det_feats = dets
+    else:
+        # else, take the embedding from tracklet class
+        track_feats = [getattr(track, feature_type) for track in tracklets]
+        det_feats = [getattr(det, 'curr_feat') for det in dets]
+    # cal culate distance base on cosine
+    if len(track_feats) > 0 and len(det_feats) > 0:
+        track_feats = np.stack(track_feats)
+        det_feats = np.stack(det_feats)
+
+        qf = torch.from_numpy(track_feats).to('cuda')
+        gf = torch.from_numpy(det_feats).to('cuda')
+        if norm:
+            qf = torch.nn.functional.normalize(qf, dim=1, p=2)
+            gf = torch.nn.functional.normalize(gf, dim=1, p=2)
+        m = qf.shape[0]
+        n = gf.shape[0]
+        dist_mat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
+                torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
+        dist_mat.addmm_(1, -2, qf, gf.t())
+        return dist_mat.cpu().numpy()
+    else:
+        return np.empty((len(track_feats), len(det_feats)), dtype=np.float64)
+
+def best_area_solider_distance(tracklets, dets, norm=False):
+
+    if len(tracklets) > 0 and len(dets) > 0:
+        # if tracklets and dets are already embeddings
+        if ((len(tracklets)>0 and isinstance(tracklets[0], np.ndarray)) 
+            or (len(dets) > 0 and isinstance(dets[0], np.ndarray))):
+            track_feats = tracklets
+            det_feats = dets
+        else:
+            # else, take the embedding from tracklet class
+            det_feats = [getattr(det, 'curr_feat') for det in dets]
+            
+            # take all best match feat based on det box area
+            track_feats = []
+            for track_idx, tracklet in enumerate(tracklets):
+                for det_idx, det in enumerate(dets):
+                    best_match_results = tracklet.best_area_feat(det._tlwh)
+                    feat = best_match_results['feat']
+                    track_feats.append(feat)
+        # MxNxC
+        track_feats = np.stack(track_feats).reshape(len(tracklets), len(dets), -1)
+        # NxC
+        det_feats = np.stack(det_feats)
+
+        qf = torch.from_numpy(track_feats).to('cuda')
+        gf = torch.from_numpy(det_feats).to('cuda')
+        
+        if norm:
+            qf = torch.nn.functional.normalize(qf, dim=-1, p=2)
+            gf = torch.nn.functional.normalize(gf, dim=-1, p=2)
+        
+        # m = qf.shape[0]
+        # n = gf.shape[0]
+        # dist_mat = torch.pow(qf, 2).sum(dim=-1, keepdim=False) + \
+        #         torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
+        # dist_mat.addmm_(1, -2, qf, gf.t())
+        
+        gf = gf.unsqueeze(0)
+        dist_matff = ((qf-gf)**2).sum(-1) #.sqrt()
+
+        return dist_matff.cpu().numpy()
+    else:
+        return np.empty((len(tracklets), len(dets)), dtype=np.float64)

models/models/trackers/reid_parallel_tracker/core/tracklet.py

@@ -0,0 +1,427 @@
+
+import numpy as np
+from collections import deque
+from .matching import iou_scores
+from .basetrack import BaseTrack, TrackState
+from .kalman_filter import KalmanFilter
+from copy import deepcopy
+import cv2
+
+RAD2DEG = 180.0/np.pi
+
+
+def add_stracks(tlista, tlistb):
+    exists = {}
+    res = []
+    for t in tlista:
+        exists[t.track_id] = 1
+        res.append(t)
+    for t in tlistb:
+        tid = t.track_id
+        if not exists.get(tid, 0):
+            exists[tid] = 1
+            res.append(t)
+    return res
+
+
+def subtract_stracks(tlista, tlistb):
+    stracks = {}
+    for t in tlista:
+        stracks[t.track_id] = t
+    for t in tlistb:
+        tid = t.track_id
+        if stracks.get(tid, 0):
+            del stracks[tid]
+    return list(stracks.values())
+
+
+def remove_duplicate_stracks(stracksa, stracksb):
+    pscores = iou_scores(stracksa, stracksb)
+    pairs = np.where(pscores > 0.85)
+    dupa, dupb = list(), list()
+    for p, q in zip(*pairs):
+        timep = stracksa[p].frame_id - stracksa[p].start_frame
+        timeq = stracksb[q].frame_id - stracksb[q].start_frame
+        if timep > timeq:
+            dupb.append(q)
+        else:
+            dupa.append(p)
+    resa = [t for i, t in enumerate(stracksa) if not i in dupa]
+    resb = [t for i, t in enumerate(stracksb) if not i in dupb]
+    return resa, resb
+
+
+class Tracklet(BaseTrack):
+    shared_kalman = KalmanFilter()
+
+    def __init__(self, tlwh, score, cls, angle=None,
+                 feat=None, feat_history=50,
+                 enable_buffer=True, obj_img=None):
+        self._tlwh = np.asarray(tlwh, dtype=np.float32)
+        # wait activate
+        self.score = score
+        self.cls = -1
+
+        self.kalman_filter = None
+        self.mean, self.covariance = None, None
+        # add trajectory, last observation, last_mean,last_cov
+        self.trajectory = None  # buffer to keep track of trajectory
+        # last observation before lost (frame_id,det_bbox)
+        self.last_observation = None
+        # last mean,last observation before lost
+        self.last_mean, self.last_covariance = None, None
+        self.last_frame_id = None  # last frame_id before lost
+
+        self.is_activated = False
+
+        self.cls_hist = []  # (cls id, freq)
+        self.update_cls(cls, score)
+
+        self.tracklet_len = 0
+
+        self.angle = angle
+
+        # reid
+        norm_feat = self.norm_feat(feat) if feat is not None else None
+        self.smooth_feat = norm_feat
+        self.curr_feat = norm_feat
+        self.feat_momentum = 0.9
+        self.obj_img = obj_img
+        self.best_det_feat = norm_feat
+        self.best_obj_img = obj_img
+
+        # buffer of feature
+        self.enable_buffer = enable_buffer
+        self.feat_buffer = list()
+        self.box_buffer = list()  # tlwh
+        self.obj_img_buffer = list()
+
+        self.lost_frame_num = 0
+
+        self.history_info = []
+
+        self.active_frames = []
+
+    def update_active_frames(self, frame):
+        if not isinstance(frame, list):
+            self.active_frames.append(frame)
+        else:
+            self.active_frames.extend(frame)
+
+    def set_lost_frame_num(self):
+        self.lost_frame_num += 1
+
+    def reset_lost_frame_num(self):
+        self.lost_frame_num = 0
+
+    @staticmethod
+    def norm_feat(feat):
+        feat /= np.linalg.norm(feat)
+        return feat
+
+    def buffer_areas(self):
+        buf_areas = np.array([_bbox[2]*_bbox[3] for _bbox in self.box_buffer])
+        return buf_areas
+
+    def update_reid_buffer(self, new_track):
+        N = 15
+        # check wheather the obj image size and ratio is different from current image sizes
+        # --> update the buffer if the condition is True
+        update = False
+        if len(self.box_buffer) == 0:
+            update = True
+        else:
+            new_tlwh = new_track._tlwh
+            # calculate the area of boxes
+            buf_areas = self.buffer_areas()
+            new_area = new_tlwh[2] * new_tlwh[3]
+
+            # Calculate the percentage value
+            percentage = N / 100
+
+            # Check if new_area is N% bigger or smaller than every box in buf_areas
+            is_bigger = np.all(new_area >= (1 + percentage) * buf_areas)
+            is_smaller = np.all(new_area <= (1 - percentage) * buf_areas)
+
+            if is_bigger or is_smaller:
+                update = True
+
+        if update:
+            self.feat_buffer.append(new_track.curr_feat)
+            self.box_buffer.append(new_track._tlwh)
+            self.obj_img_buffer.append(new_track.obj_img)
+
+    def best_area_feat(self, new_box):
+        # new_box: tlwh
+        # return a best match feature and obj image base on min area
+        if len(self.feat_buffer) == 0 or not self.enable_buffer:
+            return dict(
+                feat=self.curr_feat,
+                obj_img=self.obj_img)
+
+        new_box_area = new_box[2] * new_box[3]
+        buf_areas = self.buffer_areas()
+        dif = np.abs(buf_areas - new_box_area)
+
+        best_match_idx = np.argmin(dif)
+
+        best_match_results = dict(
+            feat=self.feat_buffer[best_match_idx],
+            obj_img=self.obj_img_buffer[best_match_idx]
+        )
+        return best_match_results
+
+    def update_features(self, new_track):
+        feat = new_track.curr_feat
+        self.smooth_feat = (self.feat_momentum * self.smooth_feat
+                            + (1 - self.feat_momentum) * feat)
+        self.smooth_feat = self.norm_feat(self.smooth_feat)
+
+        if self.enable_buffer:
+            # update feature buffer based on boxsize
+            self.update_reid_buffer(new_track)
+
+        # update current feature
+        self.curr_feat = feat
+        self.obj_img = new_track.obj_img
+
+        # update best reid feature base on detection score
+        if new_track.score >= self.score:
+            self.best_det_feat = feat
+            self.best_obj_img = new_track.obj_img
+
+    def update_cls(self, cls, score):
+        if len(self.cls_hist) > 0:
+            max_freq = 0
+            found = False
+            for c in self.cls_hist:
+                if cls == c[0]:
+                    c[1] += score
+                    found = True
+
+                if c[1] > max_freq:
+                    max_freq = c[1]
+                    self.cls = c[0]
+            if not found:
+                self.cls_hist.append([cls, score])
+                self.cls = cls
+        else:
+            self.cls_hist.append([cls, score])
+            self.cls = cls
+
+    def update_angle(self, angle, score):
+        self.angle = (1-score)*self.angle + score*angle
+
+    @staticmethod
+    def multi_predict(stracks):
+        if len(stracks) > 0:
+            multi_mean = np.asarray([st.mean.copy() for st in stracks])
+            multi_covariance = np.asarray([st.covariance for st in stracks])
+            for i, st in enumerate(stracks):
+                if st.state != TrackState.Tracked:
+                    multi_mean[i][6] = 0
+                    multi_mean[i][7] = 0
+            multi_mean, multi_covariance = Tracklet.shared_kalman.multi_predict(
+                multi_mean, multi_covariance)
+            for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+                stracks[i].mean = mean
+                stracks[i].covariance = cov
+
+    @staticmethod
+    def multi_gmc(stracks, Hmat=np.eye(3, 3)):
+        # This approach uses Homography matrix
+        if len(stracks) > 0:
+            multi_mean = [st.mean.reshape(-1, 1, 2) for st in stracks]
+            multi_covariance = np.asarray([st.covariance for st in stracks])
+
+            R, T = Hmat[:2, :2], Hmat[:2, 2]
+            R8x8 = np.kron(np.eye(4, dtype=float), R)
+            H = deepcopy(Hmat)
+            H[:2, 2] = 0  # remove translation
+
+            for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+                # w = mean.reshape(4,2).dot(V)
+                # R_ = [R/w_i for w_i in w]
+                # R8x8 = block_diag(*R_)
+                mean = cv2.perspectiveTransform(mean, H)
+                mean = mean.reshape(-1)
+                mean[:2] += T
+                cov = R8x8.dot(cov).dot(R8x8.transpose())
+
+                stracks[i].mean = mean
+                stracks[i].covariance = cov
+
+                last_obs, last_frame_id = stracks[i].last_observation,  stracks[i].last_frame_id
+                last_mean, last_cov = stracks[i].last_mean, stracks[i].last_covariance
+
+                if last_mean is not None:
+                    # Update the last observation
+                    last_obs = cv2.perspectiveTransform(
+                        last_obs.reshape(-1, 1, 2), H)
+                    last_obs = last_obs.reshape(-1)
+                    last_obs[:2] += T
+                    last_mean = cv2.perspectiveTransform(
+                        last_mean.reshape(-1, 1, 2), H)
+                    last_mean = last_mean.reshape(-1)
+                    last_mean[:2] += T
+                    last_cov = R8x8.dot(last_cov).dot(R8x8.transpose())
+                    stracks[i].last_observation = last_obs
+                    stracks[i].last_mean = last_mean
+                    stracks[i].last_covariance = last_cov
+                    # save the trajectory
+                    traj = last_mean[:4]
+                    stracks[i].trajectory = np.append(traj, last_frame_id)
+
+    def activate(self, kalman_filter, frame_id):
+        """Start a new tracklet"""
+        self.kalman_filter = kalman_filter
+        self.track_id = self.next_id()
+
+        xywh = self.tlwh_to_xywh(self._tlwh)
+        self.mean, self.covariance = self.kalman_filter.initiate(xywh)
+
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        if frame_id == 1:
+            self.is_activated = True
+        self.frame_id = frame_id
+        self.start_frame = frame_id
+
+        # Save trajectory
+        self.last_frame_id = frame_id
+        self.last_observation = deepcopy(xywh)
+        self.last_mean = deepcopy(self.mean)
+        self.last_covariance = deepcopy(self.covariance)
+
+        self.set_history_info(
+            history_info=[dict(frame_id=frame_id, track_id=self.track_id), ],
+            append_first=False)
+
+        self.update_active_frames(frame_id)
+
+    def re_activate(self, new_track, frame_id, new_id=False):
+        xywh = self.tlwh_to_xywh(new_track.tlwh)
+        self.mean, self.covariance = self.kalman_filter.update(self.mean,
+                                                               self.covariance, xywh,
+                                                               new_track.score)
+        if new_track.curr_feat is not None:
+            self.update_features(new_track)
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        self.is_activated = True
+        self.frame_id = frame_id
+        if new_id:
+            self.track_id = self.next_id()
+        self.score = new_track.score
+
+        self.update_cls(new_track.cls, new_track.score)
+        if self.angle is not None:
+            self.update_angle(new_track.angle, new_track.score)
+
+        # Save trajectory
+        self.last_frame_id = frame_id
+        self.last_observation = deepcopy(xywh)
+        self.last_mean = deepcopy(self.mean)
+        self.last_covariance = deepcopy(self.covariance)
+        self.set_history_info(
+            history_info=[dict(frame_id=frame_id, track_id=self.track_id), ],
+            append_first=False)
+        self.update_active_frames(frame_id)
+
+    def update(self, new_track, frame_id):
+        """
+        Update a matched track
+        :type new_track: Tracklet
+        :type frame_id: int
+        :type update_feature: bool
+        :return:
+        """
+        self.frame_id = frame_id
+        self.tracklet_len += 1
+
+        new_xywh = self.tlwh_to_xywh(new_track.tlwh)
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, new_xywh, new_track.score)
+
+        # update reid feature
+        if new_track.curr_feat is not None:
+            self.update_features(new_track)
+
+        self.state = TrackState.Tracked
+        self.is_activated = True
+
+        self.score = new_track.score
+        self.update_cls(new_track.cls, new_track.score)
+
+        if self.angle is not None:
+            self.update_angle(new_track.angle, new_track.score)
+
+        # Save trajectory
+        self.last_frame_id = frame_id
+        self.last_observation = deepcopy(new_xywh)
+        self.last_mean = deepcopy(self.mean)
+        self.last_covariance = deepcopy(self.covariance)
+        if self.is_activated:
+            self.set_history_info(
+                history_info=[
+                    dict(frame_id=frame_id, track_id=self.track_id), ],
+                append_first=False)
+        self.update_active_frames(frame_id)
+
+    def smooth_update(self, new_track, frame_id):
+        '''
+            As introduced in OC-SORT
+        '''
+        xywh = self.tlwh_to_xywh(new_track.tlwh)
+        # Interpolate update
+        if self.last_frame_id is not None:
+            num_missing_steps = frame_id - self.last_frame_id
+            if num_missing_steps > 1:
+                delta = (xywh - self.last_observation)/num_missing_steps
+                interpolate_tracks = [self.last_observation +
+                                      i*delta for i in range(1, num_missing_steps)]
+                mean_i, covariance_i = self.last_mean, self.last_covariance
+                for new_xywh in interpolate_tracks:
+                    # Predict
+                    mean_i, covariance_i = self.kalman_filter.predict(
+                        mean_i, covariance_i)
+                    # Update
+                    mean_i, covariance_i = self.kalman_filter.update(
+                        mean_i, covariance_i, new_xywh, new_track.score)
+                # the prediction step before last update
+                self.mean, self.covariance = self.kalman_filter.predict(
+                    mean_i, covariance_i)
+
+        # Normal update
+        self.update(new_track, frame_id)
+
+    @property
+    def velocity(self):
+        vx = self.mean[4]  # horizontal velocity
+        vy = self.mean[5]  # vertical velocity
+        vh = self.mean[7]  # height velocity
+        return [vx, vy, vh]
+
+    @property
+    def vel_dir(self):
+        if self.trajectory is None:
+            dist = self.mean[4:6]  # velocity
+        else:
+            p2 = self.mean[:2]
+            p1 = self.trajectory[:2]  # (x,y,w,h,frame_id)
+            dist = p2-p1
+        norm = np.linalg.norm(dist) + 1e-6
+        return dist/norm
+
+    def common_active_frames(self, track):
+        lst3 = [value for value in self.active_frames if value in track.active_frames]
+        return len(lst3) > 0
+
+    def set_history_info(self, history_info, append_first=False):
+        if append_first:
+            self.history_info = history_info + self.history_info
+        else:
+            self.history_info = self.history_info + history_info
+
+    def get_history_info(self):
+        return self.history_info

models/models/trackers/reid_parallel_tracker/matchers/__init__.py

@@ -0,0 +1,4 @@
+from .distances import DistSOLIDER
+from .base_matchers import SimMatcher
+from .single_stage_matcher import SingleStageMatcher
+from .prioritize_reid_matcher import PrioritizeReidMatcher

models/models/trackers/reid_parallel_tracker/matchers/base_matchers.py

@@ -0,0 +1,39 @@
+import numpy as np
+from typing import List, Tuple, Dict
+from .distances import DistCosine
+from ..core.matching import linear_assignment, topk_assignment
+from ..core.tracklet import Tracklet
+
+class SimMatcher():
+    """The baseline matcher that only use linear_assignment to match tracklets with detection boxes
+    """
+    def __init__(self,
+                 dist_cfg: Dict,
+                 match_thr: float):
+        self.dist = DistCosine(**dist_cfg)
+        self.match_thr = match_thr
+        
+    def __call__(self, 
+                 tracks: List[Tracklet], 
+                 dets: List[Tracklet]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """ Associate with tracklets with detection boxes"""
+        dists = self.dist(tracks, dets)
+        matches_idx, unmatched_tracks_idx, unmatched_dets_idx = linear_assignment(dists, thresh=self.match_thr)
+        # matches_idx, unmatched_tracks_idx, unmatched_dets_idx = topk_assignment(dists, thresh=self.match_thr, topk=2)
+        return matches_idx, unmatched_tracks_idx, unmatched_dets_idx
+    
+    def matching_dists(self, tracks: List[Tracklet],
+                 dets: List[Tracklet]) -> np.ndarray:
+        """ Compute the distance between tracklets and detections"""
+        return self.dist(tracks, dets)
+    
+    def matching_scores(self, tracks: List[Tracklet], 
+                 dets: List[Tracklet]) -> np.ndarray:
+        """ Compute the matching scores between tracklets and detections"""
+        return self.dist.matching_scores(tracks, dets)
+        
+    def assign(self, distances: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        return linear_assignment(distances, thresh=self.match_thr)
+    
+    def __repr__(self):
+        return f"SimMatcher(dist={self.dist}, match_thr={self.match_thr})"

models/models/trackers/reid_parallel_tracker/matchers/distances.py

@@ -0,0 +1,132 @@
+from typing import List, Dict
+from ..core.tracklet import Tracklet
+from ..core.matching import (embedding_cosine_distance,
+                             euclidean_distance_matrix, solider_distance,
+                             best_area_solider_distance)
+import numpy as np 
+
+class BaseDist():
+    def __init__(self, 
+                 weight=None,
+                 weight_type="dot"):
+        self.weight = weight
+        self.weight_type = weight_type.lower()
+        assert self.weight_type in ["dot","power"], "only support weight type of dot(multiply w*dist), power (dist^w)"
+        
+    def weighted_dist(self, dists):
+        if self.weight is not None:
+            dists = self.weight*dists if self.weight_type == "dot" else np.power(dists,self.weight)
+        return dists
+    
+
+class DistCosine(BaseDist):
+    # ReID distance
+    def __init__(self,
+                 location_aware=dict(search_radius_scale=10),
+                 feature_type='curr_feat',
+                 norm=False,
+                 **kwargs):
+        """
+            - use_location_aware: each object will only be matched with nearby objects 
+                                  which is inside a circle with R = Width * SRS
+        """
+        self.location_aware = location_aware
+        self.feature_type = feature_type
+        self.norm = norm
+        self.use_adaptive_search_radius = ('early_track_lost_frame' in self.location_aware) \
+                                            and ('early_lost_track_search_radius_scale' in self.location_aware) \
+                                            and ('active_track_search_radius_scale' in self.location_aware)
+        assert self.feature_type in ['curr_feat', 'smooth_feat', 'best_det_feat']
+        super().__init__(**kwargs)
+   
+    def __call__(self,
+                 tracklets:  List[Tracklet],
+                 dets:  List[Tracklet]) -> np.ndarray:
+        # calculate distance base on embedding
+        dists = embedding_cosine_distance(tracklets, dets,
+                                          self.feature_type,
+                                          norm=self.norm)
+        # if use localtion aware
+        if len(self.location_aware) and len(tracklets):
+            distance_results = euclidean_distance_matrix(tracklets, dets)
+            distance_matrix = distance_results['distance_matrix']
+            width_atracks = distance_results['width_atracks']
+
+            if not self.use_adaptive_search_radius:
+                # fixed searching radius scale 
+                invalid_reid_search_regions = distance_matrix >= width_atracks*self.location_aware['search_radius_scale']
+            else:
+                # adaptive search radius scale
+                lost_frame_nums = np.array([_track.lost_frame_num for _track in tracklets])
+                is_active = lost_frame_nums == 0
+                is_early_lost_track = (lost_frame_nums <= self.location_aware['early_track_lost_frame']) * (lost_frame_nums > 0)
+                is_long_lost_track = lost_frame_nums > self.location_aware['early_track_lost_frame']
+            
+                lost_frame_nums[is_active] = self.location_aware['active_track_search_radius_scale']
+                lost_frame_nums[is_early_lost_track] = self.location_aware['early_lost_track_search_radius_scale']
+                lost_frame_nums[is_long_lost_track] = self.location_aware['long_lost_track_search_radius_scale']
+                lost_frame_nums = lost_frame_nums.reshape(-1, 1)
+                invalid_reid_search_regions = distance_matrix >= width_atracks*lost_frame_nums
+            # update distance threshold
+            dists[invalid_reid_search_regions] = 1e4
+
+        return self.weighted_dist(dists)
+
+    
+    def matching_scores(self,
+                 tracklets:  List[Tracklet],
+                 dets:  List[Tracklet]) -> np.ndarray:
+        """ convienient function to return the matching scores instead of the distance"""
+        dists = embedding_cosine_distance(tracklets, dets,
+                                          self.feature_type,
+                                          norm=self.norm)
+        matching_scores = 1 - dists
+        return matching_scores 
+
+
+class DistSOLIDER(BaseDist):
+    # ReID distance
+    def __init__(self,
+                 location_aware=dict(search_radius_scale=10),
+                 feature_type='curr_feat',
+                 norm=False,
+                 **kwargs):
+        """
+            - use_location_aware: each object will only be matched with nearby objects 
+                                  which is inside a circle with R=Width * search_radius_scale
+        """
+        self.location_aware = location_aware
+        self.feature_type = feature_type
+        self.norm = norm
+        assert self.feature_type in ['curr_feat', 'smooth_feat', 'best_area_feat', 'best_det_feat']
+        super().__init__(**kwargs)
+   
+    def __call__(self,
+                 tracklets:  List[Tracklet],
+                 dets:  List[Tracklet]) -> np.ndarray:
+        # calculate distance base on embedding
+        if self.feature_type == "best_area_feat":
+            dists = best_area_solider_distance(tracklets, dets, self.norm)
+        else:
+            dists = solider_distance(tracklets, dets,
+                                            self.feature_type,
+                                            norm=self.norm)
+        # if use localtion aware
+        if self.location_aware is not None:
+            distance_results = euclidean_distance_matrix(tracklets, dets)
+            distance_matrix = distance_results['distance_matrix']
+            width_atracks = distance_results['width_atracks']
+            invalid_reid_search_regions = distance_matrix >= width_atracks*self.location_aware['search_radius_scale']
+            try:
+                dists[invalid_reid_search_regions] = 1.0
+            except:
+                import ipdb; ipdb.set_trace()
+        return self.weighted_dist(dists)
+
+    
+    def matching_scores(self,
+                 tracklets:  List[Tracklet],
+                 dets:  List[Tracklet]) -> np.ndarray:
+        """ convienient function to return the matching scores instead of the distance"""
+        NotImplementedError('Not implemented')
+

models/models/trackers/reid_parallel_tracker/matchers/prioritize_reid_matcher.py

@@ -0,0 +1,224 @@
+import numpy as np
+from copy import deepcopy
+from typing import List, Tuple, Dict
+from ..core.matching import linear_assignment, topk_assignment
+from ..core.tracklet import Tracklet
+from ..core.tracklet import TrackState
+from .base_matchers import SimMatcher
+from .distances import DistCosine
+
+class PrioritizeReidMatcher():
+    def __init__(self,
+                 reid_distance,
+                 iou_distance,
+                 match_with_reid_thr,
+                 islost_match_thr,
+                 isactive_match_thr):
+        
+        self.reid_dist = DistCosine(**reid_distance)
+        self.iou_dist = DistCosine(**iou_distance)
+        self.match_with_reid_thr = match_with_reid_thr
+        self.islost_match_thr = islost_match_thr
+        self.isactive_match_thr = isactive_match_thr
+
+    def __call__(self, 
+                 tracks: List[Tracklet], 
+                 dets: List[Tracklet]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """ Associate with tracklets with detection boxes"""
+
+        # 1. Match with ReID dist <0.15 
+        dist_reid = self.reid_dist(tracks, dets)
+        matches_idx_1, unmatched_tracks_idx_1, unmatched_dets_idx_1 = self.assign(dist_reid,
+                                                                                        thresh=self.match_with_reid_thr)
+
+        # split unmatch track, det
+        unmatched_tracks = [tracks[idx] for idx in unmatched_tracks_idx_1]
+        unmatched_dets = [dets[i] for i in unmatched_dets_idx_1]
+
+        # split active and lost tracks
+        unmatched_lost_tracks = []
+        unmatched_lost_tracks_idx = []
+        unmatched_active_tracks = []
+        unmatched_active_tracks_idx = []
+        for idx, _track in enumerate(unmatched_tracks):
+            if _track.state == TrackState.Lost:
+                unmatched_lost_tracks.append(_track)
+                unmatched_lost_tracks_idx.append(unmatched_tracks_idx_1[idx])
+            else:
+                unmatched_active_tracks.append(_track)
+                unmatched_active_tracks_idx.append(unmatched_tracks_idx_1[idx])
+
+        # 2. is lost; lost object expect lower IoU threshold lower ReID
+        matches_idx_2, unmatched_tracks_idx_2, unmatched_dets_idx_2 = self.reid_and_iou_matching(
+                                                unmatched_lost_tracks,
+                                                unmatched_lost_tracks_idx,
+                                                unmatched_dets,
+                                                unmatched_dets_idx_1,
+                                                self.islost_match_thr)
+
+
+        # 3. Active object expect higher IoU threshold
+        # split remain detection
+        from2_unmatched_dets = [dets[i] for i in unmatched_dets_idx_2]
+        from2_unmatched_dets_idx = unmatched_dets_idx_2
+        matches_idx_3, unmatched_tracks_idx_3, unmatched_dets_idx_3 = self.reid_and_iou_matching(
+                                                unmatched_active_tracks,
+                                                unmatched_active_tracks_idx,
+                                                from2_unmatched_dets,
+                                                from2_unmatched_dets_idx,
+                                                self.isactive_match_thr)
+        
+        # merge result
+        matches_idx = matches_idx_1.tolist() + matches_idx_2.tolist() + matches_idx_3.tolist()
+        unmatched_tracks_idx = []
+        unmatched_dets_idx = []
+
+        if len(matches_idx):
+            matches_idx = np.array(matches_idx)
+        else:
+            matches_idx = np.empty((0, 2), dtype=np.int64)
+        
+        # update unmatch tracks
+        for idx in range(len(tracks)):
+            if not idx in matches_idx[:, 0]:
+                unmatched_tracks_idx.append(idx)
+        unmatched_tracks_idx = np.array(unmatched_tracks_idx)
+        # update unmatch dets
+        for idx in range(len(dets)):
+            if not idx in matches_idx[:, 1]:
+                unmatched_dets_idx.append(idx)
+        unmatched_dets_idx = np.array(unmatched_dets_idx)
+
+        return matches_idx, unmatched_tracks_idx, unmatched_dets_idx
+
+    def calculate_occluded_ratio(self, dets):
+        det_ious = 1 - self.iou_dist(dets, dets)
+        np.fill_diagonal(det_ious, 0)
+        det_scores = np.array([det.score for det in dets])
+        det_score_matrix = (det_scores[:, None] < det_scores[None, :]).astype(np.float32)
+        occluded_ratio_matrix = det_ious * det_score_matrix
+        # TODO: Take area --> done, same performance.
+        if len(occluded_ratio_matrix):
+            occluded_ratio = np.max(occluded_ratio_matrix, 1)
+            return occluded_ratio
+        else:
+            return []
+
+    def reid_and_iou_matching(self,
+                              unmatched_tracks,
+                              unmatched_tracks_idx,
+                              unmatched_dets,
+                              unmatched_dets_idx,
+                              match_conditions):
+        # 2. is lost; lost object expect lower IoU threshold lower ReID
+        lost_tracks_reid_dist = self.reid_dist(unmatched_tracks, unmatched_dets)
+        lost_tracks_iou_dist = self.iou_dist(unmatched_tracks, unmatched_dets)
+
+        # reweight reid
+        occluded_ratio = self.calculate_occluded_ratio(unmatched_dets)
+
+        # if len(occluded_ratio):
+            ### reid_score *= e^(-occluded_ratio)
+            # reid_reweighting = np.exp(-occluded_ratio)
+            # lost_tracks_reid_dist = lost_tracks_reid_dist * reid_reweighting
+            
+            ### (x^2)/3
+            # lost_tracks_reid_dist = lost_tracks_reid_dist + (occluded_ratio**2)/3
+
+        occluded_thr = 0.3
+        used_occluded = False
+        
+        # cascade reid matching
+        if len(occluded_ratio):
+            # if occluded: Match non occluded object first, then match occluded one
+            is_occluded_idx = np.nonzero(occluded_ratio > occluded_thr)[0]
+            isnot_occluded_idx = np.nonzero(occluded_ratio <= occluded_thr)[0]
+            if len(is_occluded_idx) and len(unmatched_tracks):
+                used_occluded = True
+
+                # matching visible det objects
+                vis_lost_tracks_reid_dist = deepcopy(lost_tracks_reid_dist)
+                vis_lost_tracks_reid_dist[:, is_occluded_idx] = 1e4
+                vis_matched_idx, vis_unmatched_tracks_idx, vis_unmatched_dets_idx = self.assign(vis_lost_tracks_reid_dist,
+                                                                                                        thresh=match_conditions['reid_thr'])
+                # matching occluded det objects
+                occluded_lost_tracks_reid_dist = deepcopy(lost_tracks_reid_dist)
+                # if tracklet already match --> ignore
+                occluded_lost_tracks_reid_dist[vis_matched_idx[:, 0], :] = 1e4
+                # update visible positions = 1e4
+                occluded_lost_tracks_reid_dist[:, isnot_occluded_idx] = 1e4
+                # matching
+                occ_matched_idx, occ_unmatched_tracks_idx, occ_unmatched_dets_idx = self.assign(occluded_lost_tracks_reid_dist,
+                                                                                            thresh=match_conditions['reid_occluded_thr'])
+
+                matches_idx_2_1 = np.concatenate([vis_matched_idx, occ_matched_idx])
+                unmatched_tracks_idx_2_1 = np.concatenate([vis_unmatched_tracks_idx, occ_unmatched_tracks_idx])
+                unmatched_dets_idx_2_1 = np.concatenate([vis_unmatched_dets_idx, occ_unmatched_dets_idx])
+
+        if not used_occluded:
+            matches_idx_2_1, unmatched_tracks_idx_2_1, unmatched_dets_idx_2_1 = self.assign(lost_tracks_reid_dist,
+                                                                                            thresh=match_conditions['reid_thr'])
+        if match_conditions['iou_thr'] >= 0:
+            matches_idx_2_2, unmatched_tracks_idx_2_2, unmatched_dets_idx_2_2 = self.assign(lost_tracks_iou_dist,
+                                                                                            thresh=match_conditions['iou_thr'])
+            # take the intersection of matches_idx_2_1 and matches_idx_2_2
+            matches_idx_2_1 = matches_idx_2_1.tolist()
+            matches_idx_2_2 = matches_idx_2_2.tolist()
+            
+            # merge the two result
+            merged_matches_idx = list()
+            for _track in matches_idx_2_1:
+                if _track in matches_idx_2_2:
+                    merged_matches_idx.append(_track)
+        else:
+            merged_matches_idx = matches_idx_2_1.tolist()
+        
+        # convert local merge idx to global ones
+        global_matches_idx = []
+        for idx, match in enumerate(merged_matches_idx):
+            global_matches_idx.append([unmatched_tracks_idx[match[0]], unmatched_dets_idx[match[1]]])
+        if len(global_matches_idx):
+            global_matches_idx = np.array(global_matches_idx)
+        else:
+            global_matches_idx = np.empty((0, 2), dtype=np.int64)
+        
+        # convert to input (global) indices        
+        if len(merged_matches_idx):
+            merged_matches_idx = np.array(merged_matches_idx)
+        else:
+            merged_matches_idx = np.empty((0, 2), dtype=np.int64)
+
+        global_unmatched_tracks_idx = []
+        global_unmatched_dets_idx = []
+
+        # update unmatch track
+        for idx, _track_idx in enumerate(unmatched_tracks_idx):
+            if idx in merged_matches_idx[:, 0]:
+                continue
+            else:
+                global_unmatched_tracks_idx.append(_track_idx)      
+
+        # update unmatch det
+        for idx, _det_idx in enumerate(unmatched_dets_idx):
+            if idx in merged_matches_idx[:, 1]:
+                continue
+            else:
+                global_unmatched_dets_idx.append(_det_idx)     
+        
+        global_unmatched_tracks_idx = np.array(global_unmatched_tracks_idx)
+        global_unmatched_dets_idx = np.array(global_unmatched_dets_idx)
+        return global_matches_idx, global_unmatched_tracks_idx, global_unmatched_dets_idx
+    
+    def matching_dists(self, tracks: List[Tracklet],
+                 dets: List[Tracklet]) -> np.ndarray:
+        """ Compute the distance between tracklets and detections"""
+        return self.dist(tracks, dets)
+    
+    def matching_scores(self, tracks: List[Tracklet], 
+                 dets: List[Tracklet]) -> np.ndarray:
+        """ Compute the matching scores between tracklets and detections"""
+        return self.dist.matching_scores(tracks, dets)
+        
+    def assign(self, distances: np.ndarray, thresh: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        return linear_assignment(distances, thresh=thresh)
+

models/models/trackers/reid_parallel_tracker/matchers/single_stage_matcher.py

@@ -0,0 +1,39 @@
+import numpy as np
+from typing import List, Tuple, Dict
+from ..core.matching import linear_assignment
+from ..core.tracklet import Tracklet
+from .distances import DistCosine
+
+class SingleStageMatcher():
+    def __init__(self,
+                 dist_high_cfg: Dict,
+                 dist_low_cfg: Dict,
+                 match_thr: float):
+        """ Perform matching with high score detection boxes and low score detection boxes in a single step
+
+        Args:
+            - dist_high_cfg (Dict): distance fucntion to compute the matching score between tracklets and high score detection boxes
+            - dist_low_cfg (Dict): distance function to compute the matching score between tracklets and low score detection boxes. This function should be stricter than dist_high_cfg as low_detection is less reliable than high_detection. 
+            - match_thr (float): matching distance threshold. Lower value means stricter matching.
+        """
+        self.dist_high = DistCosine(**dist_high_cfg)
+        self.dist_low  = DistCosine(**dist_low_cfg)
+        self.match_thr = match_thr
+        
+    def __call__(self, 
+                 tracks: List[Tracklet], 
+                 dets_high: List[Tracklet],
+                 dets_low: List[Tracklet]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """ Associate with tracklets with detection boxes"""
+        dist_high = self.dist_high(tracks, dets_high)
+        dist_low = self.dist_low(tracks, dets_low)
+        if np.prod(dist_high.shape) >0 and np.prod(dist_low.shape) >0:
+            beta = dist_high.max()/dist_low.max()
+            dists = np.concatenate([dist_high,beta*dist_low],axis=1)
+        else:
+            if np.prod(dist_high.shape) >0:
+                dists = dist_high
+            else:
+                dists = dist_low
+        matches_idx, unmatched_tracks_idx, unmatched_dets_idx = linear_assignment(dists, thresh=self.match_thr)
+        return matches_idx, unmatched_tracks_idx, unmatched_dets_idx

models/models/trackers/reid_parallel_tracker/parallel_tracker.py

@@ -0,0 +1,470 @@
+
+from typing import Tuple, Dict,List
+import numpy as np
+from .core.tracklet import TrackState
+from mmcv.ops import bbox_overlaps
+import torch
+from .three_stage_tracker import ThreeStageTracker
+from .matchers.distances import DistCosine
+
+
+class ParallelTracker(ThreeStageTracker):
+    def __init__(self,
+                 reid_filter_det=None,
+                 *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # remove duplicated boxes
+        self.reid_filter_det = reid_filter_det
+        if reid_filter_det is not None:
+            self.reid_filter_det_dist = DistCosine(**reid_filter_det['dist_cfg'])
+        else:
+            self.reid_filter_det_dist = None
+        self.history_changes = []
+
+    def merge_active_lost_tracks(self,
+                                 dets_high:np.array,
+                                 active_tracks:List,
+                                 lost_tracks:List,
+                                 matched_active_track_indices:np.array,
+                                 matched_lost_track_indices:np.array):
+
+        """
+        The merge_active_lost_tracks merges active and lost tracks based on matching detection results, updating their associated 
+        indices. The merging is guided by certain criteria, such as distance comparisons between tracks and 
+        detections, and whether the tracks share active frames.
+
+        Args:
+
+            dets_high (array): An array representing high-confidence detections.
+            active_tracks (list): A list of active track objects.
+            lost_tracks (list): A list of lost track objects.
+            matched_active_track_indices (array): An array of matched indices between active tracks and detections.
+            matched_lost_track_indices (array): An array of matched indices between lost tracks and detections.
+
+        Returns:
+
+            A dictionary containing various arrays representing different aspects of the merging process:
+                unmerged_active_track_indices: Array of unmatched active track indices that are not merged with lost tracks and cannot match with any detection.
+                unmerged_lost_track_indices: Array of unmatched lost track indices that are not merged with active tracks and cannot match with any detection.
+                matched_pairs_active: Array of matched indices between active tracks and detections after the merging process.
+                matched_pairs_lost: Array of matched indices between lost tracks and detections after the merging process.
+                active_tracks_match_lost_tracks_detection: Array of active track indices that were matched to the same detection as lost tracks.
+                merged_lost_indices: Array of lost track indices that were merged with active tracks.
+        """
+
+        unmerged_active_track_indices = []
+        unmerged_lost_track_indices = []
+        matched_pairs_active = []  # Match_pairs 1
+        matched_pairs_lost = []  # Match_pairs 1
+        active_tracks_match_lost_tracks_detection = []
+        merged_lost_indices = []
+
+        # Active/lost tracks match to different detections --> match_pairs_1
+        matched_active_det_indices = matched_active_track_indices[:, 1].tolist(
+        )
+        matched_lost_det_indices = matched_lost_track_indices[:, 1].tolist()
+
+        for idx, det_idx in enumerate(matched_active_det_indices):
+            if not det_idx in matched_lost_det_indices:
+                matched_pairs_active.append(matched_active_track_indices[idx])
+
+        for idx, det_idx in enumerate(matched_lost_det_indices):
+            if not det_idx in matched_active_det_indices:
+                matched_pairs_lost.append(matched_lost_track_indices[idx])
+
+        # detection --matches--> row-th of active track in matched_active_track_indices[:, 0]
+        det_active_rowmatrix_dict = dict()
+        for idx, (_active_trackidx, _detidx) in enumerate(matched_active_track_indices):
+            det_active_rowmatrix_dict[_detidx] = idx
+
+        # detection --matches--> row-th of lost track in matched_lost_track_indices[:, 0]
+        det_lost_rowmatrix_dict = dict()
+        for idx, (_lost_trackidx, _detidx) in enumerate(matched_lost_track_indices):
+            det_lost_rowmatrix_dict[_detidx] = idx
+
+        # calculate distance between active_tracks and dets_high
+        active_dists = self.matcher_active.matching_dists(active_tracks,
+                                                          dets_high)
+        # calculate distance between lost_tracks and dets_high
+        lost_dists = self.matcher_lost.matching_dists(lost_tracks,
+                                                      dets_high)
+
+        # -------------------------- Main Loop ----------------------------------
+        for detidx, row_lost_idx in det_lost_rowmatrix_dict.items():
+            is_matched_to_the_same_detection = detidx in det_active_rowmatrix_dict
+            if is_matched_to_the_same_detection:
+                row_active_idx = det_active_rowmatrix_dict[detidx]
+                # one detection --> 2 tracks
+                active_track_idx = matched_active_track_indices[row_active_idx][0]
+                lost_track_idx = matched_lost_track_indices[row_lost_idx][0]
+
+                # Both active in the past at the same time?
+                both_active_in_the_past_at_the_same_time = active_tracks[active_track_idx].common_active_frames(
+                    lost_tracks[lost_track_idx])
+
+                if both_active_in_the_past_at_the_same_time:
+                    active_dist = active_dists[active_track_idx, detidx]
+                    lost_dist = lost_dists[lost_track_idx, detidx]
+                    # active_track-->det has smaller distance than lost_track-->det
+                    is_active_track_give_smaller_dist = active_dist <= lost_dist
+                    if is_active_track_give_smaller_dist:
+                        # assign detection to active track
+                        matched_pairs_active.append(
+                            matched_active_track_indices[row_active_idx])
+                        unmerged_lost_track_indices.append(
+                            lost_track_idx)
+                    else:
+                        # assign detection to lost track
+                        matched_pairs_lost.append(
+                            matched_lost_track_indices[row_lost_idx])
+                        unmerged_active_track_indices.append(
+                            active_track_idx)
+                else:
+                    # ------------------- Merge trajectories --------------------------------------
+                    # the active track will be removed, match this detection for lost tracklet
+                    active_tracks_match_lost_tracks_detection.append(
+                        active_track_idx)
+                    matched_pairs_lost.append(
+                        matched_lost_track_indices[row_lost_idx])
+                    merged_lost_indices.append(lost_track_idx)
+                    # update active frames for lost track
+                    lost_tracks[lost_track_idx].update_active_frames(
+                        active_tracks[active_track_idx].active_frames)
+
+        # convert to numpy array
+        unmerged_active_track_indices = np.array(
+            unmerged_active_track_indices)
+        unmerged_lost_track_indices = np.array(
+            unmerged_lost_track_indices)
+        matched_pairs_active = np.array(matched_pairs_active).reshape(-1, 2)
+        matched_pairs_lost = np.array(matched_pairs_lost).reshape(-1, 2)
+        active_tracks_match_lost_tracks_detection = np.array(
+            active_tracks_match_lost_tracks_detection)
+        merged_lost_indices = np.array(merged_lost_indices)
+
+        # return values
+        return dict(
+            unmerged_active_track_indices=unmerged_active_track_indices,
+            unmerged_lost_track_indices=unmerged_lost_track_indices,
+            matched_pairs_active=matched_pairs_active,  # Match_pairs 1
+            matched_pairs_lost=matched_pairs_lost,  # Match_pairs 1
+            active_tracks_match_lost_tracks_detection=active_tracks_match_lost_tracks_detection,
+            merged_lost_indices=merged_lost_indices,
+        )
+
+    def update(self,
+               det_result: Dict,
+               Hmat: np.array = None,
+               meta_data: Dict = None) -> Tuple[Dict, Dict]:
+        """ 
+        The update function is the main method for performing tracking. It takes a detection result, a homography transformation matrix (optional), 
+        and additional meta-data (optional) as input and returns the active tracks and lost tracks in the current frame after updating 
+        their states and associations with new detections.
+
+        Args:
+
+            det_result (dict): A dictionary representing the detection result from a detector.
+            Hmat (np.array, optional): A NumPy array representing the homography transformation matrix. 
+            meta_data (dict, optional): Additional meta-data that may be used during the tracking process.
+            
+        Returns:
+
+            active_tracks (dict): A dictionary representing the tracked active tracks in the current frame. 
+            lost_tracks (dict): A dictionary representing the lost tracks in the current frame. 
+            modifications (dict): A dictionary containing modifications made during the tracking process.
+        """
+        self.frame_id += 1
+
+        # Step 1: Split the detections into high score/lower score group
+        det_result = self.preprocess_det_result(det_result)
+
+        # remove duplicated detection by ReID
+        if self.reid_filter_det_dist is not None:
+            det_result = self.remove_det_by_reid(det_result)
+
+        dets_high, dets_low = self.split_detections_by_scores(det_result)
+
+        # Step 2: Split the tracks into active_tracks, lost_tracks, and unconfirmed (just initialize)
+        active_tracks, lost_tracks, unconfirmed = self.split_tracks_by_activation()
+
+        # - predict the current location with KF, and compensate for Camera Motion
+        self.predict_with_gmc(active_tracks, lost_tracks, unconfirmed, Hmat)
+
+        # Step 3: Matching Stage 1.1 - Association with high score detection boxes and active tracks
+        matched_active_track_indices, unmatched_active_track_indices, unmatched_active_det_indices \
+            = self.matcher_active(active_tracks, dets_high)
+
+        # Step 4: Matching Stage 1.2 - Association with high score detection boxes and lost tracks
+        matched_lost_track_indices, unmatch_lost_track_indices, unmatch_lost_det_indices = self.matcher_lost(
+            lost_tracks, dets_high)
+
+        # Step 5. assing the Lost status to the unmatch_tracks
+        lost_stracks = [active_tracks[it]
+                        for it in unmatched_active_track_indices if active_tracks[it].state != TrackState.Lost]
+        for track in lost_stracks:
+            track.mark_lost()
+
+        # Step 6. Merge active tracks and lost tracks
+        merged_results = self.merge_active_lost_tracks(dets_high,
+                                                       active_tracks,
+                                                       lost_tracks,
+                                                       matched_active_track_indices,
+                                                       matched_lost_track_indices)
+
+        unmerged_active_track_indices = merged_results[
+            'unmerged_active_track_indices']
+        unmerged_lost_track_indices = merged_results['unmerged_lost_track_indices']
+        match_pairs_1_1 = merged_results['matched_pairs_active']
+        match_pairs_1_2 = merged_results['matched_pairs_lost']
+        active_tracks_match_lost_tracks_detection = merged_results[
+            'active_tracks_match_lost_tracks_detection']
+        merged_lost_indices = merged_results['merged_lost_indices']
+
+        # Step 7. Assign the Lost status to the Merged unmatched active tracks
+        lost_stracks_1 = [active_tracks[it]
+                          for it in unmerged_active_track_indices if active_tracks[it].state != TrackState.Lost]
+        for track in lost_stracks_1:
+            track.mark_lost()
+        lost_stracks.extend(lost_stracks_1)
+
+        # Step 8. update Match_pairs 1
+        activated_stracks, refind_stracks = self.update_matched_tracks(
+            match_pairs_1_1, active_tracks, dets_high)
+
+        activated_stracks_1, refind_stracks_1 = self.update_matched_tracks(
+            match_pairs_1_2, lost_tracks, dets_high)
+
+        activated_stracks.extend(activated_stracks_1)
+        refind_stracks.extend(refind_stracks_1)
+
+        # Step 9. Remove active tracks that are merged into lost tracks
+        merged_active_tracks = [active_tracks[idx]
+                                for idx in active_tracks_match_lost_tracks_detection]
+        for track in merged_active_tracks:
+            track.mark_removed()
+
+        # for post processing: later step will know how to change switch ID back to the lost one
+        need_replaced_track_info = [
+            _m_track.history_info for _m_track in merged_active_tracks]
+        new_track_info = [
+            lost_tracks[_idx].history_info for _idx in merged_lost_indices]
+
+        # Step 10. Take Unmatched dets in both list
+        merge_unmatched_det_indices = list(set(unmatched_active_det_indices.tolist()).intersection(
+            set(unmatch_lost_det_indices.tolist())))
+
+        #          remain unmatch detection
+        unmatched_dets_2 = [
+            dets_high[_idx] for _idx in merge_unmatched_det_indices]
+
+        # Step 11: Matching Stage 2 - Association between new detections and unconfirmed tracks (tracks just initialized in the previous frame)
+        matched_pairs_2, unmatch_track_indices_2, unmatched_high_det_indices_2 = self.matcher_unconfirmed(
+            unconfirmed, unmatched_dets_2)
+
+        for itracked, idet in matched_pairs_2:
+            # Update matches_unconfirmed into active_tracks
+            unconfirmed[itracked].update(unmatched_dets_2[idet], self.frame_id)
+            activated_stracks.append(unconfirmed[itracked])
+
+        # Step 12. remove unconfirmed tracks that do not match any detections
+        removed_stracks = [unconfirmed[it] for it in unmatch_track_indices_2]
+        for track in removed_stracks:
+            track.mark_removed()
+
+        all_remove_tracks = removed_stracks + merged_active_tracks
+        # Step 13. init new stracks if they are high score and not too small boxes as tentative tracks (unconfirmed)
+        new_tracks = self.init_new_tracks(
+            unmatched_dets_2, unmatched_high_det_indices_2)
+
+        # new_tracks that have very high score and not occluded with current tracks will be directly activated
+        self.activate_new_tracks(
+            new_tracks, activated_stracks + refind_stracks)
+
+        # Step 14: remove lost tracks if they are already lost for a certain frames
+        self.lost_stracks.extend(lost_stracks)
+        removed_lost_stracks = self.remove_lost_tracks()
+        all_remove_tracks += removed_lost_stracks
+
+        # Step 15: Final result merging
+        active_tracks, lost_tracks = self.merge_results(
+            activated_stracks, refind_stracks, new_tracks, all_remove_tracks)
+
+        # update lost frame num
+        self.update_lost_frame()
+
+        modifications = dict(
+            need_replaced_track_info=need_replaced_track_info,
+            new_track_info=new_track_info
+        )
+
+        return active_tracks, lost_tracks, modifications
+
+    def remove_det_by_reid(self, det_result:Dict):
+        """
+        The remove_det_by_reid function is a method of a class that filters out detection results based on the 
+        ReID (Person Re-identification) distances between detected objects. It removes similar detections, keeping 
+        only those that are dissimilar according to the ReID distance threshold and certain object characteristics 
+        like bounding box overlap and detection scores.
+
+        Args:
+
+            det_result (dictionary): A dictionary containing the detection results.
+        Returns:
+
+            filter_det_result (dictionary): A dictionary containing the filtered detection results after applying the ReID-based filtering.
+        """
+        all_embeddings = det_result['embeddings']
+        reid_dists = self.reid_filter_det_dist(all_embeddings, all_embeddings)
+        np.fill_diagonal(reid_dists, 1e4)
+
+        # det boxes and score
+        boxes = det_result['boxes']
+        boxes = boxes.reshape(-1, 5)
+        boxes_tensor = torch.from_numpy(boxes)[:, :-1]
+
+        # score matrix
+        det_scores = boxes[:, -1]
+        det_scores_matrix = det_scores.reshape(-1,
+                                               1) > det_scores.reshape(1, -1)
+
+        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+        width = x2 - x1
+        height = y2 - y1
+        areas = np.abs(width * height)
+        det_area_matrix = areas.reshape(-1, 1) > areas.reshape(1, -1)
+
+        box_iof = bbox_overlaps(boxes_tensor, boxes_tensor, mode='iof').numpy()
+        np.fill_diagonal(box_iof, 0.0)
+        iof_non_overllaped = box_iof < self.reid_filter_det['iof_thr']
+
+        reid_dists[iof_non_overllaped] = 1e4
+        # reid_dists[det_scores_matrix] = 1e4
+        reid_dists[det_area_matrix] = 1e4
+
+        remove_det_matrix = reid_dists < self.reid_filter_det['reid_dist_thr']
+        remove_det_matrix = remove_det_matrix.sum(1) > 0
+        remove_det_idx = np.nonzero(remove_det_matrix)[0]
+        # draw_queries_galleries(det_result['obj_imgs'], det_result['obj_imgs'], reid_dists, det_result['frame_id'])
+
+        if len(remove_det_idx):
+            filter_det_result = dict()
+            filter_det_result['frame_id'] = det_result['frame_id']
+            for key, val in det_result.items():
+                filter_val = []
+                if key in ['frame_id']:
+                    continue
+                for idx in range(len(val)):
+                    if not idx in remove_det_idx:
+                        filter_val.append(val[idx])
+                filter_det_result[key] = np.array(filter_val)
+        else:
+            filter_det_result = det_result
+
+        return filter_det_result
+
+    def remove_rows(self, matrix:np.array, row_indices:List, C=2):
+        """
+        The remove_rows function takes a 2D matrix as input and removes specific rows from 
+        the matrix based on the provided row indices. The resulting matrix is then returned with a reshaped number of columns.
+
+        Args:
+
+            matrix (numpy array): A 2D numpy array representing the input matrix.
+            row_indices (list): A list of integers representing the row indices to be removed from the matrix.
+            C (int, optional): An integer representing the number of columns in the output matrix.
+        Returns:
+
+            result_matrix (numpy array): The resulting 2D numpy array after removing the specified rows and reshaping with the 
+            specified number of columns (C).
+        """
+        row_indices = sorted(row_indices, reverse=True)
+        matrix = matrix.tolist()
+        for row_idx in row_indices:
+            matrix.pop(row_idx)
+        return np.array(matrix).reshape(-1, C)
+
+    def update_unmatch_track_det(self, tracklets:List, dets:List, matches_indices:np.array, remove_indices:List=[]):
+        """
+        The update_unmatch_track_det updates and retrieves the indices of unmatched tracklets 
+        and detections based on the provided matched indices and any specified indices for removal.
+
+        Args:
+
+            tracklets (list): A list of tracklet objects representing existing tracks.
+            dets (list): A list of detection objects representing detected items.
+            matches_indices (numpy array): A 2D numpy array containing matched indices between tracklets and detections.
+            remove_indices (list, optional): A list of integers representing indices of tracklets to be removed. 
+
+        Returns:
+
+            unmatched_track_indices (numpy array): A 1D numpy array containing the indices of unmatched tracklets (tracks that do not 
+            have any matches with detections).
+            unmatched_det_indices (numpy array): A 1D numpy array containing the indices of unmatched detections (detections that do 
+            not have any matches with tracklets).
+        """
+        # update unmatch track, det
+        unmatched_track_indices, unmatched_det_indices = [], []
+        for _track_idx in range(len(tracklets)):
+            if _track_idx in remove_indices:
+                continue
+            if not _track_idx in matches_indices[:, 0]:
+                unmatched_track_indices.append(_track_idx)
+        for _det_idx in range(len(dets)):
+            if not _det_idx in matches_indices[:, 1]:
+                unmatched_det_indices.append(_det_idx)
+        unmatched_track_indices = np.array(unmatched_track_indices)
+        unmatched_det_indices = np.array(unmatched_det_indices)
+
+        return unmatched_track_indices, unmatched_det_indices
+
+    def det_local_aware_spliter(self, dets:List, mar_top:float=0.2, mar_left:float=0.25,
+                                mar_right:float=0.25, h:int=1876, w:int=2896):
+        """
+        The det_local_aware_spliter function is a method of a class. It splits a list of detections into two separate lists based on their positions within a specified region of interest in an image.
+
+        Args:
+
+            dets (list): A list of detection objects representing detected items.
+            mar_top (float, optional): A floating-point value representing the margin ratio from the top of the image.
+            mar_left (float, optional): A floating-point value representing the margin ratio from the left side of the image.
+            mar_right (float, optional): A floating-point value representing the margin ratio from the right side of the image.
+            h (int, optional): An integer representing the height of the image. 
+            w (int, optional): An integer representing the width of the image. 
+        Returns:
+
+            det_in (list): A list of detection objects that lie inside the specified region of interest.
+            det_out (list): A list of detection objects that lie outside the specified region of interest.
+        """
+        det_in, det_out = [], []
+        for det in dets:
+            t, l, b, r = det.tlbr
+            x, y = l+(r-l)/2, t+(b-t)/2
+            if (x < w*(1-mar_right)) and (x > w*mar_left) and (y > h*mar_top):
+                det_in.append(det)
+            else:
+                det_out.append(det)
+        return det_in, det_out
+
+    def split_size(self, dets:list, thr:float=0.1, h:int=1876, w:int=2896):
+        """
+        The split_size function is a method of a class. It separates a list of detection objects into two separate lists based on their height relative to a specified threshold, as a proportion of the image height.
+
+        Args:
+
+            dets (list): A list of detection objects representing detected items.
+            thr (float, optional): A floating-point value representing the threshold for splitting detections based on height.
+            h (int, optional): An integer representing the height of the image. 
+            w (int, optional): An integer representing the width of the image. 
+        Returns:
+
+            det_small (list): A list of detection objects whose height is smaller than the specified threshold.
+            det_large (list): A list of detection objects whose height is larger than or equal to the specified threshold.
+        """
+        det_small, det_large = [], []
+        for det in dets:
+            t, _, b, _ = det.tlbr
+            if b-t < thr*h:
+                det_small.append(det)
+            else:
+                det_large.append(det)
+        return det_small, det_large

models/models/trackers/reid_parallel_tracker/three_stage_tracker.py

@@ -0,0 +1,132 @@
+
+from typing import Tuple, Dict
+import numpy as np
+from .core.tracklet import TrackState
+from .core.tracklet import (Tracklet, TrackState)
+from .base_tracker import BaseTracker
+from .matchers.base_matchers import SimMatcher
+
+class ThreeStageTracker(BaseTracker):
+    def __init__(self,  
+                matcher_active_cfg = dict(match_thr=0.5),
+                matcher_det_low_cfg = dict(match_thr=0.5),
+                matcher_lost_cfg = dict(match_thr=0.5),
+                matcher_unconfirmed_cfg = dict(match_thr=0.5),
+                enable_reid_buffer = False,       
+                *args,**kwargs):
+        super().__init__(*args,**kwargs)
+        self.matcher_active = SimMatcher(**matcher_active_cfg)
+        self.matcher_det_low = SimMatcher(**matcher_det_low_cfg)
+        self.matcher_lost = SimMatcher(**matcher_lost_cfg)
+        self.matcher_unconfirmed = SimMatcher(**matcher_unconfirmed_cfg)
+        self.enable_reid_buffer = enable_reid_buffer
+
+    def split_tracks_by_activation(self):
+        """ Split the tracks into active_tracks, lost_tracks, and unconfirmed (just initialize)
+        Returns:
+            strack_pool: List[Tracklet]
+            unconfirmed: List[Tracklet]
+        """
+        unconfirmed = []
+        tracked_stracks = []  # type: list[Tracklet]
+        for track in self.tracked_stracks:
+            if not track.is_activated:
+                unconfirmed.append(track)
+            else:
+                tracked_stracks.append(track)
+        return tracked_stracks, self.lost_stracks, unconfirmed
+
+    def predict_with_gmc(self, active_tracks, lost_tracks, unconfirmed, Hmat):
+        Tracklet.multi_predict(active_tracks)
+        Tracklet.multi_predict(lost_tracks)
+        if Hmat is not None:
+            Tracklet.multi_gmc(active_tracks, Hmat)
+            Tracklet.multi_gmc(lost_tracks, Hmat)
+            Tracklet.multi_gmc(unconfirmed, Hmat)
+
+    def update(self, 
+               det_result: Dict, 
+               Hmat: np.array=None,
+               meta_data: Dict=None) -> Tuple[Dict, Dict]:
+        """ The main function to perform tracking. The pipeline is similar to ByteTrack/BoTSort, 
+            which first associate with high score detection boxes, and then associate with low score detection boxes.
+        Args:
+            det_result (dict): detection result from detector
+            Hmat (np.array, optional): Homography transformation matrix. Defaults to None.
+        Returns:
+            active_tracks (dict): tracked tracks in the current frame. See format_track_results for the format
+            lost_tracks (dict): lost tracks in the current frame. See format_track_results for the format
+        """
+        self.frame_id += 1
+        # Step 1: Split the detections into high score/lower score group 
+        det_result = self.preprocess_det_result(det_result)
+        dets_high, dets_low = self.split_detections_by_scores(det_result)
+
+        # Step 2: Split the tracks into trackpool=(active_tracks + lost_tracks) and unconfirmed (just initialize)
+        active_tracks, lost_tracks, unconfirmed = self.split_tracks_by_activation()
+        # - predict the current location with KF, and compensate for Camera Motion
+        self.predict_with_gmc(active_tracks, lost_tracks, unconfirmed, Hmat)
+        
+        # Step 3: First association with high score detection boxes and active track
+        match_idxes, unmatch_track_idxes, unmatch_det_idxes= self.matcher_active(active_tracks, dets_high)
+        activated_stracks, refind_stracks = self.update_matched_tracks(match_idxes, active_tracks, dets_high) # refind_track=[]
+
+        # Step 4: Second association with low score detection boxes"""
+        unmatch_tracks = [active_tracks[idx] for idx in unmatch_track_idxes if active_tracks[idx].state == TrackState.Tracked]
+        
+        match_idxes_2, unmatch_track_idxes_2, _ = self.matcher_det_low(unmatch_tracks, dets_low)
+        activated_stracks_2, refind_stracks_2 = self.update_matched_tracks(match_idxes_2, unmatch_tracks, dets_low)
+        activated_stracks.extend(activated_stracks_2 )
+        refind_stracks.extend(refind_stracks_2)
+
+        # - assing the Lost status to the unmatch_tracks that are not matched with low score detection boxes
+        lost_stracks = [unmatch_tracks[it] for it in unmatch_track_idxes_2 if unmatch_tracks[it].state != TrackState.Lost]
+        for track in lost_stracks: track.mark_lost()
+
+        # Step 5: Second association, between remain high detections and lost tracks
+        remain_dets = [dets_high[i] for i in unmatch_det_idxes]
+        match_idxes_3, unmatch_track_idxes_3, unmatch_det_idxes_3 = self.matcher_lost(lost_tracks, remain_dets)
+        activated_stracks_3, refind_stracks_3 = self.update_matched_tracks(match_idxes_3, lost_tracks, remain_dets)
+        activated_stracks.extend(activated_stracks_3,)
+        refind_stracks.extend(refind_stracks_3)
+        unmatch_lost_tracks = [lost_tracks[idx] for idx in unmatch_track_idxes_3]
+        lost_stracks.extend(unmatch_lost_tracks)
+
+        # Step 6: Second association, between new detections and lost tracks (tracks just initialized in the previous frame)
+        remain_dets_3 = [remain_dets[i] for i in unmatch_det_idxes_3]
+        match_idxes_4, unmatch_track_idxes_4, unmatch_det_idxes_4 = self.matcher_unconfirmed(unconfirmed, remain_dets_3)
+
+        for itracked, idet in match_idxes_4:
+            # Update matches_unconfirmed into active_tracks
+            unconfirmed[itracked].update(remain_dets_3[idet], self.frame_id)
+            activated_stracks.append(unconfirmed[itracked])
+
+        # - remove unconfirmed tracks that do not match any detections
+        removed_stracks = [unconfirmed[it] for it in unmatch_track_idxes_4]
+        for track in removed_stracks: track.mark_removed()
+
+        # - init new stracks if they are high score and not too small boxes as tentative tracks (unconfirmed)
+        new_tracks = self.init_new_tracks(remain_dets_3, unmatch_det_idxes_4)
+        # - new_tracks that have very high score and not occluded with current tracks will be directly activated
+        self.activate_new_tracks(new_tracks, activated_stracks + refind_stracks)
+
+        # Step 6: remove lost tracks if they are already lost for a certain frames
+        self.lost_stracks.extend(lost_stracks)
+        removed_lost_stracks = self.remove_lost_tracks()
+        removed_stracks += removed_lost_stracks
+        
+        # Step 7: Final result merging
+        active_tracks, lost_tracks = self.merge_results(activated_stracks, refind_stracks, new_tracks, removed_stracks)
+        
+        # update lost frame num
+        self.update_lost_frame()
+
+        return active_tracks,lost_tracks
+
+    def update_lost_frame(self):
+
+        for track in self.tracked_stracks:
+            track.reset_lost_frame_num()
+        
+        for track in self.lost_stracks:
+            track.set_lost_frame_num()

models/reids/solider.py

@@ -0,0 +1,87 @@
+from typing import Dict, Tuple, List
+from models.base.trt_base import TRT_Base
+import torch
+import cv2
+import numpy as np
+
+class SOLIDERBase(): 
+    def __init__(self,
+                    preprocess_cfg: Dict=dict(
+                        mean=[0.5, 0.5, 0.5],
+                        std=[0.5, 0.5, 0.5]
+                    ),
+                    use_torch: bool=False):
+        """ SOLIDERBase class for inference.
+
+        Args:
+            preprocess_cfg (Dict):
+                - mean (List[float, float, float]): mean offset values for preprocessing.
+                - std (List[float, float, float]): standard deviation offset values for preprocessing.
+            use_torch (bool): use torch tensor or numpy array in preprocess and postprocess function.
+        """
+        self.preprocess_cfg = preprocess_cfg
+        self.use_torch = use_torch
+    
+    def preprocess(self, input_data: np.ndarray):
+        """ Preprocess function for input data.
+
+        Args:
+            input_data (np.ndarray): batch input image.
+        """
+        tensor_data = []
+        if ((isinstance(input_data, np.ndarray)) and (len(input_data.shape) == 3)):
+            input_data = [input_data]
+        for i in range(len(input_data)):
+            img = input_data[i]
+            img = cv2.resize(img, self.input_shape[2:][::-1], interpolation=cv2.INTER_LINEAR)
+            height, width = img.shape[0], img.shape[1]
+            if self.use_torch:
+                tensor_data.append(torch.from_numpy(img).to(self.device))
+            else:
+                tensor_data.append(img)
+        if self.use_torch:
+            mean = torch.tensor(mean).to(self.device)
+            std = torch.tensor(std).to(self.device)
+            tensor_data = (torch.stack(tensor_data, dim=0)[:, :, :, [2, 1, 0]]/255.0 - mean)/std
+            tensor_data = tensor_data.permute(0, 3, 1, 2).float().contiguous()/255.0
+        else:
+            mean = np.array(mean)
+            std = np.array(std)
+            tensor_data = (np.stack(tensor_data, axis=0)[:, :, :, [2, 1, 0]]/255.0 - mean)/std
+        return tensor_data, height, width
+    
+class SOLIDERTRT(TRT_Base, SOLIDERBase):
+    def __init__(self, 
+                    preprocess_cfg: Dict,
+                    img_shape: Tuple[int, int]=(384, 128),
+                    batch_size: int=32,
+                    model_path: str="",
+                    device: str='0',):
+        """ SOLIDER TRT class for inference, which is based on TRT_Base and SOLIDERBase.
+        """
+        self.img_shape = img_shape
+        self.batch_size = batch_size
+        input_shape = (self.batch_size, *self.img_shape)
+        super().__init__(input_shape, model_path, device)
+        SOLIDERBase.__init__(self, preprocess_cfg=preprocess_cfg, use_torch=True)
+    
+    def infer_batch(self, image_batch: np.ndarray) -> List[np.ndarray]:
+        """ Batch inference function for batch input image.
+
+        Args:
+            image_batch (np.ndarray): batch of input image.
+        """
+
+        tensor_data, height, width = self.preprocess(image_batch)
+        self.change_runtime_dimension(input_shape=(len(tensor_data), 3, height, width))
+        self.model['binding_addrs']['images'] = int(tensor_data.data_ptr())
+        self.model['context'].execute_v2(list(self.model['binding_addrs'].values()))
+        feats = self.model['bindings']['feats'].data.cpu()
+    
+        reid_outputs = []
+        for idx in range(len(feats)):
+            feat = feats[idx]
+            reid_outputs.append({"feat": feat.float().numpy()})
+        return reid_outputs 
+
+

models/trackers/byte_track.py

@@ -0,0 +1,56 @@
+from mmtrack.models.trackers.byte_tracker import ByteTracker as MMByteTracker
+from mmtrack.models.builder import build_motion
+from typing import List, Dict
+import numpy as np
+import mmcv 
+
+
+class BYTETracker(MMByteTracker):
+    def __init__(self, *args, **kwargs):
+        """ ByteTracker class for tracking.
+
+        Args:
+            obj_score_thrs (Dict):
+                - high (float): if detection box > high -> high_score_detections for first association.
+                - low (float): if low < detection box < high -> low_score_detections for second association.
+            init_track_thr (float): Detection score threshold for initializing a new tracklet.
+            weight_iou_with_det_scores (bool): Whether using detection scores to weight IOU which is used for matching.
+            match_iou_thrs (Dict): IOU distance threshold for matching between two frames.
+                - high (float): Threshold of the first matching.
+                - low (float): Threshold of the second matching.
+                - tentative (float): Threshold of the matching for tentative tracklets.
+            num_frames_retain (int): If a track is disappeared more than num_frames_retain frames, it will be deleted in the memo.
+            motion (Dict): Config for motion.
+                - type (str): Motion type (KalmanFilter, LinearFilter).
+        """
+        self.motion = build_motion(kwargs.pop("motion"))
+        super().__init__(*args, **kwargs)
+        
+    def track_batch(self, start_frame_idx: int, det_results: List[Dict], conf_thres: float=0.0):
+        """ Batch inference function for batch det results.
+
+        Args:
+            start_frame_idx (int): start_frame_idx of this batch.
+            det_results (List[Dict]): detection results of this batch.
+            conf_thres (float): If a tracklet's confidence < confidence threshold, it will be removed.
+        """
+
+        track_outputs = []
+        for frame_id, frame_det_outputs in enumerate(det_results):
+            boxes, labels = frame_det_outputs.pop("boxes"), frame_det_outputs.pop("labels")
+
+            boxes, labels, track_ids = self.track(None, None, self, 
+                                                bboxes=boxes, labels=labels, 
+                                                frame_id=frame_id+start_frame_idx, 
+                                                rescale=False)
+            boxes = np.around(boxes.numpy(),decimals=3)
+            labels = labels.numpy().astype(np.uint)
+            track_ids = track_ids.numpy().astype(np.uint)
+            idxs = np.where(boxes[:, 4] > conf_thres)[0]
+            track_outputs.append({
+                    "boxes": boxes[idxs],
+                    "labels": labels[idxs],
+                    "ids": track_ids[idxs]
+                })
+        return track_outputs
+        

projects/human_detection/ReadMe.md

@@ -0,0 +1,43 @@
+# HUMAN DETECTION DEMO
+
+### 1. Prepare data and weights:
+The weights and sample test video files can be downloaded from:
++ NAS at `/5_project_internal/19_Demo_App/human_detection`.
++ On HPC2,`/data/cc-demo/human_detection/weights/`
+
+If you run the demo on other machine, please download the files and put it the same path as HPC2.
+From now on, we will assume that you have data and weights in the following structure:
+```
+/data/cc-demo/human_detection/
+├── sample_inputs
+│   ├── 1.mp4
+│   ├── 2.mp4
+│   └── 3.mp4
+└── weights
+        └── mmyolov8_s_human_dynamic_shape.onnx
+            mmyolov8_s_human_DBsx640x800.trt
+
+
+```
+### 2. Compile ONNX & TRT model 
+#### 2.1. Start docker container
+For Developer that want to add more functions to the codebase, and use Docker as Developing environment, use:
+```
+bash docker_run.sh 
+docker attach <docker_container_name>
+```
+After attaching into docker, do the following step to export onnx and trt model
+```
+bash projects/human_detection/export_onnx_trt/export_trt_mmyolov8.sh
+```
+It will take about 20minutes to comple ONNX & TRT model, and the file will be stored at folder `/data/human_detection\deploy`
+
+### 3. Start Gradio 
+a. From inside Docker
+```
+bash projects/human_detection/demo_app.sh
+```
+b. From Host machine
+```
+bash projects/human_detection/docker_run.sh
+```

projects/human_detection/demo_app.py

@@ -0,0 +1,147 @@
+import os
+
+from models.detectors.yolov7 import YOLOv7ONNX,YOLOv7TRT
+from models.detectors.mmyolov8 import MMYOLOv8ONNX,MMYOLOv8TRT
+from projects.human_detection.engine.pipeline import run_e2e_pipeline
+from mmcv import VideoReader
+import gradio as gr
+
+title = "Human Detection - CYBERCORE AI DEMO"
+description = "Human Monitoring Demo: It will run detection, tracking on input video and show output video.\nYou may click on of the examples or upload your own image."
+
+root_folder = "/data/human_detection/"
+example_video_paths = [
+    os.path.join(root_folder,"sample_inputs/1.mp4"),
+    os.path.join(root_folder,"sample_inputs/2.mp4"),
+    os.path.join(root_folder,"sample_inputs/3.mp4"),
+]
+output_dir = os.path.join(root_folder,"outputs")
+os.makedirs(output_dir, exist_ok=True)
+assert any([os.path.exists(example_video_path) for example_video_path in example_video_paths]), f"Example video does not exist. Please download example videos from NAS:[https://gofile.me/6ZWyr/q0saLr8hb] and put them in the following path 'human_detection/inputs'"
+
+
+#---------------------Model Path----------------------------------
+# model_path_yolov7_onnx = os.path.join(root_folder, "weights/yolov7_pedestrian_480x640.onnx") # model_path yolov7
+# model_path_yolov7_trt  = os.path.join(root_folder, "weights/yolov7_pedestrian_480x640.trt") # model_path yolov7
+# assert os.path.exists(model_path_yolov7_onnx) or os.path.exists(model_path_yolov7_trt), f"Model does not exist. Please download model from NAS:[path], compile TRT, and put it in the following path 'project_folder'/weights/yolov7-honda-pedestrian_deploy_v2_0.1_768x1280.onnx"
+model_path_yolov8_onnx = os.path.join(root_folder, "weights/mmyolov8_s_human_dynamic_shape.onnx") # model_path yolov8
+model_path_yolov8_trt  = os.path.join(root_folder, "weights/mmyolov8_s_human_DBsx640x800.trt") # model_path yolov8
+assert os.path.exists(model_path_yolov8_onnx) or os.path.exists(model_path_yolov8_trt), f"Model does not exist. Please download model from NAS:[path], compile TRT, and put it in the following path 'project_folder'/weights/yolov8-human.onnx"
+
+# ---------------------Configs---------------------------------- 
+use_trt_yolov8 = os.path.exists(model_path_yolov8_trt)
+yolov8_cfg = dict(
+    img_shape=(3, 640, 800),
+    batch_size=32,
+    preprocess_cfg=dict(
+        border_color=(114, 114, 114),
+        auto=False,
+        scaleFill=False,
+        scaleup=False,
+        stride=32),
+    nms_agnostic_cfg=dict(
+        type='nms',
+        iou_threshold=0.6,
+        class_agnostic=True),
+    score_thr=0.1,
+    model_path = model_path_yolov8_trt if use_trt_yolov8 else model_path_yolov8_onnx,
+    device='0'
+)
+
+# use_trt_yolov7 = os.path.exists(model_path_yolov7_trt)
+# yolov7_cfg = dict(
+#     img_shape=(3, 480, 640),
+#     batch_size=32,
+#     preprocess_cfg=dict(
+#         border_color=(114, 114, 114),
+#         auto=False,
+#         scaleFill=False,
+#         scaleup=True,
+#         stride=32),
+#     nms_agnostic_cfg=dict(
+#         type='nms',
+#         iou_threshold=0.99,
+#         class_agnostic=True),
+#     model_path = model_path_yolov7_trt if use_trt_yolov7 else model_path_yolov7_onnx,
+#     device='0'
+# )
+
+tracker_cfg = dict(
+    obj_score_thrs=dict(high=0.6, low=0.1),
+    init_track_thr=0.7,
+    weight_iou_with_det_scores=True,
+    match_iou_thrs=dict(high=0.1, low=0.5, tentative=0.3),
+    num_frames_retain=30,
+    motion=dict(type='KalmanFilter')
+)
+
+visualizer_cfg = dict(fps=-1, min_width=1280)
+
+
+# yolov7_detector = YOLOv7TRT(**yolov7_cfg) if use_trt_yolov7 else YOLOv7ONNX(**yolov7_cfg)
+yolov8_detector = MMYOLOv8TRT(**yolov8_cfg) if use_trt_yolov8 else MMYOLOv8ONNX(**yolov8_cfg)
+    
+def inference(video, model_name, conf_thres, show_conf, progress=gr.Progress()):
+    output_path = os.path.join(output_dir, os.path.basename(video))
+    if os.path.exists(output_path):
+        os.remove(output_path)
+    # detector = yolov7_detector if model_name == "pedestrian" else yolov8_detector
+    detector = yolov8_detector
+    run_e2e_pipeline(video, detector, tracker_cfg, visualizer_cfg, output_path, conf_thres, show_conf, progress)
+    return output_path
+
+def clear(*all_components):
+    outputs = [None]*len(all_components)
+    return outputs
+
+# ---------------------Gradio UI----------------------------------
+with gr.Blocks(title=title) as demo:
+    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>" + title + "</h1>")
+    gr.Markdown(description)
+    input_components = []
+    output_components = []
+
+    with gr.Row():
+        input_video = gr.Video(type="file", label="input_video")
+        output_video = gr.Video(label="output_video")
+        input_components.append(input_video)
+        output_components.append(output_video)
+
+    with gr.Row().style(equal_height=True, mobile_collapse=True):
+        with gr.Column(scale=2, variant="panel") as input_column:
+            model_dropdown = gr.Dropdown(label="Detector Model", 
+                                         choices=["pedestrian", "general-human"], 
+                                         default="general-human", 
+                                         info="Choose the application model for Human detection.")
+            prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.3, label="Confidence Threshold")
+            show_confidence = gr.Checkbox(label="Show Confidence")
+            input_components.extend([model_dropdown, prob_threshold_slider, show_confidence])
+        
+        with gr.Column(scale=2):
+            examples_handler = gr.Examples(
+                examples=[[item] for item in example_video_paths],
+                fn=inference,
+                inputs=input_components,
+                outputs=output_components,
+                examples_per_page=3
+            )
+            
+            with gr.Row():
+                submit_btn = gr.Button("Submit", variant="primary")
+                clear_btn = gr.Button("Clear")
+    
+    submit_btn.click(
+        inference,
+        input_components,
+        output_components,
+        api_name="predict",
+        scroll_to_output=True,
+    )
+
+    clear_btn.click(
+        clear,
+        input_components + output_components,
+        input_components + output_components,
+    )
+
+demo.queue(concurrency_count=3).launch(share=True, server_name='0.0.0.0', server_port=7860)

projects/human_detection/docker_run.sh

@@ -0,0 +1,11 @@
+# Docker run to start run demo app for traffic_monitoring demo
+read -p "Please enter your container name, for example 'cc-demo': " name
+read -p "Please enter your data directory path, for example /data/cc-demo/: " data
+read -p "Please enter your public port, for example 8585: " pubport
+
+docker run --name $name --shm-size=8g --gpus all --rm -it \
+    -p $pubport:7860 \
+    -v $data:/data \
+    -v $(pwd):/root/workspace/cc-demo \
+    -w /root/workspace/cc-demo \
+    cybercorecloud/cc-demo:v2.1 /bin/bash -c "cd projects && gradio human_detection/demo_app.py"

projects/human_detection/engine/pipeline.py

@@ -0,0 +1,81 @@
+from queue import Queue
+from typing import List, Dict, Tuple
+from threading import Thread, Event
+import os
+import logging
+import sys
+from models.engine.threading_func import batch_extract_thread, detect_thread, bytetrack_thread, update_progress_thread
+from .threading_func import visualize_thread
+from models.detectors.yolov7 import YOLOv7TRT
+import gradio as gr
+
+def init_logger(task: str, logging_level=logging.INFO):
+    log_format = '[%(levelname)s][%(process)d] [%(threadName)s] [%(asctime)s] %(message)s'
+
+    # Console log handler
+    console_handler = logging.StreamHandler(stream=sys.stdout)       
+    
+    logging.basicConfig(level=logging_level,
+                format=log_format,
+                datefmt='%d/%b/%Y %H:%M:%S',
+                handlers=[console_handler])
+    
+    return logging.getLogger(task)
+
+def run_e2e_pipeline(video_path: str, 
+                     detector: YOLOv7TRT, 
+                     tracker_cfg: Dict, 
+                     visualizer_cfg: Dict, 
+                     output_path: str, 
+                     conf_thres: float, 
+                     show_conf: bool=True, 
+                     progress: gr.Progress=None, 
+                     batch_size: int=32):
+    """ Function to run full pipelines (batch_extract, detection, tracking, counting, visualization, update progress). It will start threads and wait threads finish.
+
+    Args:
+        video_path (str): input video path that needs to be processed.
+        detector (YOLOv7TRT): Detector to run detection on input video.
+        tracker_cfg (Dict): tracker's config.
+        visualizer_cfg (Dict): visualizer's config.
+        output_path (str): path of output video.
+        conf_thres (float): If a tracklet's confidence < confidence threshold, it will be removed. 
+        show_conf (bool): Visualize confidence of track results or not.
+        progress (gr.Progress): Gradio's Progress that needs to be updated.
+        batch_size (int): Size of a batch input images that needs to be processed.
+    """
+
+    # init logger
+    init_logger("Infer")
+    video_name = os.path.basename(video_path)
+
+    img_batch_queue = Queue(maxsize=2)
+    vis_img_batch_queue = Queue(maxsize=2)
+    det_queue = Queue()
+    track_queue = Queue()
+    vis_queue = Queue()
+    eStop = Event()
+
+    pipeline = []
+    pipeline.append(Thread(target=batch_extract_thread,
+                        name=f'DT Batch Thread-Video Name {video_name}',
+                        args=(video_path, img_batch_queue, vis_img_batch_queue, eStop),
+                        kwargs={"batch_size":batch_size}))
+    pipeline.append(Thread(target=detect_thread,
+                        name=f'Detect Thread-Video Name {video_name}',
+                        args=(detector, img_batch_queue, det_queue, eStop)))
+    pipeline.append(Thread(target=bytetrack_thread,
+                        name=f'Track Thread-Video Name {video_name}',
+                        args=(tracker_cfg, det_queue, track_queue, eStop, conf_thres)))
+    pipeline.append(Thread(target=visualize_thread,
+                        name=f'Visualize Thread-Video Name {video_name}',
+                        args=(visualizer_cfg, output_path, vis_img_batch_queue, track_queue, vis_queue, eStop, show_conf)))
+    if progress is not None:
+        pipeline.append(Thread(target=update_progress_thread,
+                        name=f"Update Progress Thread-Video Name {video_name}",
+                        args=(vis_queue, progress, eStop)))
+
+    for p in pipeline: p.start()
+    for p in pipeline: p.join()
+    del eStop
+    del pipeline

projects/human_detection/engine/threading_func.py

@@ -0,0 +1,43 @@
+from queue import Queue
+from threading import Event
+from models.engine.threading_func import queue_clear
+
+from projects.human_detection.engine.visualizer import Visualizer
+import logging
+
+
+def visualize_thread(visualizer_cfg, output_path, vis_img_batch_queue: Queue, track_queue: Queue, visualize_queue: Queue, eStop: Event, show_conf):
+    logging.info("Start Visualize Thread")
+    visualizer = Visualizer(**visualizer_cfg)
+    input_video_info = vis_img_batch_queue.get()
+    visualizer.init_writer(input_video_info, output_path)
+    visualize_queue.put(input_video_info)
+    track_item = track_queue.get()
+    start_frame_idx = -1
+    img_batch_item = vis_img_batch_queue.get()
+    while (img_batch_item is not None and track_item is not None):
+        if eStop.is_set():  break
+        start_frame_idx, img_batch = img_batch_item
+        track_start_frame_idx, track_result = track_item
+        if (start_frame_idx != track_start_frame_idx) or (len(img_batch) != len(track_result)):
+            error_msg=[501, f"Error when runing visualization at start_frame_idx {start_frame_idx}. "]
+            log_error_message = f"Error {error_msg[0]}: {error_msg[1]}"
+            logging.error(log_error_message)
+            eStop.set()
+            break
+        for idx, (frame, frame_track_result) in enumerate(zip(img_batch, track_result)):
+
+            visualizer.visualize(frame, frame_track_result, show_conf)
+            visualize_queue.put(start_frame_idx + idx)
+        img_batch_item = vis_img_batch_queue.get()
+        track_item = track_queue.get()
+    visualizer.close()
+
+    # Finish this thread    
+    if eStop.is_set():
+        logging.warning(f"Early stop at start_frame_idx {start_frame_idx}")
+        queue_clear(visualize_queue)
+    else:
+        visualizer.convert()
+        logging.info(f"Finish visualize_thread.")
+    visualize_queue.put(None)

projects/human_detection/engine/visualizer.py

@@ -0,0 +1,51 @@
+from typing import List, Dict
+import cv2
+import numpy as np
+from models.engine.visualizer import BaseVisualizer
+
+class Visualizer(BaseVisualizer):
+
+    def __init__(self, fps: int=-1, min_width: int=-1):
+        """ Visualizer class for visualization (track_results + count_results).
+
+        Args:
+            class_map_ids (Dict): class mapping dictionary to map model's class to original class. Eg {0: 1, 1: 0, 2: 2, 3: 3} mean we swap class ID between 0 and 1.
+            fps (int): FPS for output video. If fps = -1, it will have same fps as input video.
+            min_width (int): minimum width for output video (height will be scaled to keep aspect ratio as input video). If min_width = -1, it will have same resolution as input video.
+        """
+        class_names = ['pedestrian']
+        super().__init__(class_names, fps, min_width)
+
+    def visualize(self, img: np.ndarray, dettrack_at_frame_id: List[Dict]=None, show_conf: bool=True):
+        """ Function to visualize (track_results + count_results) a frame.
+
+        Args:
+            img (np.ndarray): image need to be visualized.
+            dettrack_at_frame_id (List[Dict]): batch of track results which can be obtained from Tracker class.
+            show_conf (bool): Visualize confidence of track results or not.
+        """
+        
+         # Draw tracking
+        if (dettrack_at_frame_id):
+            boxes = dettrack_at_frame_id["boxes"]
+            # classes = dettrack_at_frame_id["labels"]
+            ids = dettrack_at_frame_id["ids"]
+            for bbox, id_ in zip(boxes, ids): 
+                id_ = int(id_)
+                score = bbox[4]
+                color = self.get_color(id_)
+                label = f'{id_}' + (f' {score:.2f}' if (show_conf) else '')
+                tl, tf = 2, 1
+                c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))
+                img = cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
+                t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
+                c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
+                img = cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)
+                img = cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+
+        if (img.shape[0] != self.height or img.shape[1] != self.width):
+            img = cv2.resize(img, (self.width, self.height))
+        self.writer.write(img)
+
+        return img
+