diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000000000000000000000000000000000..940297d70d78f5fc487b2f33a7e78378d49b4903 --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..dfdf65ebb19ef854cd36013471aeffc9fab49c8d --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..49bfe604fc4227a541fd28e9b78eb9298e676adc --- /dev/null +++ b/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" diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/base/__init__.py b/models/base/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0c66f1b39c56bbb29bc725bcc224a67c746f4c0e --- /dev/null +++ b/models/base/__init__.py @@ -0,0 +1,2 @@ +from .trt_base import TRT_Base +from .onnx_base import ONNX_Base \ No newline at end of file diff --git a/models/base/onnx_base.py b/models/base/onnx_base.py new file mode 100644 index 0000000000000000000000000000000000000000..3162588e322a04f0c1bd917b1781b90bf4e7f391 --- /dev/null +++ b/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 + + + + + + + + + \ No newline at end of file diff --git a/models/base/trt_base.py b/models/base/trt_base.py new file mode 100644 index 0000000000000000000000000000000000000000..36fe444d1cdbd026bbfcb42126557ab46b706c02 --- /dev/null +++ b/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 + + \ No newline at end of file diff --git a/models/detectors/__init__.py b/models/detectors/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/detectors/yolov7.py b/models/detectors/yolov7.py new file mode 100755 index 0000000000000000000000000000000000000000..378551cdf59c3b13bdace9793a23dc927cfbbc3f --- /dev/null +++ b/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 + + + + \ No newline at end of file diff --git a/models/models/__init__.py b/models/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/models/base/__init__.py b/models/models/base/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/models/base/onnx_base.py b/models/models/base/onnx_base.py new file mode 100644 index 0000000000000000000000000000000000000000..d3763bbbb9377669ebf37d922969f2c60bb89283 --- /dev/null +++ b/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 + + + + + + + + + \ No newline at end of file diff --git a/models/models/base/trt_base.py b/models/models/base/trt_base.py new file mode 100644 index 0000000000000000000000000000000000000000..ea7b2c3ebbd392f294a2fc9db802d175206a1f48 --- /dev/null +++ b/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 + + \ No newline at end of file diff --git a/models/models/detectors/__init__.py b/models/models/detectors/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/models/detectors/mmyolov8.py b/models/models/detectors/mmyolov8.py new file mode 100644 index 0000000000000000000000000000000000000000..5814748777cd3e262b268d1de9ee7637f4df79c9 --- /dev/null +++ b/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) diff --git a/models/models/detectors/yolov7.py b/models/models/detectors/yolov7.py new file mode 100755 index 0000000000000000000000000000000000000000..a3804ef1ba40544636f290e5a2fa2c3850ab2b59 --- /dev/null +++ b/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) + + + + + + + \ No newline at end of file diff --git a/models/models/engine/__init__.py b/models/models/engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/models/engine/threading_func.py b/models/models/engine/threading_func.py new file mode 100755 index 0000000000000000000000000000000000000000..606fc394e73b8e5c655643f63b0c5c65b34a4da2 --- /dev/null +++ b/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.") \ No newline at end of file diff --git a/models/models/engine/utils.py b/models/models/engine/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..90e436f81791c93a22fdd36e46883a31fe2258bd --- /dev/null +++ b/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 \ No newline at end of file diff --git a/models/models/engine/visualizer.py b/models/models/engine/visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..cdb6af3e2f76d6f239638662acae9d3b461d7959 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/models/models/pose/rtmpose.py b/models/models/pose/rtmpose.py new file mode 100644 index 0000000000000000000000000000000000000000..c489ceb1128ea87443dff0595d6316ac1afd934d --- /dev/null +++ b/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} + + \ No newline at end of file diff --git a/models/models/reids/__init__.py b/models/models/reids/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/models/reids/solider.py b/models/models/reids/solider.py new file mode 100755 index 0000000000000000000000000000000000000000..c8757ecee429e8549abd4c95886109156fdc3063 --- /dev/null +++ b/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 + + diff --git a/models/models/trackers/__init__.py b/models/models/trackers/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..32e37689212641e90127b7fd01ecebced326508b --- /dev/null +++ b/models/models/trackers/__init__.py @@ -0,0 +1,2 @@ +from .byte_track import BYTETracker +from .reid_parallel_tracker import ParallelTracker \ No newline at end of file diff --git a/models/models/trackers/byte_track.py b/models/models/trackers/byte_track.py new file mode 100755 index 0000000000000000000000000000000000000000..2a317fbca3129c43afc511cb2b03f198d24c24bb --- /dev/null +++ b/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 + diff --git a/models/models/trackers/reid_parallel_tracker/__init__.py b/models/models/trackers/reid_parallel_tracker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c4fa60fdf4ef37f8153776ea20c69ee9bdfa5244 --- /dev/null +++ b/models/models/trackers/reid_parallel_tracker/__init__.py @@ -0,0 +1 @@ +from .parallel_tracker import ParallelTracker \ No newline at end of file diff --git a/models/models/trackers/reid_parallel_tracker/base_tracker.py b/models/models/trackers/reid_parallel_tracker/base_tracker.py new file mode 100755 index 0000000000000000000000000000000000000000..c8d505b72f309c9f94a4f04f705e169c88efc290 --- /dev/null +++ b/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 diff --git a/models/models/trackers/reid_parallel_tracker/core/__init__.py b/models/models/trackers/reid_parallel_tracker/core/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/models/trackers/reid_parallel_tracker/core/basetrack.py b/models/models/trackers/reid_parallel_tracker/core/basetrack.py new file mode 100755 index 0000000000000000000000000000000000000000..6ab32243d08c44b0959fe987370f6da8f658653a --- /dev/null +++ b/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) diff --git a/models/models/trackers/reid_parallel_tracker/core/homography.py b/models/models/trackers/reid_parallel_tracker/core/homography.py new file mode 100755 index 0000000000000000000000000000000000000000..2277c54c83ddeb9e714c2b867b8f47611b015348 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/models/models/trackers/reid_parallel_tracker/core/kalman_filter.py b/models/models/trackers/reid_parallel_tracker/core/kalman_filter.py new file mode 100755 index 0000000000000000000000000000000000000000..6aa7caad77e512aa3438769f1e02997b02b2ea24 --- /dev/null +++ b/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') \ No newline at end of file diff --git a/models/models/trackers/reid_parallel_tracker/core/matching.py b/models/models/trackers/reid_parallel_tracker/core/matching.py new file mode 100755 index 0000000000000000000000000000000000000000..ce0a6566b80662abd9d6fb538645d405bd426b8e --- /dev/null +++ b/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) \ No newline at end of file diff --git a/models/models/trackers/reid_parallel_tracker/core/tracklet.py b/models/models/trackers/reid_parallel_tracker/core/tracklet.py new file mode 100755 index 0000000000000000000000000000000000000000..d1c892cd44cc21dfe0321d7a150b94e3f2739948 --- /dev/null +++ b/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 diff --git a/models/models/trackers/reid_parallel_tracker/matchers/__init__.py b/models/models/trackers/reid_parallel_tracker/matchers/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..089983f254e4b0d0cb0464187f4df330286edfc6 --- /dev/null +++ b/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 diff --git a/models/models/trackers/reid_parallel_tracker/matchers/base_matchers.py b/models/models/trackers/reid_parallel_tracker/matchers/base_matchers.py new file mode 100755 index 0000000000000000000000000000000000000000..db3b0fe339a4d4e0cb229f5dad0ac177e7716b94 --- /dev/null +++ b/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})" \ No newline at end of file diff --git a/models/models/trackers/reid_parallel_tracker/matchers/distances.py b/models/models/trackers/reid_parallel_tracker/matchers/distances.py new file mode 100755 index 0000000000000000000000000000000000000000..62e22dd130674c81cfbaa2999d31d993faaa845b --- /dev/null +++ b/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') + diff --git a/models/models/trackers/reid_parallel_tracker/matchers/prioritize_reid_matcher.py b/models/models/trackers/reid_parallel_tracker/matchers/prioritize_reid_matcher.py new file mode 100755 index 0000000000000000000000000000000000000000..1524e72be792def9361b4343e428db1c822806a3 --- /dev/null +++ b/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) + diff --git a/models/models/trackers/reid_parallel_tracker/matchers/single_stage_matcher.py b/models/models/trackers/reid_parallel_tracker/matchers/single_stage_matcher.py new file mode 100755 index 0000000000000000000000000000000000000000..ab2dfb9d1cb73f952929f1f45439a135aa3e75d0 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/models/models/trackers/reid_parallel_tracker/parallel_tracker.py b/models/models/trackers/reid_parallel_tracker/parallel_tracker.py new file mode 100755 index 0000000000000000000000000000000000000000..ff0b7beff7ef873e989e93e41dbbca9eac4e6c69 --- /dev/null +++ b/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 diff --git a/models/models/trackers/reid_parallel_tracker/three_stage_tracker.py b/models/models/trackers/reid_parallel_tracker/three_stage_tracker.py new file mode 100755 index 0000000000000000000000000000000000000000..6dcc1c29485749e61931ef733b6b4904483431d1 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/models/reids/__init__.py b/models/reids/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/reids/solider.py b/models/reids/solider.py new file mode 100644 index 0000000000000000000000000000000000000000..7cf85c4b6ee0440db1590390f96e81bb325dfb2e --- /dev/null +++ b/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 + + diff --git a/models/trackers/__init__.py b/models/trackers/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/trackers/byte_track.py b/models/trackers/byte_track.py new file mode 100755 index 0000000000000000000000000000000000000000..b1c7425d2e726882197c6c9bb2dc5daac8409130 --- /dev/null +++ b/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 + diff --git a/projects/human_detection/ReadMe.md b/projects/human_detection/ReadMe.md new file mode 100644 index 0000000000000000000000000000000000000000..0a22d209d50a74b586abad81db76dadc5c6fc876 --- /dev/null +++ b/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 +``` +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 +``` \ No newline at end of file diff --git a/projects/human_detection/demo_app.py b/projects/human_detection/demo_app.py new file mode 100644 index 0000000000000000000000000000000000000000..0e38dbe571c03278eb962d52cd1a2ee07292921c --- /dev/null +++ b/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("

" + title + "

") + 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) diff --git a/projects/human_detection/docker_run.sh b/projects/human_detection/docker_run.sh new file mode 100644 index 0000000000000000000000000000000000000000..bb555cf7750369eb22ba1c3e9f8b892afc3331ef --- /dev/null +++ b/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" \ No newline at end of file diff --git a/projects/human_detection/engine/pipeline.py b/projects/human_detection/engine/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..04079833a8e63ad1dab1e75d438df78d5424c2f4 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/projects/human_detection/engine/threading_func.py b/projects/human_detection/engine/threading_func.py new file mode 100644 index 0000000000000000000000000000000000000000..fa630b524dc05a725fd7d9384d5df772cbaade5d --- /dev/null +++ b/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) diff --git a/projects/human_detection/engine/visualizer.py b/projects/human_detection/engine/visualizer.py new file mode 100755 index 0000000000000000000000000000000000000000..992d2d3caea773bed872c00b712114367cc516c1 --- /dev/null +++ b/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 + diff --git a/projects/human_detection/export_onnx_trt/detection_tensorrt-fp16_dynamic.py b/projects/human_detection/export_onnx_trt/detection_tensorrt-fp16_dynamic.py new file mode 100644 index 0000000000000000000000000000000000000000..894efc7591795da05dc84167e56ad8a324b4c745 --- /dev/null +++ b/projects/human_detection/export_onnx_trt/detection_tensorrt-fp16_dynamic.py @@ -0,0 +1,36 @@ +_base_ = ['mmyolo::deploy/base_dynamic.py'] + +onnx_config = dict( + dynamic_axes={ + 'input': { + 0: 'batch', + 2: 'height', + 3: 'width' + }, + 'dets': { + 0: 'batch', + 1: 'num_dets' + }, + 'labels': { + 0: 'batch', + 1: 'num_dets' + } + }, + ) +# input_shape=[640, 800] +backend_config = dict( + type='tensorrt', + common_config=dict(fp16_mode=True, max_workspace_size=1 << 40), + model_inputs=[ + dict( + input_shapes=dict( + input=dict( + # Change into any shape you want, but recommend to use standard shape (480,640), (640,800), (768,1280), + # min_shape=[1, 3, 192, 192], + # opt_shape=[1, 3, 640, 640], + # max_shape=[1, 3, 960, 960]))) + min_shape=[1, 3, 640, 800], + opt_shape=[32, 3, 640, 800], + max_shape=[32, 3, 640, 800]))) + ]) +use_efficientnms = True # whether native NMS of TRT instead of plugin mmdeploy:TRTBatchedNMS # noqa E501 \ No newline at end of file diff --git a/projects/human_detection/export_onnx_trt/export_onnx_mmyolov8.sh b/projects/human_detection/export_onnx_trt/export_onnx_mmyolov8.sh new file mode 100644 index 0000000000000000000000000000000000000000..8cf45e6ad1430f3d7a11fe6d01104b224ecaeeee --- /dev/null +++ b/projects/human_detection/export_onnx_trt/export_onnx_mmyolov8.sh @@ -0,0 +1,28 @@ +ENV_DIR="/opt/conda/lib/python3.8/site-packages" +CC_DEMO_DIR="/root/workspace/cc-demo" + + +DEPLOY_CFG_PATH="${ENV_DIR}/mmyolo/.mim/configs/deploy/detection_onnxruntime_dynamic.py" +MODEL_CFG_PATH="${CC_DEMO_DIR}/projects/human_detection/deploy/mmyolov8_human_cfg.py" +MODEL_CHECKPOINT_PATH="/data/human_detection/weights/mmyolov8_s_human.pth" +WORK_DIR="/data/human_detection/deploy_onnx" + +INPUT_IMG="${CC_DEMO_DIR}/tests/test_data/human_det/1.jpg" +TEST_IMG="${CC_DEMO_DIR}/tests/test_data/human_det/2.jpg" +DEVICE="cpu" + +export ONNXRUNTIME_DIR=/root/workspace/onnxruntime-linux-x64-1.15.1/ +export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH + +cd /root/workspace/mmdeploy && python tools/deploy.py \ + ${DEPLOY_CFG_PATH} \ + ${MODEL_CFG_PATH} \ + ${MODEL_CHECKPOINT_PATH} \ + ${INPUT_IMG} \ + --test-img ${TEST_IMG} \ + --work-dir ${WORK_DIR} \ + --device ${DEVICE} \ + --log-level INFO \ + --show \ + --dump-info + # --calib-dataset-cfg ${CALIB_DATA_CFG} \ \ No newline at end of file diff --git a/projects/human_detection/export_onnx_trt/export_onnx_yolov7.sh b/projects/human_detection/export_onnx_trt/export_onnx_yolov7.sh new file mode 100644 index 0000000000000000000000000000000000000000..3ba111eb1eb461afe817ed3bab68ae932af3f0a6 --- /dev/null +++ b/projects/human_detection/export_onnx_trt/export_onnx_yolov7.sh @@ -0,0 +1,10 @@ +export ONNXRUNTIME_DIR=/root/workspace/onnxruntime-linux-x64-1.15.1/ +export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH + +CKPT="/data/human_detection/weights/yolov7_pedestrian.pt" + +cd /root/workspace/yolov7 && python export.py \ + --weights $CKPT \ + --dynamic --dynamic-batch --grid --end2end --simplify --device 0 --fp16 \ + --topk-all 100 --iou-thres 0.65 --conf-thres 0.1 --img-size 480 640 --max-wh 640 + # --include-nms \ No newline at end of file diff --git a/projects/human_detection/export_onnx_trt/export_trt_mmyolov8.sh b/projects/human_detection/export_onnx_trt/export_trt_mmyolov8.sh new file mode 100644 index 0000000000000000000000000000000000000000..8f76e3cc96e9e55f340085c4307138dabe968fc1 --- /dev/null +++ b/projects/human_detection/export_onnx_trt/export_trt_mmyolov8.sh @@ -0,0 +1,24 @@ +CC_DEMO_DIR="/root/workspace/cc-demo" + +DEPLOY_CFG_PATH="${CC_DEMO_DIR}/projects/human_detection/deploy/detection_tensorrt-fp16_dynamic.py" +MODEL_CFG_PATH="${CC_DEMO_DIR}/projects/human_detection/deploy/mmyolov8_human_cfg.py" +MODEL_CHECKPOINT_PATH="/data/human_detection/weights/yolov8_s_human_mmyolo.pth" +WORK_DIR="/data/human_detection/deploy" + +INPUT_IMG="${CC_DEMO_DIR}/tests/test_data/human_det/1.jpg" +TEST_IMG="${CC_DEMO_DIR}/tests/test_data/human_det/2.jpg" +DEVICE="cuda:0" + +export ONNXRUNTIME_DIR=/root/workspace/onnxruntime-linux-x64-1.15.1/ +export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH + +cd /root/workspace/mmdeploy && python tools/deploy.py \ + ${DEPLOY_CFG_PATH} \ + ${MODEL_CFG_PATH} \ + ${MODEL_CHECKPOINT_PATH} \ + ${INPUT_IMG} \ + --test-img ${TEST_IMG} \ + --work-dir ${WORK_DIR} \ + --device ${DEVICE} \ + --log-level INFO \ + --dump-info diff --git a/projects/human_detection/export_onnx_trt/export_trt_yolov7.sh b/projects/human_detection/export_onnx_trt/export_trt_yolov7.sh new file mode 100755 index 0000000000000000000000000000000000000000..92387e1f82bbb0723805189fa27d00735e428114 --- /dev/null +++ b/projects/human_detection/export_onnx_trt/export_trt_yolov7.sh @@ -0,0 +1,73 @@ +echo "Input weight ONNX path can be downloaded by link: " +echo "1. yolov7_pedestrian_480x640.onnx: https://gofile.me/6ZWyr/gAymnT4XJ" +echo "2. yolov7_pedestrian_768x1280.onnx: https://gofile.me/6ZWyr/yr1OwvkIi" +read -p "Please enter input weight ONNX path: " weight_path +read -p "Please choose the resolution : 1. 480x640, 2. 768x1280: " choice + + +weight_dir=$(dirname "$weight_path") +file_name=$(basename "$weight_path") +new_filename="${file_name%.*}.trt" +echo "Weight directory: $weight_dir" +echo "File name: $file_name" +echo "New file name: $new_filename" + +default_params=$(cat <<-END + --rm \ + --gpus all \ + -v $weight_dir:/weights \ + cybercorecloud/cc-demo:v2.1 \ + /opt/tensorrt/bin/trtexec +END +) + +IN_DOCKER=1 #1 if run in docker, 0 if run in host +if [[ $IN_DOCKER==1 ]]; then + if [[ $choice == "1" ]]; then + echo "Exporting TRT model with resolution 480x640..." + /opt/tensorrt/bin/trtexec \ + --onnx=$weight_path \ + --minShapes=images:1x3x480x640 \ + --optShapes=images:1x3x480x640 \ + --maxShapes=images:32x3x480x640 \ + --fp16 --saveEngine=$weight_dir/$new_filename --workspace=48000 --timingCacheFile=timing.cache + elif [[ $choice == "2" ]]; then + echo "Exporting TRT model with resolution 768x1280..." + /opt/tensorrt/bin/trtexec \ + --onnx=$weight_path \ + --minShapes=images:1x3x768x1280 \ + --optShapes=images:1x3x768x1280 \ + --maxShapes=images:32x3x768x1280 \ + --fp16 --saveEngine=$weight_dir/$new_filename --workspace=48000 --timingCacheFile=timing.cache + else + echo "Not support" + exit 1 + fi +else + docker pull cybercorecloud/cc-demo:v2.1 + + if [[ $choice == "1" ]]; then + echo "Exporting TRT model with resolution 480x640..." + docker run \ + $default_params \ + --onnx=/weights/$file_name \ + --minShapes=images:1x3x480x640 \ + --optShapes=images:32x3x480x640 \ + --maxShapes=images:32x3x480x640 \ + --fp16 --saveEngine=/weights/$new_filename --workspace=48000 --timingCacheFile=timing.cache + elif [[ $choice == "2" ]]; then + echo "Exporting TRT model with resolution 768x1280..." + docker run \ + $default_params \ + --onnx=/weights/$file_name \ + --minShapes=images:1x3x768x1280 \ + --optShapes=images:32x3x768x1280 \ + --maxShapes=images:32x3x768x1280 \ + --fp16 --saveEngine=/weights/$new_filename --workspace=48000 --timingCacheFile=timing.cache + else + echo "Not support" + exit 1 + fi +fi + + diff --git a/projects/human_detection/export_onnx_trt/mmyolov8_human_cfg.py b/projects/human_detection/export_onnx_trt/mmyolov8_human_cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..ab6ef0f8ad9d482a6bd0a847282d6a0961932cad --- /dev/null +++ b/projects/human_detection/export_onnx_trt/mmyolov8_human_cfg.py @@ -0,0 +1,28 @@ +_base_ = 'mmyolo::yolov8/yolov8_s_syncbn_fast_8xb16-500e_coco.py' + +img_scale = (800, 640) #(W,H) +test_pipeline = [ + # dict(type='LoadImageFromFile', backend_args=_base_.backend_args), + dict( + type='LetterResize', + scale=img_scale, + allow_scale_up=False, + use_mini_pad=False, + ), + # dict(type='LoadAnnotations', with_bbox=True, _scope_='mmdet'), + dict( + type='mmdet.PackDetInputs', + meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape', + 'scale_factor', 'pad_param')) +] + +test_dataloader = dict( + dataset=dict(pipeline=test_pipeline, batch_shapes_cfg=None)) + +model = dict( + bbox_head=dict( + head_module=dict( + num_classes=1,) + ) +) + diff --git a/projects/scripts/start.sh b/projects/scripts/start.sh new file mode 100755 index 0000000000000000000000000000000000000000..a7136529c980d4e289a810ec687f380f2f8063ab --- /dev/null +++ b/projects/scripts/start.sh @@ -0,0 +1,2 @@ +export CUDA_VISIBLE_DEVICES=3 +gradio traffic_monitoring/demo_app.py \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..a4b562219351e7c6731bf37e225676ebb3582deb --- /dev/null +++ b/requirements.txt @@ -0,0 +1,20 @@ +# Usage: pip install -r requirements.txt + + +# Base ---------------------------------------- +numpy==1.22.0 +opencv-python==4.6.0.66 +# opencv-python>=4.1.1 +Pillow>=7.1.2 +PyYAML>=5.3.1 +scipy>=1.4.1 +torch==2.0.1 +tqdm>=4.41.0 +protobuf<4.21.3 +openmim +shapely==1.8.5 +Cython +psutil +onnx>=1.14.0 +onnxruntime-gpu>=1.15.1 +pycuda \ No newline at end of file