[
{
"path": ".gitignore",
"content": "docs/test_video.mp4\ntest/test_tao.py\nconfig/tao*\ntracker/mot/tao.py\neval/error_log.txt\nconfig/got10k*\nconfig/lasot*\nconfig/tc128*\nconfig/tlp*\nconfig/trackingnet*\nconfig/vfs*\nconfig/ssib*\n\nweights/\nresults/\nout/\nvis/\n*.ipynb\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n"
},
{
"path": "LICENSE",
"content": "MIT License\n\nCopyright (c) 2021 ZhongdaoWang\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "README.md",
"content": "
testval\nInstallation \n\nStep1. Install YOLOX.\n```shell\ngit clone git@github.com:Megvii-BaseDetection/YOLOX.git\ncd YOLOX\npip3 install -U pip && pip3 install -r requirements.txt\npip3 install -v -e . # or python3 setup.py develop\n```\nStep2. Install [apex](https://github.com/NVIDIA/apex).\n\n```shell\n# skip this step if you don't want to train model.\ngit clone https://github.com/NVIDIA/apex\ncd apex\npip3 install -v --disable-pip-version-check --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./\n```\nStep3. Install [pycocotools](https://github.com/cocodataset/cocoapi).\n\n```shell\npip3 install cython; pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'\n```\n\n \n\n\nDemo \n\nStep1. Download a pretrained model from the benchmark table.\n\nStep2. Use either -n or -f to specify your detector's config. For example:\n\n```shell\npython tools/demo.py image -n yolox-s -c /path/to/your/yolox_s.pth.tar --path assets/dog.jpg --conf 0.3 --nms 0.65 --tsize 640 --save_result --device [cpu/gpu]\n```\nor\n```shell\npython tools/demo.py image -f exps/default/yolox_s.py -c /path/to/your/yolox_s.pth.tar --path assets/dog.jpg --conf 0.3 --nms 0.65 --tsize 640 --save_result --device [cpu/gpu]\n```\nDemo for video:\n```shell\npython tools/demo.py video -n yolox-s -c /path/to/your/yolox_s.pth.tar --path /path/to/your/video --conf 0.3 --nms 0.65 --tsize 640 --save_result --device [cpu/gpu]\n```\n\n\n \n\n\nReproduce our results on COCO \n\nStep1. Prepare COCO dataset\n```shell\ncd \nln -s /path/to/your/COCO ./datasets/COCO\n```\n\nStep2. Reproduce our results on COCO by specifying -n:\n\n```shell\npython tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o\n yolox-m\n yolox-l\n yolox-x\n```\n* -d: number of gpu devices\n* -b: total batch size, the recommended number for -b is num-gpu * 8\n* --fp16: mixed precision training\n\nWhen using -f, the above commands are equivalent to:\n\n```shell\npython tools/train.py -f exps/default/yolox-s.py -d 8 -b 64 --fp16 -o\n exps/default/yolox-m.py\n exps/default/yolox-l.py\n exps/default/yolox-x.py\n```\n\n \n\n\n\nEvaluation \n\nWe support batch testing for fast evaluation:\n\n```shell\npython tools/eval.py -n yolox-s -c yolox_s.pth.tar -b 64 -d 8 --conf 0.001 [--fp16] [--fuse]\n yolox-m\n yolox-l\n yolox-x\n```\n* --fuse: fuse conv and bn\n* -d: number of GPUs used for evaluation. DEFAULT: All GPUs available will be used.\n* -b: total batch size across on all GPUs\n\nTo reproduce speed test, we use the following command:\n```shell\npython tools/eval.py -n yolox-s -c yolox_s.pth.tar -b 1 -d 1 --conf 0.001 --fp16 --fuse\n yolox-m\n yolox-l\n yolox-x\n```\n\n \n\n\n\nTutorials \n\n* [Training on custom data](docs/train_custom_data.md).\n\n \n\n## Deployment\n\n\n1. [ONNX export and an ONNXRuntime](./demo/ONNXRuntime)\n2. [TensorRT in C++ and Python](./demo/TensorRT)\n3. [ncnn in C++ and Java](./demo/ncnn)\n4. [OpenVINO in C++ and Python](./demo/OpenVINO)\n\n\n## Third-party resources\n* The ncnn android app with video support: [ncnn-android-yolox](https://github.com/FeiGeChuanShu/ncnn-android-yolox) from [FeiGeChuanShu](https://github.com/FeiGeChuanShu)\n* YOLOX with Tengine support: [Tengine](https://github.com/OAID/Tengine/blob/tengine-lite/examples/tm_yolox.cpp) from [BUG1989](https://github.com/BUG1989)\n* YOLOX + ROS2 Foxy: [YOLOX-ROS](https://github.com/Ar-Ray-code/YOLOX-ROS) from [Ar-Ray](https://github.com/Ar-Ray-code)\n* YOLOX Deploy DeepStream: [YOLOX-deepstream](https://github.com/nanmi/YOLOX-deepstream) from [nanmi](https://github.com/nanmi)\n* YOLOX ONNXRuntime C++ Demo: [lite.ai](https://github.com/DefTruth/lite.ai/blob/main/ort/cv/yolox.cpp) from [DefTruth](https://github.com/DefTruth)\n\n## Cite YOLOX\nIf you use YOLOX in your research, please cite our work by using the following BibTeX entry:\n\n```latex\n @article{yolox2021,\n title={YOLOX: Exceeding YOLO Series in 2021},\n author={Ge, Zheng and Liu, Songtao and Wang, Feng and Li, Zeming and Sun, Jian},\n journal={arXiv preprint arXiv:2107.08430},\n year={2021}\n}\n```\n"
},
{
"path": "detector/YOLOX/datasets/README.md",
"content": "# Prepare datasets\n\nIf you have a dataset directory, you could use os environment variable named `YOLOX_DATADIR`. Under this directory, YOLOX will look for datasets in the structure described below, if needed.\n```\n$YOLOX_DATADIR/\n COCO/\n```\nYou can set the location for builtin datasets by\n```shell\nexport YOLOX_DATADIR=/path/to/your/datasets\n```\nIf `YOLOX_DATADIR` is not set, the default value of dataset directory is `./datasets` relative to your current working directory.\n\n## Expected dataset structure for [COCO detection](https://cocodataset.org/#download):\n\n```\nCOCO/\n annotations/\n instances_{train,val}2017.json\n {train,val}2017/\n # image files that are mentioned in the corresponding json\n```\n\nYou can use the 2014 version of the dataset as well.\n"
},
{
"path": "detector/YOLOX/demo/ONNXRuntime/README.md",
"content": "## YOLOX-ONNXRuntime in Python\n\nThis doc introduces how to convert your pytorch model into onnx, and how to run an onnxruntime demo to verify your convertion.\n\n### Download ONNX models.\n| Model | Parameters | GFLOPs | Test Size | mAP | Weights |\n|:------| :----: | :----: | :---: | :---: | :---: |\n| YOLOX-Nano | 0.91M | 1.08 | 416x416 | 25.3 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EfAGwvevU-lNhW5OqFAyHbwBJdI_7EaKu5yU04fgF5BU7w?e=gvq4hf)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_nano.onnx) |\n| YOLOX-Tiny | 5.06M | 6.45 | 416x416 |31.7 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EVigCszU1ilDn-MwLwHCF1ABsgTy06xFdVgZ04Yyo4lHVA?e=hVKiCw)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_tiny.onnx) |\n| YOLOX-S | 9.0M | 26.8 | 640x640 |39.6 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/Ec0L1d1x2UtIpbfiahgxhtgBZVjb1NCXbotO8SCOdMqpQQ?e=siyIsK)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s.onnx) |\n| YOLOX-M | 25.3M | 73.8 | 640x640 |46.4 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ERUKlQe-nlxBoTKPy1ynbxsBmAZ_h-VBEV-nnfPdzUIkZQ?e=hyQQtl)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_m.onnx) |\n| YOLOX-L | 54.2M | 155.6 | 640x640 |50.0 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ET5w926jCA5GlVfg9ixB4KEBiW0HYl7SzaHNRaRG9dYO_A?e=ISmCYX)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_l.onnx) |\n| YOLOX-Darknet53| 63.72M | 185.3 | 640x640 |47.3 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ESArloSW-MlPlLuemLh9zKkBdovgweKbfu4zkvzKAp7pPQ?e=f81Ikw)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_darknet53.onnx) |\n| YOLOX-X | 99.1M | 281.9 | 640x640 |51.2 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ERjqoeMJlFdGuM3tQfXQmhABmGHlIHydWCwhlugeWLE9AA)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox.onnx) |\n\n\n### Convert Your Model to ONNX\n\nFirst, you should move to by:\n```shell\ncd \n```\nThen, you can:\n\n1. Convert a standard YOLOX model by -n:\n```shell\npython3 tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth.tar\n```\nNotes:\n* -n: specify a model name. The model name must be one of the [yolox-s,m,l,x and yolox-nane, yolox-tiny, yolov3]\n* -c: the model you have trained\n* -o: opset version, default 11. **However, if you will further convert your onnx model to [OpenVINO](../OpenVINO/), please specify the opset version to 10.**\n* --no-onnxsim: disable onnxsim\n* To customize an input shape for onnx model, modify the following code in tools/export.py:\n\n ```python\n dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])\n ```\n\n2. Convert a standard YOLOX model by -f. When using -f, the above command is equivalent to:\n\n```shell\npython3 tools/export_onnx.py --output-name yolox_s.onnx -f exps/default/yolox_s.py -c yolox_s.pth.tar\n```\n\n3. To convert your customized model, please use -f:\n\n```shell\npython3 tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth.tar\n```\n\n### ONNXRuntime Demo\n\nStep1.\n```shell\ncd /demo/ONNXRuntime\n```\n\nStep2. \n```shell\npython3 onnx_inference.py -m -i -o -s 0.3 --input_shape 640,640\n```\nNotes:\n* -m: your converted onnx model\n* -i: input_image\n* -s: score threshold for visualization.\n* --input_shape: should be consistent with the shape you used for onnx convertion.\n"
},
{
"path": "detector/YOLOX/demo/ONNXRuntime/onnx_inference.py",
"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport os\n\nimport cv2\nimport numpy as np\n\nimport onnxruntime\n\nfrom yolox.data.data_augment import preproc as preprocess\nfrom yolox.data.datasets import COCO_CLASSES\nfrom yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"onnxruntime inference sample\")\n parser.add_argument(\n \"-m\",\n \"--model\",\n type=str,\n default=\"yolox.onnx\",\n help=\"Input your onnx model.\",\n )\n parser.add_argument(\n \"-i\",\n \"--image_path\",\n type=str,\n default='test_image.png',\n help=\"Path to your input image.\",\n )\n parser.add_argument(\n \"-o\",\n \"--output_dir\",\n type=str,\n default='demo_output',\n help=\"Path to your output directory.\",\n )\n parser.add_argument(\n \"-s\",\n \"--score_thr\",\n type=float,\n default=0.3,\n help=\"Score threshould to filter the result.\",\n )\n parser.add_argument(\n \"--input_shape\",\n type=str,\n default=\"640,640\",\n help=\"Specify an input shape for inference.\",\n )\n parser.add_argument(\n \"--with_p6\",\n action=\"store_true\",\n help=\"Whether your model uses p6 in FPN/PAN.\",\n )\n return parser\n\n\nif __name__ == '__main__':\n args = make_parser().parse_args()\n\n input_shape = tuple(map(int, args.input_shape.split(',')))\n origin_img = cv2.imread(args.image_path)\n mean = (0.485, 0.456, 0.406)\n std = (0.229, 0.224, 0.225)\n img, ratio = preprocess(origin_img, input_shape, mean, std)\n\n session = onnxruntime.InferenceSession(args.model)\n\n ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}\n output = session.run(None, ort_inputs)\n predictions = demo_postprocess(output[0], input_shape, p6=args.with_p6)[0]\n\n boxes = predictions[:, :4]\n scores = predictions[:, 4:5] * predictions[:, 5:]\n\n boxes_xyxy = np.ones_like(boxes)\n boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.\n boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.\n boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.\n boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.\n boxes_xyxy /= ratio\n dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.65, score_thr=0.1)\n \n if dets is not None:\n final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]\n origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,\n conf=args.score_thr, class_names=COCO_CLASSES)\n\n mkdir(args.output_dir)\n output_path = os.path.join(args.output_dir, args.image_path.split(\"/\")[-1])\n cv2.imwrite(output_path, origin_img)\n"
},
{
"path": "detector/YOLOX/demo/OpenVINO/README.md",
"content": "## YOLOX for OpenVINO\n\n* [C++ Demo](./cpp)\n* [Python Demo](./python)"
},
{
"path": "detector/YOLOX/demo/OpenVINO/cpp/CMakeLists.txt",
"content": "cmake_minimum_required(VERSION 3.4.1)\nset(CMAKE_CXX_STANDARD 14)\n\nproject(yolox_openvino_demo)\n\nfind_package(OpenCV REQUIRED)\nfind_package(InferenceEngine REQUIRED)\nfind_package(ngraph REQUIRED)\n\ninclude_directories(\n ${OpenCV_INCLUDE_DIRS}\n ${CMAKE_CURRENT_SOURCE_DIR}\n ${CMAKE_CURRENT_BINARY_DIR}\n)\n\nadd_executable(yolox_openvino yolox_openvino.cpp)\n\ntarget_link_libraries(\n yolox_openvino\n ${InferenceEngine_LIBRARIES}\n ${NGRAPH_LIBRARIES}\n ${OpenCV_LIBS} \n)"
},
{
"path": "detector/YOLOX/demo/OpenVINO/cpp/README.md",
"content": "# YOLOX-OpenVINO in C++\n\nThis toturial includes a C++ demo for OpenVINO, as well as some converted models.\n\n### Download OpenVINO models.\n| Model | Parameters | GFLOPs | Test Size | mAP | Weights |\n|:------| :----: | :----: | :---: | :---: | :---: |\n| [YOLOX-Nano](../../../exps/nano.py) | 0.91M | 1.08 | 416x416 | 25.3 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EeWY57o5wQZFtXYd1KJw6Z8B4vxZru649XxQHYIFgio3Qw?e=ZS81ce)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_nano_openvino.tar.gz) |\n| [YOLOX-Tiny](../../../exps/yolox_tiny.py) | 5.06M | 6.45 | 416x416 |31.7 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ETfvOoCXdVZNinoSpKA_sEYBIQVqfjjF5_M6VvHRnLVcsA?e=STL1pi)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_tiny_openvino.tar.gz) |\n| [YOLOX-S](../../../exps/yolox_s.py) | 9.0M | 26.8 | 640x640 |39.6 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EXUjf3PQnbBLrxNrXPueqaIBzVZOrYQOnJpLK1Fytj5ssA?e=GK0LOM)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s_openvino.tar.gz) |\n| [YOLOX-M](../../../exps/yolox_m.py) | 25.3M | 73.8 | 640x640 |46.4 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EcoT1BPpeRpLvE_4c441zn8BVNCQ2naxDH3rho7WqdlgLQ?e=95VaM9)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_m_openvino.tar.gz) |\n| [YOLOX-L](../../../exps/yolox_l.py) | 54.2M | 155.6 | 640x640 |50.0 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EZvmn-YLRuVPh0GAP_w3xHMB2VGvrKqQXyK_Cv5yi_DXUg?e=YRh6Eq)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_l_openvino.tar.gz) |\n| [YOLOX-Darknet53](../../../exps/yolov3.py) | 63.72M | 185.3 | 640x640 |47.3 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EQP8LSroikFHuwX0jFRetmcBOCDWSFmylHxolV7ezUPXGw?e=bEw5iq)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_darknet53_openvino.tar.gz) |\n| [YOLOX-X](../../../exps/yolox_x.py) | 99.1M | 281.9 | 640x640 |51.2 | [Download](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EZFPnLqiD-xIlt7rcZYDjQgB4YXE9wnq1qaSXQwJrsKbdg?e=83nwEz)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_x_openvino.tar.gz) |\n\n## Install OpenVINO Toolkit\n\nPlease visit [Openvino Homepage](https://docs.openvinotoolkit.org/latest/get_started_guides.html) for more details.\n\n## Set up the Environment\n\n### For Linux\n\n**Option1. Set up the environment tempororally. You need to run this command everytime you start a new shell window.**\n\n```shell\nsource /opt/intel/openvino_2021/bin/setupvars.sh\n```\n\n**Option2. Set up the environment permenantly.**\n\n*Step1.* For Linux:\n```shell\nvim ~/.bashrc \n```\n\n*Step2.* Add the following line into your file:\n\n```shell\nsource /opt/intel/openvino_2021/bin/setupvars.sh\n```\n\n*Step3.* Save and exit the file, then run:\n\n```shell\nsource ~/.bashrc\n```\n\n\n## Convert model\n\n1. Export ONNX model\n \n Please refer to the [ONNX toturial](../../ONNXRuntime). **Note that you should set --opset to 10, otherwise your next step will fail.**\n\n2. Convert ONNX to OpenVINO \n\n ``` shell\n cd /openvino_2021/deployment_tools/model_optimizer\n ```\n\n Install requirements for convert tool\n\n ```shell\n sudo ./install_prerequisites/install_prerequisites_onnx.sh\n ```\n\n Then convert model.\n ```shell\n python3 mo.py --input_model --input_shape [--data_type FP16]\n ```\n For example:\n ```shell\n python3 mo.py --input_model yolox.onnx --input_shape (1,3,640,640) --data_type FP16\n ``` \n\n## Build \n\n### Linux\n```shell\nsource /opt/intel/openvino_2021/bin/setupvars.sh\nmkdir build\ncd build\ncmake ..\nmake\n```\n\n## Demo\n\n### c++\n\n```shell\n./yolox_openvino \n```\n"
},
{
"path": "detector/YOLOX/demo/OpenVINO/cpp/yolox_openvino.cpp",
"content": "// Copyright (C) 2018-2021 Intel Corporation\n// SPDX-License-Identifier: Apache-2.0\n//\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nusing namespace InferenceEngine;\n\n/**\n * @brief Define names based depends on Unicode path support\n */\n#define tcout std::cout\n#define file_name_t std::string\n#define imread_t cv::imread\n#define NMS_THRESH 0.65\n#define BBOX_CONF_THRESH 0.3\n\nstatic const int INPUT_W = 416;\nstatic const int INPUT_H = 416;\n\ncv::Mat static_resize(cv::Mat& img) {\n float r = std::min(INPUT_W / (img.cols*1.0), INPUT_H / (img.rows*1.0));\n // r = std::min(r, 1.0f);\n int unpad_w = r * img.cols;\n int unpad_h = r * img.rows;\n cv::Mat re(unpad_h, unpad_w, CV_8UC3);\n cv::resize(img, re, re.size());\n cv::Mat out(INPUT_W, INPUT_H, CV_8UC3, cv::Scalar(114, 114, 114));\n re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows)));\n return out;\n}\n\nvoid blobFromImage(cv::Mat& img, Blob::Ptr& blob){\n cv::cvtColor(img, img, cv::COLOR_BGR2RGB);\n int channels = 3;\n int img_h = img.rows;\n int img_w = img.cols;\n std::vector mean = {0.485, 0.456, 0.406};\n std::vector std = {0.229, 0.224, 0.225};\n InferenceEngine::MemoryBlob::Ptr mblob = InferenceEngine::as(blob);\n if (!mblob) \n {\n THROW_IE_EXCEPTION << \"We expect blob to be inherited from MemoryBlob in matU8ToBlob, \"\n << \"but by fact we were not able to cast inputBlob to MemoryBlob\";\n }\n // locked memory holder should be alive all time while access to its buffer happens\n auto mblobHolder = mblob->wmap();\n\n float *blob_data = mblobHolder.as();\n\n for (size_t c = 0; c < channels; c++) \n {\n for (size_t h = 0; h < img_h; h++) \n {\n for (size_t w = 0; w < img_w; w++) \n {\n blob_data[c * img_w * img_h + h * img_w + w] =\n (((float)img.at(h, w)[c]) / 255.0f - mean[c]) / std[c];\n }\n }\n }\n}\n\n\nstruct Object\n{\n cv::Rect_ rect;\n int label;\n float prob;\n};\n\nstruct GridAndStride\n{\n int grid0;\n int grid1;\n int stride;\n};\n\nstatic void generate_grids_and_stride(const int target_size, std::vector& strides, std::vector& grid_strides)\n{\n for (auto stride : strides)\n {\n int num_grid = target_size / stride;\n for (int g1 = 0; g1 < num_grid; g1++)\n {\n for (int g0 = 0; g0 < num_grid; g0++)\n {\n grid_strides.push_back((GridAndStride){g0, g1, stride});\n }\n }\n }\n}\n\n\nstatic void generate_yolox_proposals(std::vector grid_strides, const float* feat_ptr, float prob_threshold, std::vector& objects)\n{\n const int num_class = 80; // COCO has 80 classes. Modify this value on your own dataset.\n\n const int num_anchors = grid_strides.size();\n\n for (int anchor_idx = 0; anchor_idx < num_anchors; anchor_idx++)\n {\n const int grid0 = grid_strides[anchor_idx].grid0;\n const int grid1 = grid_strides[anchor_idx].grid1;\n const int stride = grid_strides[anchor_idx].stride;\n\n\tconst int basic_pos = anchor_idx * 85;\n\n // yolox/models/yolo_head.py decode logic\n // outputs[..., :2] = (outputs[..., :2] + grids) * strides\n // outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides\n float x_center = (feat_ptr[basic_pos + 0] + grid0) * stride;\n float y_center = (feat_ptr[basic_pos + 1] + grid1) * stride;\n float w = exp(feat_ptr[basic_pos + 2]) * stride;\n float h = exp(feat_ptr[basic_pos + 3]) * stride;\n float x0 = x_center - w * 0.5f;\n float y0 = y_center - h * 0.5f;\n\n float box_objectness = feat_ptr[basic_pos + 4];\n for (int class_idx = 0; class_idx < num_class; class_idx++)\n {\n float box_cls_score = feat_ptr[basic_pos + 5 + class_idx];\n float box_prob = box_objectness * box_cls_score;\n if (box_prob > prob_threshold)\n {\n Object obj;\n obj.rect.x = x0;\n obj.rect.y = y0;\n obj.rect.width = w;\n obj.rect.height = h;\n obj.label = class_idx;\n obj.prob = box_prob;\n\n objects.push_back(obj);\n }\n\n } // class loop\n\n } // point anchor loop\n}\n\nstatic inline float intersection_area(const Object& a, const Object& b)\n{\n cv::Rect_ inter = a.rect & b.rect;\n return inter.area();\n}\n\nstatic void qsort_descent_inplace(std::vector& faceobjects, int left, int right)\n{\n int i = left;\n int j = right;\n float p = faceobjects[(left + right) / 2].prob;\n\n while (i <= j)\n {\n while (faceobjects[i].prob > p)\n i++;\n\n while (faceobjects[j].prob < p)\n j--;\n\n if (i <= j)\n {\n // swap\n std::swap(faceobjects[i], faceobjects[j]);\n\n i++;\n j--;\n }\n }\n\n #pragma omp parallel sections\n {\n #pragma omp section\n {\n if (left < j) qsort_descent_inplace(faceobjects, left, j);\n }\n #pragma omp section\n {\n if (i < right) qsort_descent_inplace(faceobjects, i, right);\n }\n }\n}\n\n\nstatic void qsort_descent_inplace(std::vector& objects)\n{\n if (objects.empty())\n return;\n\n qsort_descent_inplace(objects, 0, objects.size() - 1);\n}\n\nstatic void nms_sorted_bboxes(const std::vector& faceobjects, std::vector& picked, float nms_threshold)\n{\n picked.clear();\n\n const int n = faceobjects.size();\n\n std::vector areas(n);\n for (int i = 0; i < n; i++)\n {\n areas[i] = faceobjects[i].rect.area();\n }\n\n for (int i = 0; i < n; i++)\n {\n const Object& a = faceobjects[i];\n\n int keep = 1;\n for (int j = 0; j < (int)picked.size(); j++)\n {\n const Object& b = faceobjects[picked[j]];\n\n // intersection over union\n float inter_area = intersection_area(a, b);\n float union_area = areas[i] + areas[picked[j]] - inter_area;\n // float IoU = inter_area / union_area\n if (inter_area / union_area > nms_threshold)\n keep = 0;\n }\n\n if (keep)\n picked.push_back(i);\n }\n}\n\n\nstatic void decode_outputs(const float* prob, std::vector& objects, float scale, const int img_w, const int img_h) {\n std::vector proposals;\n std::vector strides = {8, 16, 32};\n std::vector grid_strides;\n\n generate_grids_and_stride(INPUT_W, strides, grid_strides);\n generate_yolox_proposals(grid_strides, prob, BBOX_CONF_THRESH, proposals);\n qsort_descent_inplace(proposals);\n\n std::vector picked;\n nms_sorted_bboxes(proposals, picked, NMS_THRESH);\n int count = picked.size();\n objects.resize(count);\n\n for (int i = 0; i < count; i++)\n {\n objects[i] = proposals[picked[i]];\n\n // adjust offset to original unpadded\n float x0 = (objects[i].rect.x) / scale;\n float y0 = (objects[i].rect.y) / scale;\n float x1 = (objects[i].rect.x + objects[i].rect.width) / scale;\n float y1 = (objects[i].rect.y + objects[i].rect.height) / scale;\n\n // clip\n x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f);\n y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f);\n x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f);\n y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f);\n\n objects[i].rect.x = x0;\n objects[i].rect.y = y0;\n objects[i].rect.width = x1 - x0;\n objects[i].rect.height = y1 - y0;\n }\n}\n\nconst float color_list[80][3] =\n{\n {0.000, 0.447, 0.741},\n {0.850, 0.325, 0.098},\n {0.929, 0.694, 0.125},\n {0.494, 0.184, 0.556},\n {0.466, 0.674, 0.188},\n {0.301, 0.745, 0.933},\n {0.635, 0.078, 0.184},\n {0.300, 0.300, 0.300},\n {0.600, 0.600, 0.600},\n {1.000, 0.000, 0.000},\n {1.000, 0.500, 0.000},\n {0.749, 0.749, 0.000},\n {0.000, 1.000, 0.000},\n {0.000, 0.000, 1.000},\n {0.667, 0.000, 1.000},\n {0.333, 0.333, 0.000},\n {0.333, 0.667, 0.000},\n {0.333, 1.000, 0.000},\n {0.667, 0.333, 0.000},\n {0.667, 0.667, 0.000},\n {0.667, 1.000, 0.000},\n {1.000, 0.333, 0.000},\n {1.000, 0.667, 0.000},\n {1.000, 1.000, 0.000},\n {0.000, 0.333, 0.500},\n {0.000, 0.667, 0.500},\n {0.000, 1.000, 0.500},\n {0.333, 0.000, 0.500},\n {0.333, 0.333, 0.500},\n {0.333, 0.667, 0.500},\n {0.333, 1.000, 0.500},\n {0.667, 0.000, 0.500},\n {0.667, 0.333, 0.500},\n {0.667, 0.667, 0.500},\n {0.667, 1.000, 0.500},\n {1.000, 0.000, 0.500},\n {1.000, 0.333, 0.500},\n {1.000, 0.667, 0.500},\n {1.000, 1.000, 0.500},\n {0.000, 0.333, 1.000},\n {0.000, 0.667, 1.000},\n {0.000, 1.000, 1.000},\n {0.333, 0.000, 1.000},\n {0.333, 0.333, 1.000},\n {0.333, 0.667, 1.000},\n {0.333, 1.000, 1.000},\n {0.667, 0.000, 1.000},\n {0.667, 0.333, 1.000},\n {0.667, 0.667, 1.000},\n {0.667, 1.000, 1.000},\n {1.000, 0.000, 1.000},\n {1.000, 0.333, 1.000},\n {1.000, 0.667, 1.000},\n {0.333, 0.000, 0.000},\n {0.500, 0.000, 0.000},\n {0.667, 0.000, 0.000},\n {0.833, 0.000, 0.000},\n {1.000, 0.000, 0.000},\n {0.000, 0.167, 0.000},\n {0.000, 0.333, 0.000},\n {0.000, 0.500, 0.000},\n {0.000, 0.667, 0.000},\n {0.000, 0.833, 0.000},\n {0.000, 1.000, 0.000},\n {0.000, 0.000, 0.167},\n {0.000, 0.000, 0.333},\n {0.000, 0.000, 0.500},\n {0.000, 0.000, 0.667},\n {0.000, 0.000, 0.833},\n {0.000, 0.000, 1.000},\n {0.000, 0.000, 0.000},\n {0.143, 0.143, 0.143},\n {0.286, 0.286, 0.286},\n {0.429, 0.429, 0.429},\n {0.571, 0.571, 0.571},\n {0.714, 0.714, 0.714},\n {0.857, 0.857, 0.857},\n {0.000, 0.447, 0.741},\n {0.314, 0.717, 0.741},\n {0.50, 0.5, 0}\n};\n\nstatic void draw_objects(const cv::Mat& bgr, const std::vector& objects)\n{\n static const char* class_names[] = {\n \"person\", \"bicycle\", \"car\", \"motorcycle\", \"airplane\", \"bus\", \"train\", \"truck\", \"boat\", \"traffic light\",\n \"fire hydrant\", \"stop sign\", \"parking meter\", \"bench\", \"bird\", \"cat\", \"dog\", \"horse\", \"sheep\", \"cow\",\n \"elephant\", \"bear\", \"zebra\", \"giraffe\", \"backpack\", \"umbrella\", \"handbag\", \"tie\", \"suitcase\", \"frisbee\",\n \"skis\", \"snowboard\", \"sports ball\", \"kite\", \"baseball bat\", \"baseball glove\", \"skateboard\", \"surfboard\",\n \"tennis racket\", \"bottle\", \"wine glass\", \"cup\", \"fork\", \"knife\", \"spoon\", \"bowl\", \"banana\", \"apple\",\n \"sandwich\", \"orange\", \"broccoli\", \"carrot\", \"hot dog\", \"pizza\", \"donut\", \"cake\", \"chair\", \"couch\",\n \"potted plant\", \"bed\", \"dining table\", \"toilet\", \"tv\", \"laptop\", \"mouse\", \"remote\", \"keyboard\", \"cell phone\",\n \"microwave\", \"oven\", \"toaster\", \"sink\", \"refrigerator\", \"book\", \"clock\", \"vase\", \"scissors\", \"teddy bear\",\n \"hair drier\", \"toothbrush\"\n };\n\n cv::Mat image = bgr.clone();\n\n for (size_t i = 0; i < objects.size(); i++)\n {\n const Object& obj = objects[i];\n\n fprintf(stderr, \"%d = %.5f at %.2f %.2f %.2f x %.2f\\n\", obj.label, obj.prob,\n obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height);\n\n cv::Scalar color = cv::Scalar(color_list[obj.label][0], color_list[obj.label][1], color_list[obj.label][2]);\n float c_mean = cv::mean(color)[0];\n cv::Scalar txt_color;\n if (c_mean > 0.5){\n txt_color = cv::Scalar(0, 0, 0);\n }else{\n txt_color = cv::Scalar(255, 255, 255);\n }\n\n cv::rectangle(image, obj.rect, color * 255, 2);\n\n char text[256];\n sprintf(text, \"%s %.1f%%\", class_names[obj.label], obj.prob * 100);\n\n int baseLine = 0;\n cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);\n\n cv::Scalar txt_bk_color = color * 0.7 * 255;\n\n int x = obj.rect.x;\n int y = obj.rect.y + 1;\n //int y = obj.rect.y - label_size.height - baseLine;\n if (y > image.rows)\n y = image.rows;\n //if (x + label_size.width > image.cols)\n //x = image.cols - label_size.width;\n\n cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),\n txt_bk_color, -1);\n\n cv::putText(image, text, cv::Point(x, y + label_size.height),\n cv::FONT_HERSHEY_SIMPLEX, 0.4, txt_color, 1);\n }\n\n cv::imwrite(\"_demo.jpg\" , image);\n fprintf(stderr, \"save vis file\\n\");\n /* cv::imshow(\"image\", image); */\n /* cv::waitKey(0); */\n}\n\n\nint main(int argc, char* argv[]) {\n try {\n // ------------------------------ Parsing and validation of input arguments\n // ---------------------------------\n if (argc != 4) {\n tcout << \"Usage : \" << argv[0] << \" \" << std::endl;\n return EXIT_FAILURE;\n }\n\n const file_name_t input_model {argv[1]};\n const file_name_t input_image_path {argv[2]};\n const std::string device_name {argv[3]};\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 1. Initialize inference engine core\n // -------------------------------------\n Core ie;\n // -----------------------------------------------------------------------------------------------------\n\n // Step 2. Read a model in OpenVINO Intermediate Representation (.xml and\n // .bin files) or ONNX (.onnx file) format\n CNNNetwork network = ie.ReadNetwork(input_model);\n if (network.getOutputsInfo().size() != 1)\n throw std::logic_error(\"Sample supports topologies with 1 output only\");\n if (network.getInputsInfo().size() != 1)\n throw std::logic_error(\"Sample supports topologies with 1 input only\");\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 3. Configure input & output\n // ---------------------------------------------\n // --------------------------- Prepare input blobs\n // -----------------------------------------------------\n InputInfo::Ptr input_info = network.getInputsInfo().begin()->second;\n std::string input_name = network.getInputsInfo().begin()->first;\n\n /* Mark input as resizable by setting of a resize algorithm.\n * In this case we will be able to set an input blob of any shape to an\n * infer request. Resize and layout conversions are executed automatically\n * during inference */\n //input_info->getPreProcess().setResizeAlgorithm(RESIZE_BILINEAR);\n //input_info->setLayout(Layout::NHWC);\n //input_info->setPrecision(Precision::FP32);\n\n // --------------------------- Prepare output blobs\n // ----------------------------------------------------\n if (network.getOutputsInfo().empty()) {\n std::cerr << \"Network outputs info is empty\" << std::endl;\n return EXIT_FAILURE;\n }\n DataPtr output_info = network.getOutputsInfo().begin()->second;\n std::string output_name = network.getOutputsInfo().begin()->first;\n\n output_info->setPrecision(Precision::FP32);\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 4. Loading a model to the device\n // ------------------------------------------\n ExecutableNetwork executable_network = ie.LoadNetwork(network, device_name);\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 5. Create an infer request\n // -------------------------------------------------\n InferRequest infer_request = executable_network.CreateInferRequest();\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 6. Prepare input\n // --------------------------------------------------------\n /* Read input image to a blob and set it to an infer request without resize\n * and layout conversions. */\n cv::Mat image = imread_t(input_image_path);\n\t cv::Mat pr_img = static_resize(image);\n Blob::Ptr imgBlob = infer_request.GetBlob(input_name); // just wrap Mat data by Blob::Ptr\n\t blobFromImage(pr_img, imgBlob);\n\n // infer_request.SetBlob(input_name, imgBlob); // infer_request accepts input blob of any size\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 7. Do inference\n // --------------------------------------------------------\n /* Running the request synchronously */\n infer_request.Infer();\n // -----------------------------------------------------------------------------------------------------\n\n // --------------------------- Step 8. Process output\n // ------------------------------------------------------\n const Blob::Ptr output_blob = infer_request.GetBlob(output_name);\n MemoryBlob::CPtr moutput = as(output_blob);\n if (!moutput) {\n throw std::logic_error(\"We expect output to be inherited from MemoryBlob, \"\n \"but by fact we were not able to cast output to MemoryBlob\");\n }\n // locked memory holder should be alive all time while access to its buffer\n // happens\n auto moutputHolder = moutput->rmap();\n const float* net_pred = moutputHolder.as::value_type*>();\n \n const int image_size = 416;\n\t int img_w = image.cols;\n int img_h = image.rows;\n\t float scale = std::min(INPUT_W / (image.cols*1.0), INPUT_H / (image.rows*1.0));\n std::vector objects;\n\n decode_outputs(net_pred, objects, scale, img_w, img_h);\n draw_objects(image, objects);\n\n // -----------------------------------------------------------------------------------------------------\n } catch (const std::exception& ex) {\n std::cerr << ex.what() << std::endl;\n return EXIT_FAILURE;\n }\n return EXIT_SUCCESS;\n}\n"
},
{
"path": "detector/YOLOX/demo/OpenVINO/python/README.md",
"content": "# YOLOX-OpenVINO in Python\n\nThis toturial includes a Python demo for OpenVINO, as well as some converted models.\n\n### Download OpenVINO models.\n| Model | Parameters | GFLOPs | Test Size | mAP | Weights |\n|:------| :----: | :----: | :---: | :---: | :---: |\n| [YOLOX-Nano](../../../exps/default/nano.py) | 0.91M | 1.08 | 416x416 | 25.3 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EeWY57o5wQZFtXYd1KJw6Z8B4vxZru649XxQHYIFgio3Qw?e=ZS81ce)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_nano_openvino.tar.gz) |\n| [YOLOX-Tiny](../../../exps/default/yolox_tiny.py) | 5.06M | 6.45 | 416x416 |31.7 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ETfvOoCXdVZNinoSpKA_sEYBIQVqfjjF5_M6VvHRnLVcsA?e=STL1pi)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_tiny_openvino.tar.gz) |\n| [YOLOX-S](../../../exps/default/yolox_s.py) | 9.0M | 26.8 | 640x640 |39.6 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EXUjf3PQnbBLrxNrXPueqaIBzVZOrYQOnJpLK1Fytj5ssA?e=GK0LOM)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s_openvino.tar.gz) |\n| [YOLOX-M](../../../exps/default/yolox_m.py) | 25.3M | 73.8 | 640x640 |46.4 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EcoT1BPpeRpLvE_4c441zn8BVNCQ2naxDH3rho7WqdlgLQ?e=95VaM9)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_m_openvino.tar.gz) |\n| [YOLOX-L](../../../exps/default/yolox_l.py) | 54.2M | 155.6 | 640x640 |50.0 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EZvmn-YLRuVPh0GAP_w3xHMB2VGvrKqQXyK_Cv5yi_DXUg?e=YRh6Eq)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_l_openvino.tar.gz) |\n| [YOLOX-Darknet53](../../../exps/default/yolov3.py) | 63.72M | 185.3 | 640x640 |47.3 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EQP8LSroikFHuwX0jFRetmcBOCDWSFmylHxolV7ezUPXGw?e=bEw5iq)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_darknet53_openvino.tar.gz) | \n| [YOLOX-X](../../../exps/default/yolox_x.py) | 99.1M | 281.9 | 640x640 |51.2 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EZFPnLqiD-xIlt7rcZYDjQgB4YXE9wnq1qaSXQwJrsKbdg?e=83nwEz)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_x_openvino.tar.gz) |\n\n## Install OpenVINO Toolkit\n\nPlease visit [Openvino Homepage](https://docs.openvinotoolkit.org/latest/get_started_guides.html) for more details.\n\n## Set up the Environment\n\n### For Linux\n\n**Option1. Set up the environment tempororally. You need to run this command everytime you start a new shell window.**\n\n```shell\nsource /opt/intel/openvino_2021/bin/setupvars.sh\n```\n\n**Option2. Set up the environment permenantly.**\n\n*Step1.* For Linux:\n```shell\nvim ~/.bashrc\n```\n\n*Step2.* Add the following line into your file:\n\n```shell\nsource /opt/intel/openvino_2021/bin/setupvars.sh\n```\n\n*Step3.* Save and exit the file, then run:\n\n```shell\nsource ~/.bashrc\n```\n\n\n## Convert model\n\n1. Export ONNX model\n\n Please refer to the [ONNX toturial](../../ONNXRuntime). **Note that you should set --opset to 10, otherwise your next step will fail.**\n\n2. Convert ONNX to OpenVINO\n\n ``` shell\n cd /openvino_2021/deployment_tools/model_optimizer\n ```\n\n Install requirements for convert tool\n\n ```shell\n sudo ./install_prerequisites/install_prerequisites_onnx.sh\n ```\n\n Then convert model.\n ```shell\n python3 mo.py --input_model --input_shape [--data_type FP16]\n ```\n For example:\n ```shell\n python3 mo.py --input_model yolox.onnx --input_shape [1,3,640,640] --data_type FP16 --output_dir converted_output\n ```\n\n## Demo\n\n### python\n\n```shell\npython openvino_inference.py -m -i \n```\nor\n```shell\npython openvino_inference.py -m -i -o -s -d \n```\n\n"
},
{
"path": "detector/YOLOX/demo/OpenVINO/python/openvino_inference.py",
"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n# Copyright (C) 2018-2021 Intel Corporation\n# SPDX-License-Identifier: Apache-2.0\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport logging as log\nimport os\nimport sys\n\nimport cv2\nimport numpy as np\n\nfrom openvino.inference_engine import IECore\n\nfrom yolox.data.data_augment import preproc as preprocess\nfrom yolox.data.datasets import COCO_CLASSES\nfrom yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis\n\n\ndef parse_args() -> argparse.Namespace:\n \"\"\"Parse and return command line arguments\"\"\"\n parser = argparse.ArgumentParser(add_help=False)\n args = parser.add_argument_group('Options')\n args.add_argument(\n '-h',\n '--help',\n action='help',\n help='Show this help message and exit.')\n args.add_argument(\n '-m',\n '--model',\n required=True,\n type=str,\n help='Required. Path to an .xml or .onnx file with a trained model.')\n args.add_argument(\n '-i',\n '--input',\n required=True,\n type=str,\n help='Required. Path to an image file.')\n args.add_argument(\n '-o',\n '--output_dir',\n type=str,\n default='demo_output',\n help='Path to your output dir.')\n args.add_argument(\n '-s',\n '--score_thr',\n type=float,\n default=0.3,\n help=\"Score threshould to visualize the result.\")\n args.add_argument(\n '-d',\n '--device',\n default='CPU',\n type=str,\n help='Optional. Specify the target device to infer on; CPU, GPU, \\\n MYRIAD, HDDL or HETERO: is acceptable. The sample will look \\\n for a suitable plugin for device specified. Default value \\\n is CPU.')\n args.add_argument(\n '--labels',\n default=None,\n type=str,\n help='Option:al. Path to a labels mapping file.')\n args.add_argument(\n '-nt',\n '--number_top',\n default=10,\n type=int,\n help='Optional. Number of top results.')\n return parser.parse_args()\n\n\ndef main():\n log.basicConfig(format='[ %(levelname)s ] %(message)s', level=log.INFO, stream=sys.stdout)\n args = parse_args()\n\n # ---------------------------Step 1. Initialize inference engine core--------------------------------------------------\n log.info('Creating Inference Engine')\n ie = IECore()\n\n # ---------------------------Step 2. Read a model in OpenVINO Intermediate Representation or ONNX format---------------\n log.info(f'Reading the network: {args.model}')\n # (.xml and .bin files) or (.onnx file)\n net = ie.read_network(model=args.model)\n\n if len(net.input_info) != 1:\n log.error('Sample supports only single input topologies')\n return -1\n if len(net.outputs) != 1:\n log.error('Sample supports only single output topologies')\n return -1\n\n # ---------------------------Step 3. Configure input & output----------------------------------------------------------\n log.info('Configuring input and output blobs')\n # Get names of input and output blobs\n input_blob = next(iter(net.input_info))\n out_blob = next(iter(net.outputs))\n\n # Set input and output precision manually\n net.input_info[input_blob].precision = 'FP32'\n net.outputs[out_blob].precision = 'FP16'\n\n # Get a number of classes recognized by a model\n num_of_classes = max(net.outputs[out_blob].shape)\n\n # ---------------------------Step 4. Loading model to the device-------------------------------------------------------\n log.info('Loading the model to the plugin')\n exec_net = ie.load_network(network=net, device_name=args.device)\n\n # ---------------------------Step 5. Create infer request--------------------------------------------------------------\n # load_network() method of the IECore class with a specified number of requests (default 1) returns an ExecutableNetwork\n # instance which stores infer requests. So you already created Infer requests in the previous step.\n\n # ---------------------------Step 6. Prepare input---------------------------------------------------------------------\n origin_img = cv2.imread(args.input)\n _, _, h, w = net.input_info[input_blob].input_data.shape\n mean = (0.485, 0.456, 0.406)\n std = (0.229, 0.224, 0.225)\n image, ratio = preprocess(origin_img, (h, w), mean, std)\n\n # ---------------------------Step 7. Do inference----------------------------------------------------------------------\n log.info('Starting inference in synchronous mode')\n res = exec_net.infer(inputs={input_blob: image})\n\n # ---------------------------Step 8. Process output--------------------------------------------------------------------\n res = res[out_blob]\n\n predictions = demo_postprocess(res, (h, w), p6=False)[0]\n\n boxes = predictions[:, :4]\n scores = predictions[:, 4, None] * predictions[:, 5:]\n\n boxes_xyxy = np.ones_like(boxes)\n boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.\n boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.\n boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.\n boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.\n boxes_xyxy /= ratio\n dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.65, score_thr=0.1)\n \n if dets is not None:\n final_boxes = dets[:, :4]\n final_scores, final_cls_inds = dets[:, 4], dets[:, 5]\n origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,\n conf=args.score_thr, class_names=COCO_CLASSES)\n\n mkdir(args.output_dir)\n output_path = os.path.join(args.output_dir, args.image_path.split(\"/\")[-1])\n cv2.imwrite(output_path, origin_img)\n\n\nif __name__ == '__main__':\n sys.exit(main())\n"
},
{
"path": "detector/YOLOX/demo/TensorRT/cpp/CMakeLists.txt",
"content": "cmake_minimum_required(VERSION 2.6)\n\nproject(yolox)\n\nadd_definitions(-std=c++11)\n\noption(CUDA_USE_STATIC_CUDA_RUNTIME OFF)\nset(CMAKE_CXX_STANDARD 11)\nset(CMAKE_BUILD_TYPE Debug)\n\nfind_package(CUDA REQUIRED)\n\ninclude_directories(${PROJECT_SOURCE_DIR}/include)\n# include and link dirs of cuda and tensorrt, you need adapt them if yours are different\n# cuda\ninclude_directories(/data/cuda/cuda-10.2/cuda/include)\nlink_directories(/data/cuda/cuda-10.2/cuda/lib64)\n# cudnn\ninclude_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/include)\nlink_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/lib64)\n# tensorrt\ninclude_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/include)\nlink_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/lib)\n\nset(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED\")\n\nfind_package(OpenCV)\ninclude_directories(${OpenCV_INCLUDE_DIRS})\n\nadd_executable(yolox ${PROJECT_SOURCE_DIR}/yolox.cpp)\ntarget_link_libraries(yolox nvinfer)\ntarget_link_libraries(yolox cudart)\ntarget_link_libraries(yolox ${OpenCV_LIBS})\n\nadd_definitions(-O2 -pthread)\n\n"
},
{
"path": "detector/YOLOX/demo/TensorRT/cpp/README.md",
"content": "# YOLOX-TensorRT in C++\n\nAs YOLOX models is easy to converted to tensorrt using [torch2trt gitrepo](https://github.com/NVIDIA-AI-IOT/torch2trt), \nour C++ demo will not include the model converting or constructing like other tenorrt demos.\n\n\n## Step 1: Prepare serialized engine file\n\nFollow the trt [python demo README](../python/README.md) to convert and save the serialized engine file.\n\nCheck the 'model_trt.engine' file generated from Step 1, which will automatically saved at the current demo dir.\n\n\n## Step 2: build the demo\n\nPlease follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) to install TensorRT.\n\nInstall opencv with ```sudo apt-get install libopencv-dev```.\n\nbuild the demo:\n\n```shell\nmkdir build\ncd build\ncmake ..\nmake\n```\n\nThen run the demo:\n\n```shell\n./yolox ../model_trt.engine -i ../../../../assets/dog.jpg\n```\n\nor\n\n```shell\n./yolox -i \n```\n"
},
{
"path": "detector/YOLOX/demo/TensorRT/cpp/logging.h",
"content": "/*\n * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#ifndef TENSORRT_LOGGING_H\n#define TENSORRT_LOGGING_H\n\n#include \"NvInferRuntimeCommon.h\"\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nusing Severity = nvinfer1::ILogger::Severity;\n\nclass LogStreamConsumerBuffer : public std::stringbuf\n{\npublic:\n LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)\n : mOutput(stream)\n , mPrefix(prefix)\n , mShouldLog(shouldLog)\n {\n }\n\n LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other)\n : mOutput(other.mOutput)\n {\n }\n\n ~LogStreamConsumerBuffer()\n {\n // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence\n // std::streambuf::pptr() gives a pointer to the current position of the output sequence\n // if the pointer to the beginning is not equal to the pointer to the current position,\n // call putOutput() to log the output to the stream\n if (pbase() != pptr())\n {\n putOutput();\n }\n }\n\n // synchronizes the stream buffer and returns 0 on success\n // synchronizing the stream buffer consists of inserting the buffer contents into the stream,\n // resetting the buffer and flushing the stream\n virtual int sync()\n {\n putOutput();\n return 0;\n }\n\n void putOutput()\n {\n if (mShouldLog)\n {\n // prepend timestamp\n std::time_t timestamp = std::time(nullptr);\n tm* tm_local = std::localtime(×tamp);\n std::cout << \"[\";\n std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << \"/\";\n std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << \"/\";\n std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << \"-\";\n std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << \":\";\n std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << \":\";\n std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << \"] \";\n // std::stringbuf::str() gets the string contents of the buffer\n // insert the buffer contents pre-appended by the appropriate prefix into the stream\n mOutput << mPrefix << str();\n // set the buffer to empty\n str(\"\");\n // flush the stream\n mOutput.flush();\n }\n }\n\n void setShouldLog(bool shouldLog)\n {\n mShouldLog = shouldLog;\n }\n\nprivate:\n std::ostream& mOutput;\n std::string mPrefix;\n bool mShouldLog;\n};\n\n//!\n//! \\class LogStreamConsumerBase\n//! \\brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer\n//!\nclass LogStreamConsumerBase\n{\npublic:\n LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)\n : mBuffer(stream, prefix, shouldLog)\n {\n }\n\nprotected:\n LogStreamConsumerBuffer mBuffer;\n};\n\n//!\n//! \\class LogStreamConsumer\n//! \\brief Convenience object used to facilitate use of C++ stream syntax when logging messages.\n//! Order of base classes is LogStreamConsumerBase and then std::ostream.\n//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field\n//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream.\n//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.\n//! Please do not change the order of the parent classes.\n//!\nclass LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream\n{\npublic:\n //! \\brief Creates a LogStreamConsumer which logs messages with level severity.\n //! Reportable severity determines if the messages are severe enough to be logged.\n LogStreamConsumer(Severity reportableSeverity, Severity severity)\n : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity)\n , std::ostream(&mBuffer) // links the stream buffer with the stream\n , mShouldLog(severity <= reportableSeverity)\n , mSeverity(severity)\n {\n }\n\n LogStreamConsumer(LogStreamConsumer&& other)\n : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog)\n , std::ostream(&mBuffer) // links the stream buffer with the stream\n , mShouldLog(other.mShouldLog)\n , mSeverity(other.mSeverity)\n {\n }\n\n void setReportableSeverity(Severity reportableSeverity)\n {\n mShouldLog = mSeverity <= reportableSeverity;\n mBuffer.setShouldLog(mShouldLog);\n }\n\nprivate:\n static std::ostream& severityOstream(Severity severity)\n {\n return severity >= Severity::kINFO ? std::cout : std::cerr;\n }\n\n static std::string severityPrefix(Severity severity)\n {\n switch (severity)\n {\n case Severity::kINTERNAL_ERROR: return \"[F] \";\n case Severity::kERROR: return \"[E] \";\n case Severity::kWARNING: return \"[W] \";\n case Severity::kINFO: return \"[I] \";\n case Severity::kVERBOSE: return \"[V] \";\n default: assert(0); return \"\";\n }\n }\n\n bool mShouldLog;\n Severity mSeverity;\n};\n\n//! \\class Logger\n//!\n//! \\brief Class which manages logging of TensorRT tools and samples\n//!\n//! \\details This class provides a common interface for TensorRT tools and samples to log information to the console,\n//! and supports logging two types of messages:\n//!\n//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)\n//! - Test pass/fail messages\n//!\n//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is\n//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.\n//!\n//! In the future, this class could be extended to support dumping test results to a file in some standard format\n//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).\n//!\n//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger\n//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT\n//! library and messages coming from the sample.\n//!\n//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the\n//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger\n//! object.\n\nclass Logger : public nvinfer1::ILogger\n{\npublic:\n Logger(Severity severity = Severity::kWARNING)\n : mReportableSeverity(severity)\n {\n }\n\n //!\n //! \\enum TestResult\n //! \\brief Represents the state of a given test\n //!\n enum class TestResult\n {\n kRUNNING, //!< The test is running\n kPASSED, //!< The test passed\n kFAILED, //!< The test failed\n kWAIVED //!< The test was waived\n };\n\n //!\n //! \\brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger\n //! \\return The nvinfer1::ILogger associated with this Logger\n //!\n //! TODO Once all samples are updated to use this method to register the logger with TensorRT,\n //! we can eliminate the inheritance of Logger from ILogger\n //!\n nvinfer1::ILogger& getTRTLogger()\n {\n return *this;\n }\n\n //!\n //! \\brief Implementation of the nvinfer1::ILogger::log() virtual method\n //!\n //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the\n //! inheritance from nvinfer1::ILogger\n //!\n void log(Severity severity, const char* msg) override\n {\n LogStreamConsumer(mReportableSeverity, severity) << \"[TRT] \" << std::string(msg) << std::endl;\n }\n\n //!\n //! \\brief Method for controlling the verbosity of logging output\n //!\n //! \\param severity The logger will only emit messages that have severity of this level or higher.\n //!\n void setReportableSeverity(Severity severity)\n {\n mReportableSeverity = severity;\n }\n\n //!\n //! \\brief Opaque handle that holds logging information for a particular test\n //!\n //! This object is an opaque handle to information used by the Logger to print test results.\n //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used\n //! with Logger::reportTest{Start,End}().\n //!\n class TestAtom\n {\n public:\n TestAtom(TestAtom&&) = default;\n\n private:\n friend class Logger;\n\n TestAtom(bool started, const std::string& name, const std::string& cmdline)\n : mStarted(started)\n , mName(name)\n , mCmdline(cmdline)\n {\n }\n\n bool mStarted;\n std::string mName;\n std::string mCmdline;\n };\n\n //!\n //! \\brief Define a test for logging\n //!\n //! \\param[in] name The name of the test. This should be a string starting with\n //! \"TensorRT\" and containing dot-separated strings containing\n //! the characters [A-Za-z0-9_].\n //! For example, \"TensorRT.sample_googlenet\"\n //! \\param[in] cmdline The command line used to reproduce the test\n //\n //! \\return a TestAtom that can be used in Logger::reportTest{Start,End}().\n //!\n static TestAtom defineTest(const std::string& name, const std::string& cmdline)\n {\n return TestAtom(false, name, cmdline);\n }\n\n //!\n //! \\brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments\n //! as input\n //!\n //! \\param[in] name The name of the test\n //! \\param[in] argc The number of command-line arguments\n //! \\param[in] argv The array of command-line arguments (given as C strings)\n //!\n //! \\return a TestAtom that can be used in Logger::reportTest{Start,End}().\n static TestAtom defineTest(const std::string& name, int argc, char const* const* argv)\n {\n auto cmdline = genCmdlineString(argc, argv);\n return defineTest(name, cmdline);\n }\n\n //!\n //! \\brief Report that a test has started.\n //!\n //! \\pre reportTestStart() has not been called yet for the given testAtom\n //!\n //! \\param[in] testAtom The handle to the test that has started\n //!\n static void reportTestStart(TestAtom& testAtom)\n {\n reportTestResult(testAtom, TestResult::kRUNNING);\n assert(!testAtom.mStarted);\n testAtom.mStarted = true;\n }\n\n //!\n //! \\brief Report that a test has ended.\n //!\n //! \\pre reportTestStart() has been called for the given testAtom\n //!\n //! \\param[in] testAtom The handle to the test that has ended\n //! \\param[in] result The result of the test. Should be one of TestResult::kPASSED,\n //! TestResult::kFAILED, TestResult::kWAIVED\n //!\n static void reportTestEnd(const TestAtom& testAtom, TestResult result)\n {\n assert(result != TestResult::kRUNNING);\n assert(testAtom.mStarted);\n reportTestResult(testAtom, result);\n }\n\n static int reportPass(const TestAtom& testAtom)\n {\n reportTestEnd(testAtom, TestResult::kPASSED);\n return EXIT_SUCCESS;\n }\n\n static int reportFail(const TestAtom& testAtom)\n {\n reportTestEnd(testAtom, TestResult::kFAILED);\n return EXIT_FAILURE;\n }\n\n static int reportWaive(const TestAtom& testAtom)\n {\n reportTestEnd(testAtom, TestResult::kWAIVED);\n return EXIT_SUCCESS;\n }\n\n static int reportTest(const TestAtom& testAtom, bool pass)\n {\n return pass ? reportPass(testAtom) : reportFail(testAtom);\n }\n\n Severity getReportableSeverity() const\n {\n return mReportableSeverity;\n }\n\nprivate:\n //!\n //! \\brief returns an appropriate string for prefixing a log message with the given severity\n //!\n static const char* severityPrefix(Severity severity)\n {\n switch (severity)\n {\n case Severity::kINTERNAL_ERROR: return \"[F] \";\n case Severity::kERROR: return \"[E] \";\n case Severity::kWARNING: return \"[W] \";\n case Severity::kINFO: return \"[I] \";\n case Severity::kVERBOSE: return \"[V] \";\n default: assert(0); return \"\";\n }\n }\n\n //!\n //! \\brief returns an appropriate string for prefixing a test result message with the given result\n //!\n static const char* testResultString(TestResult result)\n {\n switch (result)\n {\n case TestResult::kRUNNING: return \"RUNNING\";\n case TestResult::kPASSED: return \"PASSED\";\n case TestResult::kFAILED: return \"FAILED\";\n case TestResult::kWAIVED: return \"WAIVED\";\n default: assert(0); return \"\";\n }\n }\n\n //!\n //! \\brief returns an appropriate output stream (cout or cerr) to use with the given severity\n //!\n static std::ostream& severityOstream(Severity severity)\n {\n return severity >= Severity::kINFO ? std::cout : std::cerr;\n }\n\n //!\n //! \\brief method that implements logging test results\n //!\n static void reportTestResult(const TestAtom& testAtom, TestResult result)\n {\n severityOstream(Severity::kINFO) << \"&&&& \" << testResultString(result) << \" \" << testAtom.mName << \" # \"\n << testAtom.mCmdline << std::endl;\n }\n\n //!\n //! \\brief generate a command line string from the given (argc, argv) values\n //!\n static std::string genCmdlineString(int argc, char const* const* argv)\n {\n std::stringstream ss;\n for (int i = 0; i < argc; i++)\n {\n if (i > 0)\n ss << \" \";\n ss << argv[i];\n }\n return ss.str();\n }\n\n Severity mReportableSeverity;\n};\n\nnamespace\n{\n\n//!\n//! \\brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE\n//!\n//! Example usage:\n//!\n//! LOG_VERBOSE(logger) << \"hello world\" << std::endl;\n//!\ninline LogStreamConsumer LOG_VERBOSE(const Logger& logger)\n{\n return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);\n}\n\n//!\n//! \\brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO\n//!\n//! Example usage:\n//!\n//! LOG_INFO(logger) << \"hello world\" << std::endl;\n//!\ninline LogStreamConsumer LOG_INFO(const Logger& logger)\n{\n return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);\n}\n\n//!\n//! \\brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING\n//!\n//! Example usage:\n//!\n//! LOG_WARN(logger) << \"hello world\" << std::endl;\n//!\ninline LogStreamConsumer LOG_WARN(const Logger& logger)\n{\n return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);\n}\n\n//!\n//! \\brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR\n//!\n//! Example usage:\n//!\n//! LOG_ERROR(logger) << \"hello world\" << std::endl;\n//!\ninline LogStreamConsumer LOG_ERROR(const Logger& logger)\n{\n return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);\n}\n\n//!\n//! \\brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR\n// (\"fatal\" severity)\n//!\n//! Example usage:\n//!\n//! LOG_FATAL(logger) << \"hello world\" << std::endl;\n//!\ninline LogStreamConsumer LOG_FATAL(const Logger& logger)\n{\n return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);\n}\n\n} // anonymous namespace\n\n#endif // TENSORRT_LOGGING_H\n"
},
{
"path": "detector/YOLOX/demo/TensorRT/cpp/yolox.cpp",
"content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \"NvInfer.h\"\n#include \"cuda_runtime_api.h\"\n#include \"logging.h\"\n\n#define CHECK(status) \\\n do\\\n {\\\n auto ret = (status);\\\n if (ret != 0)\\\n {\\\n std::cerr << \"Cuda failure: \" << ret << std::endl;\\\n abort();\\\n }\\\n } while (0)\n\n#define DEVICE 0 // GPU id\n#define NMS_THRESH 0.65\n#define BBOX_CONF_THRESH 0.3\n\nusing namespace nvinfer1;\n\n// stuff we know about the network and the input/output blobs\nstatic const int INPUT_W = 640;\nstatic const int INPUT_H = 640;\nconst char* INPUT_BLOB_NAME = \"input_0\";\nconst char* OUTPUT_BLOB_NAME = \"output_0\";\nstatic Logger gLogger;\n\ncv::Mat static_resize(cv::Mat& img) {\n float r = std::min(INPUT_W / (img.cols*1.0), INPUT_H / (img.rows*1.0));\n // r = std::min(r, 1.0f);\n int unpad_w = r * img.cols;\n int unpad_h = r * img.rows;\n cv::Mat re(unpad_h, unpad_w, CV_8UC3);\n cv::resize(img, re, re.size());\n cv::Mat out(INPUT_W, INPUT_H, CV_8UC3, cv::Scalar(114, 114, 114));\n re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows)));\n return out;\n}\n\nstruct Object\n{\n cv::Rect_ rect;\n int label;\n float prob;\n};\n\nstruct GridAndStride\n{\n int grid0;\n int grid1;\n int stride;\n};\n\nstatic void generate_grids_and_stride(const int target_size, std::vector& strides, std::vector& grid_strides)\n{\n for (auto stride : strides)\n {\n int num_grid = target_size / stride;\n for (int g1 = 0; g1 < num_grid; g1++)\n {\n for (int g0 = 0; g0 < num_grid; g0++)\n {\n grid_strides.push_back((GridAndStride){g0, g1, stride});\n }\n }\n }\n}\n\nstatic inline float intersection_area(const Object& a, const Object& b)\n{\n cv::Rect_ inter = a.rect & b.rect;\n return inter.area();\n}\n\nstatic void qsort_descent_inplace(std::vector& faceobjects, int left, int right)\n{\n int i = left;\n int j = right;\n float p = faceobjects[(left + right) / 2].prob;\n\n while (i <= j)\n {\n while (faceobjects[i].prob > p)\n i++;\n\n while (faceobjects[j].prob < p)\n j--;\n\n if (i <= j)\n {\n // swap\n std::swap(faceobjects[i], faceobjects[j]);\n\n i++;\n j--;\n }\n }\n\n #pragma omp parallel sections\n {\n #pragma omp section\n {\n if (left < j) qsort_descent_inplace(faceobjects, left, j);\n }\n #pragma omp section\n {\n if (i < right) qsort_descent_inplace(faceobjects, i, right);\n }\n }\n}\n\nstatic void qsort_descent_inplace(std::vector& objects)\n{\n if (objects.empty())\n return;\n\n qsort_descent_inplace(objects, 0, objects.size() - 1);\n}\n\nstatic void nms_sorted_bboxes(const std::vector& faceobjects, std::vector& picked, float nms_threshold)\n{\n picked.clear();\n\n const int n = faceobjects.size();\n\n std::vector areas(n);\n for (int i = 0; i < n; i++)\n {\n areas[i] = faceobjects[i].rect.area();\n }\n\n for (int i = 0; i < n; i++)\n {\n const Object& a = faceobjects[i];\n\n int keep = 1;\n for (int j = 0; j < (int)picked.size(); j++)\n {\n const Object& b = faceobjects[picked[j]];\n\n // intersection over union\n float inter_area = intersection_area(a, b);\n float union_area = areas[i] + areas[picked[j]] - inter_area;\n // float IoU = inter_area / union_area\n if (inter_area / union_area > nms_threshold)\n keep = 0;\n }\n\n if (keep)\n picked.push_back(i);\n }\n}\n\n\nstatic void generate_yolox_proposals(std::vector grid_strides, float* feat_blob, float prob_threshold, std::vector& objects)\n{\n const int num_class = 80;\n\n const int num_anchors = grid_strides.size();\n\n for (int anchor_idx = 0; anchor_idx < num_anchors; anchor_idx++)\n {\n const int grid0 = grid_strides[anchor_idx].grid0;\n const int grid1 = grid_strides[anchor_idx].grid1;\n const int stride = grid_strides[anchor_idx].stride;\n\n const int basic_pos = anchor_idx * 85;\n\n // yolox/models/yolo_head.py decode logic\n float x_center = (feat_blob[basic_pos+0] + grid0) * stride;\n float y_center = (feat_blob[basic_pos+1] + grid1) * stride;\n float w = exp(feat_blob[basic_pos+2]) * stride;\n float h = exp(feat_blob[basic_pos+3]) * stride;\n float x0 = x_center - w * 0.5f;\n float y0 = y_center - h * 0.5f;\n\n float box_objectness = feat_blob[basic_pos+4];\n for (int class_idx = 0; class_idx < num_class; class_idx++)\n {\n float box_cls_score = feat_blob[basic_pos + 5 + class_idx];\n float box_prob = box_objectness * box_cls_score;\n if (box_prob > prob_threshold)\n {\n Object obj;\n obj.rect.x = x0;\n obj.rect.y = y0;\n obj.rect.width = w;\n obj.rect.height = h;\n obj.label = class_idx;\n obj.prob = box_prob;\n\n objects.push_back(obj);\n }\n\n } // class loop\n\n } // point anchor loop\n}\n\nfloat* blobFromImage(cv::Mat& img){\n cv::cvtColor(img, img, cv::COLOR_BGR2RGB);\n\n float* blob = new float[img.total()*3];\n int channels = 3;\n int img_h = 640;\n int img_w = 640;\n std::vector mean = {0.485, 0.456, 0.406};\n std::vector std = {0.229, 0.224, 0.225};\n for (size_t c = 0; c < channels; c++) \n {\n for (size_t h = 0; h < img_h; h++) \n {\n for (size_t w = 0; w < img_w; w++) \n {\n blob[c * img_w * img_h + h * img_w + w] =\n (((float)img.at(h, w)[c]) / 255.0f - mean[c]) / std[c];\n }\n }\n }\n return blob;\n}\n\n\nstatic void decode_outputs(float* prob, std::vector& objects, float scale, const int img_w, const int img_h) {\n std::vector proposals;\n std::vector strides = {8, 16, 32};\n std::vector grid_strides;\n generate_grids_and_stride(INPUT_W, strides, grid_strides);\n generate_yolox_proposals(grid_strides, prob, BBOX_CONF_THRESH, proposals);\n std::cout << \"num of boxes before nms: \" << proposals.size() << std::endl;\n\n qsort_descent_inplace(proposals);\n\n std::vector picked;\n nms_sorted_bboxes(proposals, picked, NMS_THRESH);\n\n\n int count = picked.size();\n\n std::cout << \"num of boxes: \" << count << std::endl;\n\n objects.resize(count);\n for (int i = 0; i < count; i++)\n {\n objects[i] = proposals[picked[i]];\n\n // adjust offset to original unpadded\n float x0 = (objects[i].rect.x) / scale;\n float y0 = (objects[i].rect.y) / scale;\n float x1 = (objects[i].rect.x + objects[i].rect.width) / scale;\n float y1 = (objects[i].rect.y + objects[i].rect.height) / scale;\n\n // clip\n x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f);\n y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f);\n x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f);\n y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f);\n\n objects[i].rect.x = x0;\n objects[i].rect.y = y0;\n objects[i].rect.width = x1 - x0;\n objects[i].rect.height = y1 - y0;\n }\n}\n\nconst float color_list[80][3] =\n{\n {0.000, 0.447, 0.741},\n {0.850, 0.325, 0.098},\n {0.929, 0.694, 0.125},\n {0.494, 0.184, 0.556},\n {0.466, 0.674, 0.188},\n {0.301, 0.745, 0.933},\n {0.635, 0.078, 0.184},\n {0.300, 0.300, 0.300},\n {0.600, 0.600, 0.600},\n {1.000, 0.000, 0.000},\n {1.000, 0.500, 0.000},\n {0.749, 0.749, 0.000},\n {0.000, 1.000, 0.000},\n {0.000, 0.000, 1.000},\n {0.667, 0.000, 1.000},\n {0.333, 0.333, 0.000},\n {0.333, 0.667, 0.000},\n {0.333, 1.000, 0.000},\n {0.667, 0.333, 0.000},\n {0.667, 0.667, 0.000},\n {0.667, 1.000, 0.000},\n {1.000, 0.333, 0.000},\n {1.000, 0.667, 0.000},\n {1.000, 1.000, 0.000},\n {0.000, 0.333, 0.500},\n {0.000, 0.667, 0.500},\n {0.000, 1.000, 0.500},\n {0.333, 0.000, 0.500},\n {0.333, 0.333, 0.500},\n {0.333, 0.667, 0.500},\n {0.333, 1.000, 0.500},\n {0.667, 0.000, 0.500},\n {0.667, 0.333, 0.500},\n {0.667, 0.667, 0.500},\n {0.667, 1.000, 0.500},\n {1.000, 0.000, 0.500},\n {1.000, 0.333, 0.500},\n {1.000, 0.667, 0.500},\n {1.000, 1.000, 0.500},\n {0.000, 0.333, 1.000},\n {0.000, 0.667, 1.000},\n {0.000, 1.000, 1.000},\n {0.333, 0.000, 1.000},\n {0.333, 0.333, 1.000},\n {0.333, 0.667, 1.000},\n {0.333, 1.000, 1.000},\n {0.667, 0.000, 1.000},\n {0.667, 0.333, 1.000},\n {0.667, 0.667, 1.000},\n {0.667, 1.000, 1.000},\n {1.000, 0.000, 1.000},\n {1.000, 0.333, 1.000},\n {1.000, 0.667, 1.000},\n {0.333, 0.000, 0.000},\n {0.500, 0.000, 0.000},\n {0.667, 0.000, 0.000},\n {0.833, 0.000, 0.000},\n {1.000, 0.000, 0.000},\n {0.000, 0.167, 0.000},\n {0.000, 0.333, 0.000},\n {0.000, 0.500, 0.000},\n {0.000, 0.667, 0.000},\n {0.000, 0.833, 0.000},\n {0.000, 1.000, 0.000},\n {0.000, 0.000, 0.167},\n {0.000, 0.000, 0.333},\n {0.000, 0.000, 0.500},\n {0.000, 0.000, 0.667},\n {0.000, 0.000, 0.833},\n {0.000, 0.000, 1.000},\n {0.000, 0.000, 0.000},\n {0.143, 0.143, 0.143},\n {0.286, 0.286, 0.286},\n {0.429, 0.429, 0.429},\n {0.571, 0.571, 0.571},\n {0.714, 0.714, 0.714},\n {0.857, 0.857, 0.857},\n {0.000, 0.447, 0.741},\n {0.314, 0.717, 0.741},\n {0.50, 0.5, 0}\n};\n\nstatic void draw_objects(const cv::Mat& bgr, const std::vector& objects, std::string f)\n{\n static const char* class_names[] = {\n \"person\", \"bicycle\", \"car\", \"motorcycle\", \"airplane\", \"bus\", \"train\", \"truck\", \"boat\", \"traffic light\",\n \"fire hydrant\", \"stop sign\", \"parking meter\", \"bench\", \"bird\", \"cat\", \"dog\", \"horse\", \"sheep\", \"cow\",\n \"elephant\", \"bear\", \"zebra\", \"giraffe\", \"backpack\", \"umbrella\", \"handbag\", \"tie\", \"suitcase\", \"frisbee\",\n \"skis\", \"snowboard\", \"sports ball\", \"kite\", \"baseball bat\", \"baseball glove\", \"skateboard\", \"surfboard\",\n \"tennis racket\", \"bottle\", \"wine glass\", \"cup\", \"fork\", \"knife\", \"spoon\", \"bowl\", \"banana\", \"apple\",\n \"sandwich\", \"orange\", \"broccoli\", \"carrot\", \"hot dog\", \"pizza\", \"donut\", \"cake\", \"chair\", \"couch\",\n \"potted plant\", \"bed\", \"dining table\", \"toilet\", \"tv\", \"laptop\", \"mouse\", \"remote\", \"keyboard\", \"cell phone\",\n \"microwave\", \"oven\", \"toaster\", \"sink\", \"refrigerator\", \"book\", \"clock\", \"vase\", \"scissors\", \"teddy bear\",\n \"hair drier\", \"toothbrush\"\n };\n\n cv::Mat image = bgr.clone();\n\n for (size_t i = 0; i < objects.size(); i++)\n {\n const Object& obj = objects[i];\n\n fprintf(stderr, \"%d = %.5f at %.2f %.2f %.2f x %.2f\\n\", obj.label, obj.prob,\n obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height);\n\n cv::Scalar color = cv::Scalar(color_list[obj.label][0], color_list[obj.label][1], color_list[obj.label][2]);\n float c_mean = cv::mean(color)[0];\n cv::Scalar txt_color;\n if (c_mean > 0.5){\n txt_color = cv::Scalar(0, 0, 0);\n }else{\n txt_color = cv::Scalar(255, 255, 255);\n }\n\n cv::rectangle(image, obj.rect, color * 255, 2);\n\n char text[256];\n sprintf(text, \"%s %.1f%%\", class_names[obj.label], obj.prob * 100);\n\n int baseLine = 0;\n cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);\n\n cv::Scalar txt_bk_color = color * 0.7 * 255;\n\n int x = obj.rect.x;\n int y = obj.rect.y + 1;\n //int y = obj.rect.y - label_size.height - baseLine;\n if (y > image.rows)\n y = image.rows;\n //if (x + label_size.width > image.cols)\n //x = image.cols - label_size.width;\n\n cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),\n txt_bk_color, -1);\n\n cv::putText(image, text, cv::Point(x, y + label_size.height),\n cv::FONT_HERSHEY_SIMPLEX, 0.4, txt_color, 1);\n }\n\n cv::imwrite(\"det_res.jpg\", image);\n fprintf(stderr, \"save vis file\\n\");\n /* cv::imshow(\"image\", image); */\n /* cv::waitKey(0); */\n}\n\n\nvoid doInference(IExecutionContext& context, float* input, float* output, const int output_size, cv::Size input_shape) {\n const ICudaEngine& engine = context.getEngine();\n\n // Pointers to input and output device buffers to pass to engine.\n // Engine requires exactly IEngine::getNbBindings() number of buffers.\n assert(engine.getNbBindings() == 2);\n void* buffers[2];\n\n // In order to bind the buffers, we need to know the names of the input and output tensors.\n // Note that indices are guaranteed to be less than IEngine::getNbBindings()\n const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);\n\n assert(engine.getBindingDataType(inputIndex) == nvinfer1::DataType::kFLOAT);\n const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);\n assert(engine.getBindingDataType(outputIndex) == nvinfer1::DataType::kFLOAT);\n int mBatchSize = engine.getMaxBatchSize();\n\n // Create GPU buffers on device\n CHECK(cudaMalloc(&buffers[inputIndex], 3 * input_shape.height * input_shape.width * sizeof(float)));\n CHECK(cudaMalloc(&buffers[outputIndex], output_size*sizeof(float)));\n\n // Create stream\n cudaStream_t stream;\n CHECK(cudaStreamCreate(&stream));\n\n // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host\n CHECK(cudaMemcpyAsync(buffers[inputIndex], input, 3 * input_shape.height * input_shape.width * sizeof(float), cudaMemcpyHostToDevice, stream));\n context.enqueue(1, buffers, stream, nullptr);\n CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size * sizeof(float), cudaMemcpyDeviceToHost, stream));\n cudaStreamSynchronize(stream);\n\n // Release stream and buffers\n cudaStreamDestroy(stream);\n CHECK(cudaFree(buffers[inputIndex]));\n CHECK(cudaFree(buffers[outputIndex]));\n}\n\nint main(int argc, char** argv) {\n cudaSetDevice(DEVICE);\n // create a model using the API directly and serialize it to a stream\n char *trtModelStream{nullptr};\n size_t size{0};\n\n if (argc == 4 && std::string(argv[2]) == \"-i\") {\n const std::string engine_file_path {argv[1]};\n std::ifstream file(engine_file_path, std::ios::binary);\n if (file.good()) {\n file.seekg(0, file.end);\n size = file.tellg();\n file.seekg(0, file.beg);\n trtModelStream = new char[size];\n assert(trtModelStream);\n file.read(trtModelStream, size);\n file.close();\n }\n } else {\n std::cerr << \"arguments not right!\" << std::endl;\n std::cerr << \"run 'python3 yolox/deploy/trt.py -n yolox-{tiny, s, m, l, x}' to serialize model first!\" << std::endl;\n std::cerr << \"Then use the following command:\" << std::endl;\n std::cerr << \"./yolox ../model_trt.engine -i ../../../assets/dog.jpg // deserialize file and run inference\" << std::endl;\n return -1;\n }\n const std::string input_image_path {argv[3]};\n\n //std::vector file_names;\n //if (read_files_in_dir(argv[2], file_names) < 0) {\n //std::cout << \"read_files_in_dir failed.\" << std::endl;\n //return -1;\n //}\n\n IRuntime* runtime = createInferRuntime(gLogger);\n assert(runtime != nullptr);\n ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);\n assert(engine != nullptr); \n IExecutionContext* context = engine->createExecutionContext();\n assert(context != nullptr);\n delete[] trtModelStream;\n auto out_dims = engine->getBindingDimensions(1);\n auto output_size = 1;\n for(int j=0;j(end - start).count() << \"ms\" << std::endl;\n\n std::vector objects;\n decode_outputs(prob, objects, scale, img_w, img_h);\n draw_objects(img, objects, input_image_path);\n\n // destroy the engine\n context->destroy();\n engine->destroy();\n runtime->destroy();\n return 0;\n}\n"
},
{
"path": "detector/YOLOX/demo/TensorRT/python/README.md",
"content": "# YOLOX-TensorRT in Python\n\nThis toturial includes a Python demo for TensorRT.\n\n## Install TensorRT Toolkit\n\nPlease follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) and [torch2trt gitrepo](https://github.com/NVIDIA-AI-IOT/torch2trt) to install TensorRT and torch2trt.\n\n## Convert model\n\nYOLOX models can be easily conveted to TensorRT models using torch2trt\n\n If you want to convert our model, use the flag -n to specify a model name:\n ```shell\n python tools/trt.py -n -c \n ```\n For example:\n ```shell\n python tools/trt.py -n yolox-s -c your_ckpt.pth.tar\n ```\n can be: yolox-nano, yolox-tiny. yolox-s, yolox-m, yolox-l, yolox-x.\n\n If you want to convert your customized model, use the flag -f to specify you exp file:\n ```shell\n python tools/trt.py -f -c \n ```\n For example:\n ```shell\n python tools/trt.py -f /path/to/your/yolox/exps/yolox_s.py -c your_ckpt.pth.tar\n ```\n *yolox_s.py* can be any exp file modified by you.\n\nThe converted model and the serialized engine file (for C++ demo) will be saved on your experiment output dir. \n\n## Demo\n\nThe TensorRT python demo is merged on our pytorch demo file, so you can run the pytorch demo command with ```--trt```.\n\n```shell\npython tools/demo.py image -n yolox-s --trt --save_result\n```\nor\n```shell\npython tools/demo.py image -f exps/default/yolox_s.py --trt --save_result\n```\n\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/README.md",
"content": "# YOLOX-Android-ncnn\n\nAndoird app of YOLOX object detection base on [ncnn](https://github.com/Tencent/ncnn)\n\n\n## Tutorial\n\n### Step1\n\nDownload ncnn-android-vulkan.zip from [releases of ncnn](https://github.com/Tencent/ncnn/releases). This repo us\n[20210525 release](https://github.com/Tencent/ncnn/releases/download/20210525/ncnn-20210525-android-vulkan.zip) for building.\n\n### Step2\n\nAfter downloading, please extract your zip file. Then, there are two ways to finish this step:\n* put your extracted directory into app/src/main/jni\n* change the ncnn_DIR path in app/src/main/jni/CMakeLists.txt to your extracted directory.\n\n### Step3\nDownload example param and bin file from [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ESXBH_GSSmFMszWJ6YG2VkQB5cWDfqVWXgk0D996jH0rpQ?e=qzEqUh) or [github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s_ncnn.tar.gz). Unzip the file to app/src/main/assets.\n\n### Step4\nOpen this project with Android Studio, build it and enjoy!\n\n## Reference\n\n* [ncnn-android-yolov5](https://github.com/nihui/ncnn-android-yolov5)\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/build.gradle",
"content": "apply plugin: 'com.android.application'\n\nandroid {\n compileSdkVersion 24\n buildToolsVersion \"29.0.2\"\n\n defaultConfig {\n applicationId \"com.megvii.yoloXncnn\"\n archivesBaseName = \"$applicationId\"\n\n ndk {\n moduleName \"ncnn\"\n abiFilters \"armeabi-v7a\", \"arm64-v8a\"\n }\n minSdkVersion 24\n }\n\n externalNativeBuild {\n cmake {\n version \"3.10.2\"\n path file('src/main/jni/CMakeLists.txt')\n }\n }\n}\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/AndroidManifest.xml",
"content": "\n\n \n \n \n \n \n \n \n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/assets/yolox.param",
"content": "7767517\n220 250\nInput images 0 1 images\nYoloV5Focus focus 1 1 images 503\nConvolution Conv_41 1 1 503 877 0=32 1=3 4=1 5=1 6=3456\nSwish Mul_43 1 1 877 507\nConvolution Conv_44 1 1 507 880 0=64 1=3 3=2 4=1 5=1 6=18432\nSwish Mul_46 1 1 880 511\nSplit splitncnn_0 1 2 511 511_splitncnn_0 511_splitncnn_1\nConvolution Conv_47 1 1 511_splitncnn_1 883 0=32 1=1 5=1 6=2048\nSwish Mul_49 1 1 883 515\nSplit splitncnn_1 1 2 515 515_splitncnn_0 515_splitncnn_1\nConvolution Conv_50 1 1 511_splitncnn_0 886 0=32 1=1 5=1 6=2048\nSwish Mul_52 1 1 886 519\nConvolution Conv_53 1 1 515_splitncnn_1 889 0=32 1=1 5=1 6=1024\nSwish Mul_55 1 1 889 523\nConvolution Conv_56 1 1 523 892 0=32 1=3 4=1 5=1 6=9216\nSwish Mul_58 1 1 892 527\nBinaryOp Add_59 2 1 527 515_splitncnn_0 528\nConcat Concat_60 2 1 528 519 529\nConvolution Conv_61 1 1 529 895 0=64 1=1 5=1 6=4096\nSwish Mul_63 1 1 895 533\nConvolution Conv_64 1 1 533 898 0=128 1=3 3=2 4=1 5=1 6=73728\nSwish Mul_66 1 1 898 537\nSplit splitncnn_2 1 2 537 537_splitncnn_0 537_splitncnn_1\nConvolution Conv_67 1 1 537_splitncnn_1 901 0=64 1=1 5=1 6=8192\nSwish Mul_69 1 1 901 541\nSplit splitncnn_3 1 2 541 541_splitncnn_0 541_splitncnn_1\nConvolution Conv_70 1 1 537_splitncnn_0 904 0=64 1=1 5=1 6=8192\nSwish Mul_72 1 1 904 545\nConvolution Conv_73 1 1 541_splitncnn_1 907 0=64 1=1 5=1 6=4096\nSwish Mul_75 1 1 907 549\nConvolution Conv_76 1 1 549 910 0=64 1=3 4=1 5=1 6=36864\nSwish Mul_78 1 1 910 553\nBinaryOp Add_79 2 1 553 541_splitncnn_0 554\nSplit splitncnn_4 1 2 554 554_splitncnn_0 554_splitncnn_1\nConvolution Conv_80 1 1 554_splitncnn_1 913 0=64 1=1 5=1 6=4096\nSwish Mul_82 1 1 913 558\nConvolution Conv_83 1 1 558 916 0=64 1=3 4=1 5=1 6=36864\nSwish Mul_85 1 1 916 562\nBinaryOp Add_86 2 1 562 554_splitncnn_0 563\nSplit splitncnn_5 1 2 563 563_splitncnn_0 563_splitncnn_1\nConvolution Conv_87 1 1 563_splitncnn_1 919 0=64 1=1 5=1 6=4096\nSwish Mul_89 1 1 919 567\nConvolution Conv_90 1 1 567 922 0=64 1=3 4=1 5=1 6=36864\nSwish Mul_92 1 1 922 571\nBinaryOp Add_93 2 1 571 563_splitncnn_0 572\nConcat Concat_94 2 1 572 545 573\nConvolution Conv_95 1 1 573 925 0=128 1=1 5=1 6=16384\nSwish Mul_97 1 1 925 577\nSplit splitncnn_6 1 2 577 577_splitncnn_0 577_splitncnn_1\nConvolution Conv_98 1 1 577_splitncnn_1 928 0=256 1=3 3=2 4=1 5=1 6=294912\nSwish Mul_100 1 1 928 581\nSplit splitncnn_7 1 2 581 581_splitncnn_0 581_splitncnn_1\nConvolution Conv_101 1 1 581_splitncnn_1 931 0=128 1=1 5=1 6=32768\nSwish Mul_103 1 1 931 585\nSplit splitncnn_8 1 2 585 585_splitncnn_0 585_splitncnn_1\nConvolution Conv_104 1 1 581_splitncnn_0 934 0=128 1=1 5=1 6=32768\nSwish Mul_106 1 1 934 589\nConvolution Conv_107 1 1 585_splitncnn_1 937 0=128 1=1 5=1 6=16384\nSwish Mul_109 1 1 937 593\nConvolution Conv_110 1 1 593 940 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_112 1 1 940 597\nBinaryOp Add_113 2 1 597 585_splitncnn_0 598\nSplit splitncnn_9 1 2 598 598_splitncnn_0 598_splitncnn_1\nConvolution Conv_114 1 1 598_splitncnn_1 943 0=128 1=1 5=1 6=16384\nSwish Mul_116 1 1 943 602\nConvolution Conv_117 1 1 602 946 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_119 1 1 946 606\nBinaryOp Add_120 2 1 606 598_splitncnn_0 607\nSplit splitncnn_10 1 2 607 607_splitncnn_0 607_splitncnn_1\nConvolution Conv_121 1 1 607_splitncnn_1 949 0=128 1=1 5=1 6=16384\nSwish Mul_123 1 1 949 611\nConvolution Conv_124 1 1 611 952 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_126 1 1 952 615\nBinaryOp Add_127 2 1 615 607_splitncnn_0 616\nConcat Concat_128 2 1 616 589 617\nConvolution Conv_129 1 1 617 955 0=256 1=1 5=1 6=65536\nSwish Mul_131 1 1 955 621\nSplit splitncnn_11 1 2 621 621_splitncnn_0 621_splitncnn_1\nConvolution Conv_132 1 1 621_splitncnn_1 958 0=512 1=3 3=2 4=1 5=1 6=1179648\nSwish Mul_134 1 1 958 625\nConvolution Conv_135 1 1 625 961 0=256 1=1 5=1 6=131072\nSwish Mul_137 1 1 961 629\nSplit splitncnn_12 1 4 629 629_splitncnn_0 629_splitncnn_1 629_splitncnn_2 629_splitncnn_3\nPooling MaxPool_138 1 1 629_splitncnn_3 630 1=5 3=2 5=1\nPooling MaxPool_139 1 1 629_splitncnn_2 631 1=9 3=4 5=1\nPooling MaxPool_140 1 1 629_splitncnn_1 632 1=13 3=6 5=1\nConcat Concat_141 4 1 629_splitncnn_0 630 631 632 633\nConvolution Conv_142 1 1 633 964 0=512 1=1 5=1 6=524288\nSwish Mul_144 1 1 964 637\nSplit splitncnn_13 1 2 637 637_splitncnn_0 637_splitncnn_1\nConvolution Conv_145 1 1 637_splitncnn_1 967 0=256 1=1 5=1 6=131072\nSwish Mul_147 1 1 967 641\nConvolution Conv_148 1 1 637_splitncnn_0 970 0=256 1=1 5=1 6=131072\nSwish Mul_150 1 1 970 645\nConvolution Conv_151 1 1 641 973 0=256 1=1 5=1 6=65536\nSwish Mul_153 1 1 973 649\nConvolution Conv_154 1 1 649 976 0=256 1=3 4=1 5=1 6=589824\nSwish Mul_156 1 1 976 653\nConcat Concat_157 2 1 653 645 654\nConvolution Conv_158 1 1 654 979 0=512 1=1 5=1 6=262144\nSwish Mul_160 1 1 979 658\nConvolution Conv_161 1 1 658 982 0=256 1=1 5=1 6=131072\nSwish Mul_163 1 1 982 662\nSplit splitncnn_14 1 2 662 662_splitncnn_0 662_splitncnn_1\nInterp Resize_165 1 1 662_splitncnn_1 667 0=1 1=2.000000e+00 2=2.000000e+00\nConcat Concat_166 2 1 667 621_splitncnn_0 668\nSplit splitncnn_15 1 2 668 668_splitncnn_0 668_splitncnn_1\nConvolution Conv_167 1 1 668_splitncnn_1 985 0=128 1=1 5=1 6=65536\nSwish Mul_169 1 1 985 672\nConvolution Conv_170 1 1 668_splitncnn_0 988 0=128 1=1 5=1 6=65536\nSwish Mul_172 1 1 988 676\nConvolution Conv_173 1 1 672 991 0=128 1=1 5=1 6=16384\nSwish Mul_175 1 1 991 680\nConvolution Conv_176 1 1 680 994 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_178 1 1 994 684\nConcat Concat_179 2 1 684 676 685\nConvolution Conv_180 1 1 685 997 0=256 1=1 5=1 6=65536\nSwish Mul_182 1 1 997 689\nConvolution Conv_183 1 1 689 1000 0=128 1=1 5=1 6=32768\nSwish Mul_185 1 1 1000 693\nSplit splitncnn_16 1 2 693 693_splitncnn_0 693_splitncnn_1\nInterp Resize_187 1 1 693_splitncnn_1 698 0=1 1=2.000000e+00 2=2.000000e+00\nConcat Concat_188 2 1 698 577_splitncnn_0 699\nSplit splitncnn_17 1 2 699 699_splitncnn_0 699_splitncnn_1\nConvolution Conv_189 1 1 699_splitncnn_1 1003 0=64 1=1 5=1 6=16384\nSwish Mul_191 1 1 1003 703\nConvolution Conv_192 1 1 699_splitncnn_0 1006 0=64 1=1 5=1 6=16384\nSwish Mul_194 1 1 1006 707\nConvolution Conv_195 1 1 703 1009 0=64 1=1 5=1 6=4096\nSwish Mul_197 1 1 1009 711\nConvolution Conv_198 1 1 711 1012 0=64 1=3 4=1 5=1 6=36864\nSwish Mul_200 1 1 1012 715\nConcat Concat_201 2 1 715 707 716\nConvolution Conv_202 1 1 716 1015 0=128 1=1 5=1 6=16384\nSwish Mul_204 1 1 1015 720\nSplit splitncnn_18 1 2 720 720_splitncnn_0 720_splitncnn_1\nConvolution Conv_205 1 1 720_splitncnn_1 1018 0=128 1=3 3=2 4=1 5=1 6=147456\nSwish Mul_207 1 1 1018 724\nConcat Concat_208 2 1 724 693_splitncnn_0 725\nSplit splitncnn_19 1 2 725 725_splitncnn_0 725_splitncnn_1\nConvolution Conv_209 1 1 725_splitncnn_1 1021 0=128 1=1 5=1 6=32768\nSwish Mul_211 1 1 1021 729\nConvolution Conv_212 1 1 725_splitncnn_0 1024 0=128 1=1 5=1 6=32768\nSwish Mul_214 1 1 1024 733\nConvolution Conv_215 1 1 729 1027 0=128 1=1 5=1 6=16384\nSwish Mul_217 1 1 1027 737\nConvolution Conv_218 1 1 737 1030 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_220 1 1 1030 741\nConcat Concat_221 2 1 741 733 742\nConvolution Conv_222 1 1 742 1033 0=256 1=1 5=1 6=65536\nSwish Mul_224 1 1 1033 746\nSplit splitncnn_20 1 2 746 746_splitncnn_0 746_splitncnn_1\nConvolution Conv_225 1 1 746_splitncnn_1 1036 0=256 1=3 3=2 4=1 5=1 6=589824\nSwish Mul_227 1 1 1036 750\nConcat Concat_228 2 1 750 662_splitncnn_0 751\nSplit splitncnn_21 1 2 751 751_splitncnn_0 751_splitncnn_1\nConvolution Conv_229 1 1 751_splitncnn_1 1039 0=256 1=1 5=1 6=131072\nSwish Mul_231 1 1 1039 755\nConvolution Conv_232 1 1 751_splitncnn_0 1042 0=256 1=1 5=1 6=131072\nSwish Mul_234 1 1 1042 759\nConvolution Conv_235 1 1 755 1045 0=256 1=1 5=1 6=65536\nSwish Mul_237 1 1 1045 763\nConvolution Conv_238 1 1 763 1048 0=256 1=3 4=1 5=1 6=589824\nSwish Mul_240 1 1 1048 767\nConcat Concat_241 2 1 767 759 768\nConvolution Conv_242 1 1 768 1051 0=512 1=1 5=1 6=262144\nSwish Mul_244 1 1 1051 772\nConvolution Conv_245 1 1 720_splitncnn_0 1054 0=128 1=1 5=1 6=16384\nSwish Mul_247 1 1 1054 776\nSplit splitncnn_22 1 2 776 776_splitncnn_0 776_splitncnn_1\nConvolution Conv_248 1 1 776_splitncnn_1 1057 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_250 1 1 1057 780\nConvolution Conv_251 1 1 780 1060 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_253 1 1 1060 784\nConvolution Conv_254 1 1 784 797 0=80 1=1 5=1 6=10240 9=4\nConvolution Conv_255 1 1 776_splitncnn_0 1063 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_257 1 1 1063 789\nConvolution Conv_258 1 1 789 1066 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_260 1 1 1066 793\nSplit splitncnn_23 1 2 793 793_splitncnn_0 793_splitncnn_1\nConvolution Conv_261 1 1 793_splitncnn_1 794 0=4 1=1 5=1 6=512\nConvolution Conv_262 1 1 793_splitncnn_0 796 0=1 1=1 5=1 6=128 9=4\nConcat Concat_265 3 1 794 796 797 798\nConvolution Conv_266 1 1 746_splitncnn_0 1069 0=128 1=1 5=1 6=32768\nSwish Mul_268 1 1 1069 802\nSplit splitncnn_24 1 2 802 802_splitncnn_0 802_splitncnn_1\nConvolution Conv_269 1 1 802_splitncnn_1 1072 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_271 1 1 1072 806\nConvolution Conv_272 1 1 806 1075 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_274 1 1 1075 810\nConvolution Conv_275 1 1 810 823 0=80 1=1 5=1 6=10240 9=4\nConvolution Conv_276 1 1 802_splitncnn_0 1078 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_278 1 1 1078 815\nConvolution Conv_279 1 1 815 1081 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_281 1 1 1081 819\nSplit splitncnn_25 1 2 819 819_splitncnn_0 819_splitncnn_1\nConvolution Conv_282 1 1 819_splitncnn_1 820 0=4 1=1 5=1 6=512\nConvolution Conv_283 1 1 819_splitncnn_0 822 0=1 1=1 5=1 6=128 9=4\nConcat Concat_286 3 1 820 822 823 824\nConvolution Conv_287 1 1 772 1084 0=128 1=1 5=1 6=65536\nSwish Mul_289 1 1 1084 828\nSplit splitncnn_26 1 2 828 828_splitncnn_0 828_splitncnn_1\nConvolution Conv_290 1 1 828_splitncnn_1 1087 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_292 1 1 1087 832\nConvolution Conv_293 1 1 832 1090 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_295 1 1 1090 836\nConvolution Conv_296 1 1 836 849 0=80 1=1 5=1 6=10240 9=4\nConvolution Conv_297 1 1 828_splitncnn_0 1093 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_299 1 1 1093 841\nConvolution Conv_300 1 1 841 1096 0=128 1=3 4=1 5=1 6=147456\nSwish Mul_302 1 1 1096 845\nSplit splitncnn_27 1 2 845 845_splitncnn_0 845_splitncnn_1\nConvolution Conv_303 1 1 845_splitncnn_1 846 0=4 1=1 5=1 6=512\nConvolution Conv_304 1 1 845_splitncnn_0 848 0=1 1=1 5=1 6=128 9=4\nConcat Concat_307 3 1 846 848 849 850\nReshape Reshape_315 1 1 798 858 0=-1 1=85\nReshape Reshape_323 1 1 824 866 0=-1 1=85\nReshape Reshape_331 1 1 850 874 0=-1 1=85\nConcat Concat_332 3 1 858 866 874 875 0=1\nPermute Transpose_333 1 1 875 output 0=1\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/MainActivity.java",
"content": "// Some code in this file is based on:\n// https://github.com/nihui/ncnn-android-yolov5/blob/master/app/src/main/java/com/tencent/yolov5ncnn/MainActivity.java\n// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.\n// Copyright (C) Megvii, Inc. and its affiliates. All rights reserved.\n\npackage com.megvii.yoloXncnn;\n\nimport android.app.Activity;\nimport android.content.Intent;\nimport android.graphics.Bitmap;\nimport android.graphics.BitmapFactory;\nimport android.graphics.Canvas;\nimport android.graphics.Color;\nimport android.graphics.Paint;\nimport android.media.ExifInterface;\nimport android.graphics.Matrix;\nimport android.net.Uri;\nimport android.os.Bundle;\nimport android.util.Log;\nimport android.view.View;\nimport android.widget.Button;\nimport android.widget.ImageView;\n\nimport java.io.FileNotFoundException;\nimport java.io.InputStream;\nimport java.io.IOException;\n\npublic class MainActivity extends Activity\n{\n private static final int SELECT_IMAGE = 1;\n\n private ImageView imageView;\n private Bitmap bitmap = null;\n private Bitmap yourSelectedImage = null;\n\n private YOLOXncnn yoloX = new YOLOXncnn();\n\n /** Called when the activity is first created. */\n @Override\n public void onCreate(Bundle savedInstanceState)\n {\n super.onCreate(savedInstanceState);\n setContentView(R.layout.main);\n\n boolean ret_init = yoloX.Init(getAssets());\n if (!ret_init)\n {\n Log.e(\"MainActivity\", \"yoloXncnn Init failed\");\n }\n\n imageView = (ImageView) findViewById(R.id.imageView);\n\n Button buttonImage = (Button) findViewById(R.id.buttonImage);\n buttonImage.setOnClickListener(new View.OnClickListener() {\n @Override\n public void onClick(View arg0) {\n Intent i = new Intent(Intent.ACTION_PICK);\n i.setType(\"image/*\");\n startActivityForResult(i, SELECT_IMAGE);\n }\n });\n\n Button buttonDetect = (Button) findViewById(R.id.buttonDetect);\n buttonDetect.setOnClickListener(new View.OnClickListener() {\n @Override\n public void onClick(View arg0) {\n if (yourSelectedImage == null)\n return;\n YOLOXncnn.Obj[] objects = yoloX.Detect(yourSelectedImage, false);\n\n showObjects(objects);\n }\n });\n\n Button buttonDetectGPU = (Button) findViewById(R.id.buttonDetectGPU);\n buttonDetectGPU.setOnClickListener(new View.OnClickListener() {\n @Override\n public void onClick(View arg0) {\n if (yourSelectedImage == null)\n return;\n\n YOLOXncnn.Obj[] objects = yoloX.Detect(yourSelectedImage, true);\n\n showObjects(objects);\n }\n });\n }\n\n private void showObjects(YOLOXncnn.Obj[] objects)\n {\n if (objects == null)\n {\n imageView.setImageBitmap(bitmap);\n return;\n }\n\n // draw objects on bitmap\n Bitmap rgba = bitmap.copy(Bitmap.Config.ARGB_8888, true);\n\n final int[] colors = new int[] {\n Color.rgb( 54, 67, 244),\n Color.rgb( 99, 30, 233),\n Color.rgb(176, 39, 156),\n Color.rgb(183, 58, 103),\n Color.rgb(181, 81, 63),\n Color.rgb(243, 150, 33),\n Color.rgb(244, 169, 3),\n Color.rgb(212, 188, 0),\n Color.rgb(136, 150, 0),\n Color.rgb( 80, 175, 76),\n Color.rgb( 74, 195, 139),\n Color.rgb( 57, 220, 205),\n Color.rgb( 59, 235, 255),\n Color.rgb( 7, 193, 255),\n Color.rgb( 0, 152, 255),\n Color.rgb( 34, 87, 255),\n Color.rgb( 72, 85, 121),\n Color.rgb(158, 158, 158),\n Color.rgb(139, 125, 96)\n };\n\n Canvas canvas = new Canvas(rgba);\n\n Paint paint = new Paint();\n paint.setStyle(Paint.Style.STROKE);\n paint.setStrokeWidth(4);\n\n Paint textbgpaint = new Paint();\n textbgpaint.setColor(Color.WHITE);\n textbgpaint.setStyle(Paint.Style.FILL);\n\n Paint textpaint = new Paint();\n textpaint.setColor(Color.BLACK);\n textpaint.setTextSize(26);\n textpaint.setTextAlign(Paint.Align.LEFT);\n\n for (int i = 0; i < objects.length; i++)\n {\n paint.setColor(colors[i % 19]);\n\n canvas.drawRect(objects[i].x, objects[i].y, objects[i].x + objects[i].w, objects[i].y + objects[i].h, paint);\n\n // draw filled text inside image\n {\n String text = objects[i].label + \" = \" + String.format(\"%.1f\", objects[i].prob * 100) + \"%\";\n\n float text_width = textpaint.measureText(text);\n float text_height = - textpaint.ascent() + textpaint.descent();\n\n float x = objects[i].x;\n float y = objects[i].y - text_height;\n if (y < 0)\n y = 0;\n if (x + text_width > rgba.getWidth())\n x = rgba.getWidth() - text_width;\n\n canvas.drawRect(x, y, x + text_width, y + text_height, textbgpaint);\n\n canvas.drawText(text, x, y - textpaint.ascent(), textpaint);\n }\n }\n\n imageView.setImageBitmap(rgba);\n }\n\n @Override\n protected void onActivityResult(int requestCode, int resultCode, Intent data)\n {\n super.onActivityResult(requestCode, resultCode, data);\n\n if (resultCode == RESULT_OK && null != data) {\n Uri selectedImage = data.getData();\n\n try\n {\n if (requestCode == SELECT_IMAGE) {\n bitmap = decodeUri(selectedImage);\n\n yourSelectedImage = bitmap.copy(Bitmap.Config.ARGB_8888, true);\n\n imageView.setImageBitmap(bitmap);\n }\n }\n catch (FileNotFoundException e)\n {\n Log.e(\"MainActivity\", \"FileNotFoundException\");\n return;\n }\n }\n }\n\n private Bitmap decodeUri(Uri selectedImage) throws FileNotFoundException\n {\n // Decode image size\n BitmapFactory.Options o = new BitmapFactory.Options();\n o.inJustDecodeBounds = true;\n BitmapFactory.decodeStream(getContentResolver().openInputStream(selectedImage), null, o);\n\n // The new size we want to scale to\n final int REQUIRED_SIZE = 640;\n\n // Find the correct scale value. It should be the power of 2.\n int width_tmp = o.outWidth, height_tmp = o.outHeight;\n int scale = 1;\n while (true) {\n if (width_tmp / 2 < REQUIRED_SIZE || height_tmp / 2 < REQUIRED_SIZE) {\n break;\n }\n width_tmp /= 2;\n height_tmp /= 2;\n scale *= 2;\n }\n\n // Decode with inSampleSize\n BitmapFactory.Options o2 = new BitmapFactory.Options();\n o2.inSampleSize = scale;\n Bitmap bitmap = BitmapFactory.decodeStream(getContentResolver().openInputStream(selectedImage), null, o2);\n\n // Rotate according to EXIF\n int rotate = 0;\n try\n {\n ExifInterface exif = new ExifInterface(getContentResolver().openInputStream(selectedImage));\n int orientation = exif.getAttributeInt(ExifInterface.TAG_ORIENTATION, ExifInterface.ORIENTATION_NORMAL);\n switch (orientation) {\n case ExifInterface.ORIENTATION_ROTATE_270:\n rotate = 270;\n break;\n case ExifInterface.ORIENTATION_ROTATE_180:\n rotate = 180;\n break;\n case ExifInterface.ORIENTATION_ROTATE_90:\n rotate = 90;\n break;\n }\n }\n catch (IOException e)\n {\n Log.e(\"MainActivity\", \"ExifInterface IOException\");\n }\n\n Matrix matrix = new Matrix();\n matrix.postRotate(rotate);\n return Bitmap.createBitmap(bitmap, 0, 0, bitmap.getWidth(), bitmap.getHeight(), matrix, true);\n }\n\n}\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/YOLOXncnn.java",
"content": "// Copyright (C) Megvii, Inc. and its affiliates. All rights reserved.\n\npackage com.megvii.yoloXncnn;\n\nimport android.content.res.AssetManager;\nimport android.graphics.Bitmap;\n\npublic class YOLOXncnn\n{\n public native boolean Init(AssetManager mgr);\n\n public class Obj\n {\n public float x;\n public float y;\n public float w;\n public float h;\n public String label;\n public float prob;\n }\n\n public native Obj[] Detect(Bitmap bitmap, boolean use_gpu);\n\n static {\n System.loadLibrary(\"yoloXncnn\");\n }\n}\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/yoloXncnn.java",
"content": "// Copyright (C) Megvii, Inc. and its affiliates. All rights reserved.\n\npackage com.megvii.yoloXncnn;\n\nimport android.content.res.AssetManager;\nimport android.graphics.Bitmap;\n\npublic class YOLOXncnn\n{\n public native boolean Init(AssetManager mgr);\n\n public class Obj\n {\n public float x;\n public float y;\n public float w;\n public float h;\n public String label;\n public float prob;\n }\n\n public native Obj[] Detect(Bitmap bitmap, boolean use_gpu);\n\n static {\n System.loadLibrary(\"yoloXncnn\");\n }\n}\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/jni/CMakeLists.txt",
"content": "project(yoloXncnn)\n\ncmake_minimum_required(VERSION 3.4.1)\n\nset(ncnn_DIR ${CMAKE_SOURCE_DIR}/ncnn-20210525-android-vulkan/${ANDROID_ABI}/lib/cmake/ncnn)\nfind_package(ncnn REQUIRED)\n\nadd_library(yoloXncnn SHARED yoloXncnn_jni.cpp)\n\ntarget_link_libraries(yoloXncnn\n ncnn\n\n jnigraphics\n)\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/jni/yoloXncnn_jni.cpp",
"content": "// Some code in this file is based on:\n// https://github.com/nihui/ncnn-android-yolov5/blob/master/app/src/main/jni/yolov5ncnn_jni.cpp\n// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.\n// Copyright (C) Megvii, Inc. and its affiliates. All rights reserved.\n\n#include \n#include \n#include \n\n#include \n\n#include \n#include \n\n// ncnn\n#include \"layer.h\"\n#include \"net.h\"\n#include \"benchmark.h\"\n\nstatic ncnn::UnlockedPoolAllocator g_blob_pool_allocator;\nstatic ncnn::PoolAllocator g_workspace_pool_allocator;\n\nstatic ncnn::Net yoloX;\n\nclass YoloV5Focus : public ncnn::Layer\n{\npublic:\n YoloV5Focus()\n {\n one_blob_only = true;\n }\n\n virtual int forward(const ncnn::Mat& bottom_blob, ncnn::Mat& top_blob, const ncnn::Option& opt) const\n {\n int w = bottom_blob.w;\n int h = bottom_blob.h;\n int channels = bottom_blob.c;\n\n int outw = w / 2;\n int outh = h / 2;\n int outc = channels * 4;\n\n top_blob.create(outw, outh, outc, 4u, 1, opt.blob_allocator);\n if (top_blob.empty())\n return -100;\n\n #pragma omp parallel for num_threads(opt.num_threads)\n for (int p = 0; p < outc; p++)\n {\n const float* ptr = bottom_blob.channel(p % channels).row((p / channels) % 2) + ((p / channels) / 2);\n float* outptr = top_blob.channel(p);\n\n for (int i = 0; i < outh; i++)\n {\n for (int j = 0; j < outw; j++)\n {\n *outptr = *ptr;\n\n outptr += 1;\n ptr += 2;\n }\n\n ptr += w;\n }\n }\n\n return 0;\n }\n};\n\nDEFINE_LAYER_CREATOR(YoloV5Focus)\n\nstruct Object\n{\n float x;\n float y;\n float w;\n float h;\n int label;\n float prob;\n};\n\nstruct GridAndStride\n{\n int grid0;\n int grid1;\n int stride;\n};\n\nstatic inline float intersection_area(const Object& a, const Object& b)\n{\n if (a.x > b.x + b.w || a.x + a.w < b.x || a.y > b.y + b.h || a.y + a.h < b.y)\n {\n // no intersection\n return 0.f;\n }\n\n float inter_width = std::min(a.x + a.w, b.x + b.w) - std::max(a.x, b.x);\n float inter_height = std::min(a.y + a.h, b.y + b.h) - std::max(a.y, b.y);\n\n return inter_width * inter_height;\n}\n\nstatic void qsort_descent_inplace(std::vector& faceobjects, int left, int right)\n{\n int i = left;\n int j = right;\n float p = faceobjects[(left + right) / 2].prob;\n\n while (i <= j)\n {\n while (faceobjects[i].prob > p)\n i++;\n\n while (faceobjects[j].prob < p)\n j--;\n\n if (i <= j)\n {\n // swap\n std::swap(faceobjects[i], faceobjects[j]);\n\n i++;\n j--;\n }\n }\n\n #pragma omp parallel sections\n {\n #pragma omp section\n {\n if (left < j) qsort_descent_inplace(faceobjects, left, j);\n }\n #pragma omp section\n {\n if (i < right) qsort_descent_inplace(faceobjects, i, right);\n }\n }\n}\n\nstatic void qsort_descent_inplace(std::vector& faceobjects)\n{\n if (faceobjects.empty())\n return;\n\n qsort_descent_inplace(faceobjects, 0, faceobjects.size() - 1);\n}\n\nstatic void nms_sorted_bboxes(const std::vector& faceobjects, std::vector& picked, float nms_threshold)\n{\n picked.clear();\n\n const int n = faceobjects.size();\n\n std::vector areas(n);\n for (int i = 0; i < n; i++)\n {\n areas[i] = faceobjects[i].w * faceobjects[i].h;\n }\n\n for (int i = 0; i < n; i++)\n {\n const Object& a = faceobjects[i];\n\n int keep = 1;\n for (int j = 0; j < (int)picked.size(); j++)\n {\n const Object& b = faceobjects[picked[j]];\n\n // intersection over union\n float inter_area = intersection_area(a, b);\n float union_area = areas[i] + areas[picked[j]] - inter_area;\n // float IoU = inter_area / union_area\n if (inter_area / union_area > nms_threshold)\n keep = 0;\n }\n\n if (keep)\n picked.push_back(i);\n }\n}\n\nstatic void generate_grids_and_stride(const int target_size, std::vector& strides, std::vector& grid_strides)\n{\n for (auto stride : strides)\n {\n int num_grid = target_size / stride;\n for (int g1 = 0; g1 < num_grid; g1++)\n {\n for (int g0 = 0; g0 < num_grid; g0++)\n {\n grid_strides.push_back((GridAndStride){g0, g1, stride});\n }\n }\n }\n}\n\nstatic void generate_yolox_proposals(std::vector grid_strides, const ncnn::Mat& feat_blob, float prob_threshold, std::vector& objects)\n{\n const int num_grid = feat_blob.h;\n fprintf(stderr, \"output height: %d, width: %d, channels: %d, dims:%d\\n\", feat_blob.h, feat_blob.w, feat_blob.c, feat_blob.dims);\n\n const int num_class = feat_blob.w - 5;\n\n const int num_anchors = grid_strides.size();\n\n const float* feat_ptr = feat_blob.channel(0);\n for (int anchor_idx = 0; anchor_idx < num_anchors; anchor_idx++)\n {\n const int grid0 = grid_strides[anchor_idx].grid0;\n const int grid1 = grid_strides[anchor_idx].grid1;\n const int stride = grid_strides[anchor_idx].stride;\n\n // yolox/models/yolo_head.py decode logic\n // outputs[..., :2] = (outputs[..., :2] + grids) * strides\n // outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides\n float x_center = (feat_ptr[0] + grid0) * stride;\n float y_center = (feat_ptr[1] + grid1) * stride;\n float w = exp(feat_ptr[2]) * stride;\n float h = exp(feat_ptr[3]) * stride;\n float x0 = x_center - w * 0.5f;\n float y0 = y_center - h * 0.5f;\n\n float box_objectness = feat_ptr[4];\n for (int class_idx = 0; class_idx < num_class; class_idx++)\n {\n float box_cls_score = feat_ptr[5 + class_idx];\n float box_prob = box_objectness * box_cls_score;\n if (box_prob > prob_threshold)\n {\n Object obj;\n obj.x = x0;\n obj.y = y0;\n obj.w = w;\n obj.h = h;\n obj.label = class_idx;\n obj.prob = box_prob;\n\n objects.push_back(obj);\n }\n\n } // class loop\n feat_ptr += feat_blob.w;\n\n } // point anchor loop\n}\n\n\nextern \"C\" {\n\n// FIXME DeleteGlobalRef is missing for objCls\nstatic jclass objCls = NULL;\nstatic jmethodID constructortorId;\nstatic jfieldID xId;\nstatic jfieldID yId;\nstatic jfieldID wId;\nstatic jfieldID hId;\nstatic jfieldID labelId;\nstatic jfieldID probId;\n\nJNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved)\n{\n __android_log_print(ANDROID_LOG_DEBUG, \"YOLOXncnn\", \"JNI_OnLoad\");\n\n ncnn::create_gpu_instance();\n\n return JNI_VERSION_1_4;\n}\n\nJNIEXPORT void JNI_OnUnload(JavaVM* vm, void* reserved)\n{\n __android_log_print(ANDROID_LOG_DEBUG, \"YOLOXncnn\", \"JNI_OnUnload\");\n\n ncnn::destroy_gpu_instance();\n}\n\n// public native boolean Init(AssetManager mgr);\nJNIEXPORT jboolean JNICALL Java_com_megvii_yoloXncnn_YOLOXncnn_Init(JNIEnv* env, jobject thiz, jobject assetManager)\n{\n ncnn::Option opt;\n opt.lightmode = true;\n opt.num_threads = 4;\n opt.blob_allocator = &g_blob_pool_allocator;\n opt.workspace_allocator = &g_workspace_pool_allocator;\n opt.use_packing_layout = true;\n\n // use vulkan compute\n if (ncnn::get_gpu_count() != 0)\n opt.use_vulkan_compute = true;\n\n AAssetManager* mgr = AAssetManager_fromJava(env, assetManager);\n\n yoloX.opt = opt;\n\n yoloX.register_custom_layer(\"YoloV5Focus\", YoloV5Focus_layer_creator);\n\n // init param\n {\n int ret = yoloX.load_param(mgr, \"yolox.param\");\n if (ret != 0)\n {\n __android_log_print(ANDROID_LOG_DEBUG, \"YOLOXncnn\", \"load_param failed\");\n return JNI_FALSE;\n }\n }\n\n // init bin\n {\n int ret = yoloX.load_model(mgr, \"yolox.bin\");\n if (ret != 0)\n {\n __android_log_print(ANDROID_LOG_DEBUG, \"YOLOXncnn\", \"load_model failed\");\n return JNI_FALSE;\n }\n }\n\n // init jni glue\n jclass localObjCls = env->FindClass(\"com/megvii/yoloXncnn/YOLOXncnn$Obj\");\n objCls = reinterpret_cast(env->NewGlobalRef(localObjCls));\n\n constructortorId = env->GetMethodID(objCls, \"\", \"(Lcom/megvii/yoloXncnn/YOLOXncnn;)V\");\n\n xId = env->GetFieldID(objCls, \"x\", \"F\");\n yId = env->GetFieldID(objCls, \"y\", \"F\");\n wId = env->GetFieldID(objCls, \"w\", \"F\");\n hId = env->GetFieldID(objCls, \"h\", \"F\");\n labelId = env->GetFieldID(objCls, \"label\", \"Ljava/lang/String;\");\n probId = env->GetFieldID(objCls, \"prob\", \"F\");\n\n return JNI_TRUE;\n}\n\n// public native Obj[] Detect(Bitmap bitmap, boolean use_gpu);\nJNIEXPORT jobjectArray JNICALL Java_com_megvii_yoloXncnn_YOLOXncnn_Detect(JNIEnv* env, jobject thiz, jobject bitmap, jboolean use_gpu)\n{\n if (use_gpu == JNI_TRUE && ncnn::get_gpu_count() == 0)\n {\n return NULL;\n //return env->NewStringUTF(\"no vulkan capable gpu\");\n }\n\n double start_time = ncnn::get_current_time();\n\n AndroidBitmapInfo info;\n AndroidBitmap_getInfo(env, bitmap, &info);\n const int width = info.width;\n const int height = info.height;\n if (info.format != ANDROID_BITMAP_FORMAT_RGBA_8888)\n return NULL;\n\n // parameters which might change for different model\n const int target_size = 640;\n const float prob_threshold = 0.3f;\n const float nms_threshold = 0.65f;\n std::vector strides = {8, 16, 32}; // might have stride=64\n // python 0-1 input tensor with rgb_means = (0.485, 0.456, 0.406), std = (0.229, 0.224, 0.225)\n // so for 0-255 input image, rgb_mean should multiply 255 and norm should div by std.\n const float mean_vals[3] = {255.f * 0.485f, 255.f * 0.456, 255.f * 0.406f};\n const float norm_vals[3] = {1 / (255.f * 0.229f), 1 / (255.f * 0.224f), 1 / (255.f * 0.225f)};\n\n int w = width;\n int h = height;\n float scale = 1.f;\n if (w > h)\n {\n scale = (float)target_size / w;\n w = target_size;\n h = h * scale;\n }\n else\n {\n scale = (float)target_size / h;\n h = target_size;\n w = w * scale;\n }\n\n ncnn::Mat in = ncnn::Mat::from_android_bitmap_resize(env, bitmap, ncnn::Mat::PIXEL_RGB, w, h);\n\n // pad to target_size rectangle\n int wpad = target_size - w;\n int hpad = target_size - h;\n ncnn::Mat in_pad;\n // different from yolov5, yolox only pad on bottom and right side,\n // which means users don't need to extra padding info to decode boxes coordinate.\n ncnn::copy_make_border(in, in_pad, 0, hpad, 0, wpad, ncnn::BORDER_CONSTANT, 114.f);\n\n // yolox\n std::vector objects;\n {\n\n in_pad.substract_mean_normalize(mean_vals, norm_vals);\n\n ncnn::Extractor ex = yoloX.create_extractor();\n\n ex.set_vulkan_compute(use_gpu);\n\n ex.input(\"images\", in_pad);\n\n std::vector proposals;\n\n // yolox decode and generate proposal logic\n {\n ncnn::Mat out;\n ex.extract(\"output\", out);\n\n std::vector grid_strides;\n generate_grids_and_stride(target_size, strides, grid_strides);\n generate_yolox_proposals(grid_strides, out, prob_threshold, proposals);\n\n }\n\n // sort all proposals by score from highest to lowest\n qsort_descent_inplace(proposals);\n\n // apply nms with nms_threshold\n std::vector picked;\n nms_sorted_bboxes(proposals, picked, nms_threshold);\n\n int count = picked.size();\n\n objects.resize(count);\n for (int i = 0; i < count; i++)\n {\n objects[i] = proposals[picked[i]];\n\n // adjust offset to original unpadded\n float x0 = (objects[i].x) / scale;\n float y0 = (objects[i].y) / scale;\n float x1 = (objects[i].x + objects[i].w) / scale;\n float y1 = (objects[i].y + objects[i].h) / scale;\n\n // clip\n x0 = std::max(std::min(x0, (float)(width - 1)), 0.f);\n y0 = std::max(std::min(y0, (float)(height - 1)), 0.f);\n x1 = std::max(std::min(x1, (float)(width - 1)), 0.f);\n y1 = std::max(std::min(y1, (float)(height - 1)), 0.f);\n\n objects[i].x = x0;\n objects[i].y = y0;\n objects[i].w = x1 - x0;\n objects[i].h = y1 - y0;\n }\n }\n\n // objects to Obj[]\n static const char* class_names[] = {\n \"person\", \"bicycle\", \"car\", \"motorcycle\", \"airplane\", \"bus\", \"train\", \"truck\", \"boat\", \"traffic light\",\n \"fire hydrant\", \"stop sign\", \"parking meter\", \"bench\", \"bird\", \"cat\", \"dog\", \"horse\", \"sheep\", \"cow\",\n \"elephant\", \"bear\", \"zebra\", \"giraffe\", \"backpack\", \"umbrella\", \"handbag\", \"tie\", \"suitcase\", \"frisbee\",\n \"skis\", \"snowboard\", \"sports ball\", \"kite\", \"baseball bat\", \"baseball glove\", \"skateboard\", \"surfboard\",\n \"tennis racket\", \"bottle\", \"wine glass\", \"cup\", \"fork\", \"knife\", \"spoon\", \"bowl\", \"banana\", \"apple\",\n \"sandwich\", \"orange\", \"broccoli\", \"carrot\", \"hot dog\", \"pizza\", \"donut\", \"cake\", \"chair\", \"couch\",\n \"potted plant\", \"bed\", \"dining table\", \"toilet\", \"tv\", \"laptop\", \"mouse\", \"remote\", \"keyboard\", \"cell phone\",\n \"microwave\", \"oven\", \"toaster\", \"sink\", \"refrigerator\", \"book\", \"clock\", \"vase\", \"scissors\", \"teddy bear\",\n \"hair drier\", \"toothbrush\"\n };\n\n jobjectArray jObjArray = env->NewObjectArray(objects.size(), objCls, NULL);\n\n for (size_t i=0; iNewObject(objCls, constructortorId, thiz);\n\n env->SetFloatField(jObj, xId, objects[i].x);\n env->SetFloatField(jObj, yId, objects[i].y);\n env->SetFloatField(jObj, wId, objects[i].w);\n env->SetFloatField(jObj, hId, objects[i].h);\n env->SetObjectField(jObj, labelId, env->NewStringUTF(class_names[objects[i].label]));\n env->SetFloatField(jObj, probId, objects[i].prob);\n\n env->SetObjectArrayElement(jObjArray, i, jObj);\n }\n\n double elasped = ncnn::get_current_time() - start_time;\n __android_log_print(ANDROID_LOG_DEBUG, \"YOLOXncnn\", \"%.2fms detect\", elasped);\n\n return jObjArray;\n}\n\n}\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/app/src/main/res/layout/main.xml",
"content": "\n\n\n \n\n \n\n \n yoloXncnn \n \n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/build.gradle",
"content": "// Top-level build file where you can add configuration options common to all sub-projects/modules.\nbuildscript {\n repositories {\n jcenter()\n google()\n }\n dependencies {\n classpath 'com.android.tools.build:gradle:3.5.0'\n }\n}\n\nallprojects {\n repositories {\n jcenter()\n google()\n }\n}\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/gradle/wrapper/gradle-wrapper.properties",
"content": "#Sun Aug 25 10:34:48 CST 2019\ndistributionBase=GRADLE_USER_HOME\ndistributionPath=wrapper/dists\nzipStoreBase=GRADLE_USER_HOME\nzipStorePath=wrapper/dists\ndistributionUrl=https\\://services.gradle.org/distributions/gradle-5.4.1-all.zip\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/gradlew",
"content": "#!/usr/bin/env sh\n\n##############################################################################\n##\n## Gradle start up script for UN*X\n##\n##############################################################################\n\n# Attempt to set APP_HOME\n# Resolve links: $0 may be a link\nPRG=\"$0\"\n# Need this for relative symlinks.\nwhile [ -h \"$PRG\" ] ; do\n ls=`ls -ld \"$PRG\"`\n link=`expr \"$ls\" : '.*-> \\(.*\\)$'`\n if expr \"$link\" : '/.*' > /dev/null; then\n PRG=\"$link\"\n else\n PRG=`dirname \"$PRG\"`\"/$link\"\n fi\ndone\nSAVED=\"`pwd`\"\ncd \"`dirname \\\"$PRG\\\"`/\" >/dev/null\nAPP_HOME=\"`pwd -P`\"\ncd \"$SAVED\" >/dev/null\n\nAPP_NAME=\"Gradle\"\nAPP_BASE_NAME=`basename \"$0\"`\n\n# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.\nDEFAULT_JVM_OPTS=\"\"\n\n# Use the maximum available, or set MAX_FD != -1 to use that value.\nMAX_FD=\"maximum\"\n\nwarn () {\n echo \"$*\"\n}\n\ndie () {\n echo\n echo \"$*\"\n echo\n exit 1\n}\n\n# OS specific support (must be 'true' or 'false').\ncygwin=false\nmsys=false\ndarwin=false\nnonstop=false\ncase \"`uname`\" in\n CYGWIN* )\n cygwin=true\n ;;\n Darwin* )\n darwin=true\n ;;\n MINGW* )\n msys=true\n ;;\n NONSTOP* )\n nonstop=true\n ;;\nesac\n\nCLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar\n\n# Determine the Java command to use to start the JVM.\nif [ -n \"$JAVA_HOME\" ] ; then\n if [ -x \"$JAVA_HOME/jre/sh/java\" ] ; then\n # IBM's JDK on AIX uses strange locations for the executables\n JAVACMD=\"$JAVA_HOME/jre/sh/java\"\n else\n JAVACMD=\"$JAVA_HOME/bin/java\"\n fi\n if [ ! -x \"$JAVACMD\" ] ; then\n die \"ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME\n\nPlease set the JAVA_HOME variable in your environment to match the\nlocation of your Java installation.\"\n fi\nelse\n JAVACMD=\"java\"\n which java >/dev/null 2>&1 || die \"ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.\n\nPlease set the JAVA_HOME variable in your environment to match the\nlocation of your Java installation.\"\nfi\n\n# Increase the maximum file descriptors if we can.\nif [ \"$cygwin\" = \"false\" -a \"$darwin\" = \"false\" -a \"$nonstop\" = \"false\" ] ; then\n MAX_FD_LIMIT=`ulimit -H -n`\n if [ $? -eq 0 ] ; then\n if [ \"$MAX_FD\" = \"maximum\" -o \"$MAX_FD\" = \"max\" ] ; then\n MAX_FD=\"$MAX_FD_LIMIT\"\n fi\n ulimit -n $MAX_FD\n if [ $? -ne 0 ] ; then\n warn \"Could not set maximum file descriptor limit: $MAX_FD\"\n fi\n else\n warn \"Could not query maximum file descriptor limit: $MAX_FD_LIMIT\"\n fi\nfi\n\n# For Darwin, add options to specify how the application appears in the dock\nif $darwin; then\n GRADLE_OPTS=\"$GRADLE_OPTS \\\"-Xdock:name=$APP_NAME\\\" \\\"-Xdock:icon=$APP_HOME/media/gradle.icns\\\"\"\nfi\n\n# For Cygwin, switch paths to Windows format before running java\nif $cygwin ; then\n APP_HOME=`cygpath --path --mixed \"$APP_HOME\"`\n CLASSPATH=`cygpath --path --mixed \"$CLASSPATH\"`\n JAVACMD=`cygpath --unix \"$JAVACMD\"`\n\n # We build the pattern for arguments to be converted via cygpath\n ROOTDIRSRAW=`find -L / -maxdepth 1 -mindepth 1 -type d 2>/dev/null`\n SEP=\"\"\n for dir in $ROOTDIRSRAW ; do\n ROOTDIRS=\"$ROOTDIRS$SEP$dir\"\n SEP=\"|\"\n done\n OURCYGPATTERN=\"(^($ROOTDIRS))\"\n # Add a user-defined pattern to the cygpath arguments\n if [ \"$GRADLE_CYGPATTERN\" != \"\" ] ; then\n OURCYGPATTERN=\"$OURCYGPATTERN|($GRADLE_CYGPATTERN)\"\n fi\n # Now convert the arguments - kludge to limit ourselves to /bin/sh\n i=0\n for arg in \"$@\" ; do\n CHECK=`echo \"$arg\"|egrep -c \"$OURCYGPATTERN\" -`\n CHECK2=`echo \"$arg\"|egrep -c \"^-\"` ### Determine if an option\n\n if [ $CHECK -ne 0 ] && [ $CHECK2 -eq 0 ] ; then ### Added a condition\n eval `echo args$i`=`cygpath --path --ignore --mixed \"$arg\"`\n else\n eval `echo args$i`=\"\\\"$arg\\\"\"\n fi\n i=$((i+1))\n done\n case $i in\n (0) set -- ;;\n (1) set -- \"$args0\" ;;\n (2) set -- \"$args0\" \"$args1\" ;;\n (3) set -- \"$args0\" \"$args1\" \"$args2\" ;;\n (4) set -- \"$args0\" \"$args1\" \"$args2\" \"$args3\" ;;\n (5) set -- \"$args0\" \"$args1\" \"$args2\" \"$args3\" \"$args4\" ;;\n (6) set -- \"$args0\" \"$args1\" \"$args2\" \"$args3\" \"$args4\" \"$args5\" ;;\n (7) set -- \"$args0\" \"$args1\" \"$args2\" \"$args3\" \"$args4\" \"$args5\" \"$args6\" ;;\n (8) set -- \"$args0\" \"$args1\" \"$args2\" \"$args3\" \"$args4\" \"$args5\" \"$args6\" \"$args7\" ;;\n (9) set -- \"$args0\" \"$args1\" \"$args2\" \"$args3\" \"$args4\" \"$args5\" \"$args6\" \"$args7\" \"$args8\" ;;\n esac\nfi\n\n# Escape application args\nsave () {\n for i do printf %s\\\\n \"$i\" | sed \"s/'/'\\\\\\\\''/g;1s/^/'/;\\$s/\\$/' \\\\\\\\/\" ; done\n echo \" \"\n}\nAPP_ARGS=$(save \"$@\")\n\n# Collect all arguments for the java command, following the shell quoting and substitution rules\neval set -- $DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS \"\\\"-Dorg.gradle.appname=$APP_BASE_NAME\\\"\" -classpath \"\\\"$CLASSPATH\\\"\" org.gradle.wrapper.GradleWrapperMain \"$APP_ARGS\"\n\n# by default we should be in the correct project dir, but when run from Finder on Mac, the cwd is wrong\nif [ \"$(uname)\" = \"Darwin\" ] && [ \"$HOME\" = \"$PWD\" ]; then\n cd \"$(dirname \"$0\")\"\nfi\n\nexec \"$JAVACMD\" \"$@\"\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/gradlew.bat",
"content": "@if \"%DEBUG%\" == \"\" @echo off\n@rem ##########################################################################\n@rem\n@rem Gradle startup script for Windows\n@rem\n@rem ##########################################################################\n\n@rem Set local scope for the variables with windows NT shell\nif \"%OS%\"==\"Windows_NT\" setlocal\n\nset DIRNAME=%~dp0\nif \"%DIRNAME%\" == \"\" set DIRNAME=.\nset APP_BASE_NAME=%~n0\nset APP_HOME=%DIRNAME%\n\n@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.\nset DEFAULT_JVM_OPTS=\n\n@rem Find java.exe\nif defined JAVA_HOME goto findJavaFromJavaHome\n\nset JAVA_EXE=java.exe\n%JAVA_EXE% -version >NUL 2>&1\nif \"%ERRORLEVEL%\" == \"0\" goto init\n\necho.\necho ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.\necho.\necho Please set the JAVA_HOME variable in your environment to match the\necho location of your Java installation.\n\ngoto fail\n\n:findJavaFromJavaHome\nset JAVA_HOME=%JAVA_HOME:\"=%\nset JAVA_EXE=%JAVA_HOME%/bin/java.exe\n\nif exist \"%JAVA_EXE%\" goto init\n\necho.\necho ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%\necho.\necho Please set the JAVA_HOME variable in your environment to match the\necho location of your Java installation.\n\ngoto fail\n\n:init\n@rem Get command-line arguments, handling Windows variants\n\nif not \"%OS%\" == \"Windows_NT\" goto win9xME_args\n\n:win9xME_args\n@rem Slurp the command line arguments.\nset CMD_LINE_ARGS=\nset _SKIP=2\n\n:win9xME_args_slurp\nif \"x%~1\" == \"x\" goto execute\n\nset CMD_LINE_ARGS=%*\n\n:execute\n@rem Setup the command line\n\nset CLASSPATH=%APP_HOME%\\gradle\\wrapper\\gradle-wrapper.jar\n\n@rem Execute Gradle\n\"%JAVA_EXE%\" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% \"-Dorg.gradle.appname=%APP_BASE_NAME%\" -classpath \"%CLASSPATH%\" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%\n\n:end\n@rem End local scope for the variables with windows NT shell\nif \"%ERRORLEVEL%\"==\"0\" goto mainEnd\n\n:fail\nrem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of\nrem the _cmd.exe /c_ return code!\nif not \"\" == \"%GRADLE_EXIT_CONSOLE%\" exit 1\nexit /b 1\n\n:mainEnd\nif \"%OS%\"==\"Windows_NT\" endlocal\n\n:omega\n"
},
{
"path": "detector/YOLOX/demo/ncnn/android/settings.gradle",
"content": "include ':app'\n"
},
{
"path": "detector/YOLOX/demo/ncnn/cpp/README.md",
"content": "# YOLOX-CPP-ncnn\n\nCpp file compile of YOLOX object detection base on [ncnn](https://github.com/Tencent/ncnn).\n\n## Tutorial\n\n### Step1\nClone [ncnn](https://github.com/Tencent/ncnn) first, then please following [build tutorial of ncnn](https://github.com/Tencent/ncnn/wiki/how-to-build) to build on your own device.\n\n### Step2\nUse provided tools to generate onnx file.\nFor example, if you want to generate onnx file of yolox-s, please run the following command:\n```shell\ncd \npython3 tools/export_onnx.py -n yolox-s\n```\nThen, a yolox.onnx file is generated.\n\n### Step3\nGenerate ncnn param and bin file.\n```shell\ncd \ncd build/tools/ncnn\n./onnx2ncnn yolox.onnx yolox.param yolox.bin\n```\n\nSince Focus module is not supported in ncnn. Warnings like:\n```shell\nUnsupported slice step ! \n```\nwill be printed. However, don't worry! C++ version of Focus layer is already implemented in yolox.cpp.\n\n### Step4\nOpen **yolox.param**, and modify it.\nBefore (just an example):\n```\n295 328\nInput images 0 1 images\nSplit splitncnn_input0 1 4 images images_splitncnn_0 images_splitncnn_1 images_splitncnn_2 images_splitncnn_3\nCrop Slice_4 1 1 images_splitncnn_3 647 -23309=1,0 -23310=1,2147483647 -23311=1,1\nCrop Slice_9 1 1 647 652 -23309=1,0 -23310=1,2147483647 -23311=1,2\nCrop Slice_14 1 1 images_splitncnn_2 657 -23309=1,0 -23310=1,2147483647 -23311=1,1\nCrop Slice_19 1 1 657 662 -23309=1,1 -23310=1,2147483647 -23311=1,2\nCrop Slice_24 1 1 images_splitncnn_1 667 -23309=1,1 -23310=1,2147483647 -23311=1,1\nCrop Slice_29 1 1 667 672 -23309=1,0 -23310=1,2147483647 -23311=1,2\nCrop Slice_34 1 1 images_splitncnn_0 677 -23309=1,1 -23310=1,2147483647 -23311=1,1\nCrop Slice_39 1 1 677 682 -23309=1,1 -23310=1,2147483647 -23311=1,2\nConcat Concat_40 4 1 652 672 662 682 683 0=0\n...\n```\n* Change first number for 295 to 295 - 9 = 286(since we will remove 10 layers and add 1 layers, total layers number should minus 9). \n* Then remove 10 lines of code from Split to Concat, but remember the last but 2nd number: 683.\n* Add YoloV5Focus layer After Input (using previous number 683):\n```\nYoloV5Focus focus 1 1 images 683\n```\nAfter(just an exmaple):\n```\n286 328\nInput images 0 1 images\nYoloV5Focus focus 1 1 images 683\n...\n```\n\n### Step5\nUse onnx_optimze to generate new param and bin:\n```shell\n# suppose you are still under ncnn/build/tools/ncnn dir.\n../ncnnoptimize model.param model.bin yolox.param yolox.bin 65536\n```\n\n### Step6\nCopy or Move yolox.cpp file into ncnn/examples, modify the CMakeList.txt, then build yolox\n\n### Step7\nInference image with executable file yolox, enjoy the detect result:\n```shell\n./yolox demo.jpg\n```\n\n## Acknowledgement\n\n* [ncnn](https://github.com/Tencent/ncnn)\n"
},
{
"path": "detector/YOLOX/demo/ncnn/cpp/yolox.cpp",
"content": "// This file is wirtten base on the following file:\n// https://github.com/Tencent/ncnn/blob/master/examples/yolov5.cpp\n// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.\n// Licensed under the BSD 3-Clause License (the \"License\"); you may not use this file except\n// in compliance with the License. You may obtain a copy of the License at\n//\n// https://opensource.org/licenses/BSD-3-Clause\n//\n// Unless required by applicable law or agreed to in writing, software distributed\n// under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR\n// CONDITIONS OF ANY KIND, either express or implied. See the License for the\n// specific language governing permissions and limitations under the License.\n// ------------------------------------------------------------------------------\n// Copyright (C) 2020-2021, Megvii Inc. All rights reserved.\n\n#include \"layer.h\"\n#include \"net.h\"\n\n#if defined(USE_NCNN_SIMPLEOCV)\n#include \"simpleocv.h\"\n#else\n#include \n#include \n#include \n#endif\n#include \n#include \n#include \n\n// YOLOX use the same focus in yolov5\nclass YoloV5Focus : public ncnn::Layer\n{\npublic:\n YoloV5Focus()\n {\n one_blob_only = true;\n }\n\n virtual int forward(const ncnn::Mat& bottom_blob, ncnn::Mat& top_blob, const ncnn::Option& opt) const\n {\n int w = bottom_blob.w;\n int h = bottom_blob.h;\n int channels = bottom_blob.c;\n\n int outw = w / 2;\n int outh = h / 2;\n int outc = channels * 4;\n\n top_blob.create(outw, outh, outc, 4u, 1, opt.blob_allocator);\n if (top_blob.empty())\n return -100;\n\n #pragma omp parallel for num_threads(opt.num_threads)\n for (int p = 0; p < outc; p++)\n {\n const float* ptr = bottom_blob.channel(p % channels).row((p / channels) % 2) + ((p / channels) / 2);\n float* outptr = top_blob.channel(p);\n\n for (int i = 0; i < outh; i++)\n {\n for (int j = 0; j < outw; j++)\n {\n *outptr = *ptr;\n\n outptr += 1;\n ptr += 2;\n }\n\n ptr += w;\n }\n }\n\n return 0;\n }\n};\n\nDEFINE_LAYER_CREATOR(YoloV5Focus)\n\nstruct Object\n{\n cv::Rect_ rect;\n int label;\n float prob;\n};\n\nstruct GridAndStride\n{\n int grid0;\n int grid1;\n int stride;\n};\n\nstatic inline float intersection_area(const Object& a, const Object& b)\n{\n cv::Rect_ inter = a.rect & b.rect;\n return inter.area();\n}\n\nstatic void qsort_descent_inplace(std::vector& faceobjects, int left, int right)\n{\n int i = left;\n int j = right;\n float p = faceobjects[(left + right) / 2].prob;\n\n while (i <= j)\n {\n while (faceobjects[i].prob > p)\n i++;\n\n while (faceobjects[j].prob < p)\n j--;\n\n if (i <= j)\n {\n // swap\n std::swap(faceobjects[i], faceobjects[j]);\n\n i++;\n j--;\n }\n }\n\n #pragma omp parallel sections\n {\n #pragma omp section\n {\n if (left < j) qsort_descent_inplace(faceobjects, left, j);\n }\n #pragma omp section\n {\n if (i < right) qsort_descent_inplace(faceobjects, i, right);\n }\n }\n}\n\nstatic void qsort_descent_inplace(std::vector& objects)\n{\n if (objects.empty())\n return;\n\n qsort_descent_inplace(objects, 0, objects.size() - 1);\n}\n\nstatic void nms_sorted_bboxes(const std::vector& faceobjects, std::vector& picked, float nms_threshold)\n{\n picked.clear();\n\n const int n = faceobjects.size();\n\n std::vector areas(n);\n for (int i = 0; i < n; i++)\n {\n areas[i] = faceobjects[i].rect.area();\n }\n\n for (int i = 0; i < n; i++)\n {\n const Object& a = faceobjects[i];\n\n int keep = 1;\n for (int j = 0; j < (int)picked.size(); j++)\n {\n const Object& b = faceobjects[picked[j]];\n\n // intersection over union\n float inter_area = intersection_area(a, b);\n float union_area = areas[i] + areas[picked[j]] - inter_area;\n // float IoU = inter_area / union_area\n if (inter_area / union_area > nms_threshold)\n keep = 0;\n }\n\n if (keep)\n picked.push_back(i);\n }\n}\n\nstatic void generate_grids_and_stride(const int target_size, std::vector& strides, std::vector& grid_strides)\n{\n for (auto stride : strides)\n {\n int num_grid = target_size / stride;\n for (int g1 = 0; g1 < num_grid; g1++)\n {\n for (int g0 = 0; g0 < num_grid; g0++)\n {\n grid_strides.push_back((GridAndStride){g0, g1, stride});\n }\n }\n }\n}\n\nstatic void generate_yolox_proposals(std::vector grid_strides, const ncnn::Mat& feat_blob, float prob_threshold, std::vector& objects)\n{\n const int num_grid = feat_blob.h;\n fprintf(stderr, \"output height: %d, width: %d, channels: %d, dims:%d\\n\", feat_blob.h, feat_blob.w, feat_blob.c, feat_blob.dims);\n\n const int num_class = feat_blob.w - 5;\n\n const int num_anchors = grid_strides.size();\n\n const float* feat_ptr = feat_blob.channel(0);\n for (int anchor_idx = 0; anchor_idx < num_anchors; anchor_idx++)\n {\n const int grid0 = grid_strides[anchor_idx].grid0;\n const int grid1 = grid_strides[anchor_idx].grid1;\n const int stride = grid_strides[anchor_idx].stride;\n\n // yolox/models/yolo_head.py decode logic\n // outputs[..., :2] = (outputs[..., :2] + grids) * strides\n // outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides\n float x_center = (feat_ptr[0] + grid0) * stride;\n float y_center = (feat_ptr[1] + grid1) * stride;\n float w = exp(feat_ptr[2]) * stride;\n float h = exp(feat_ptr[3]) * stride;\n float x0 = x_center - w * 0.5f;\n float y0 = y_center - h * 0.5f;\n\n float box_objectness = feat_ptr[4];\n for (int class_idx = 0; class_idx < num_class; class_idx++)\n {\n float box_cls_score = feat_ptr[5 + class_idx];\n float box_prob = box_objectness * box_cls_score;\n if (box_prob > prob_threshold)\n {\n Object obj;\n obj.rect.x = x0;\n obj.rect.y = y0;\n obj.rect.width = w;\n obj.rect.height = h;\n obj.label = class_idx;\n obj.prob = box_prob;\n\n objects.push_back(obj);\n }\n\n } // class loop\n feat_ptr += feat_blob.w;\n\n } // point anchor loop\n}\n \nstatic int detect_yolox(const cv::Mat& bgr, std::vector& objects)\n{\n ncnn::Net yolox;\n\n yolox.opt.use_vulkan_compute = true;\n // yolox.opt.use_bf16_storage = true;\n\n yolox.register_custom_layer(\"YoloV5Focus\", YoloV5Focus_layer_creator);\n\n // original pretrained model from https://github.com/yolox\n // TODO ncnn model https://github.com/nihui/ncnn-assets/tree/master/models\n yolox.load_param(\"yolox.param\");\n yolox.load_model(\"yolox.bin\");\n\n const int target_size = 416;\n const float prob_threshold = 0.3f;\n const float nms_threshold = 0.65f;\n\n int img_w = bgr.cols;\n int img_h = bgr.rows;\n\n int w = img_w;\n int h = img_h;\n float scale = 1.f;\n if (w > h)\n {\n scale = (float)target_size / w;\n w = target_size;\n h = h * scale;\n }\n else\n {\n scale = (float)target_size / h;\n h = target_size;\n w = w * scale;\n }\n ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, img_w, img_h, w, h);\n\n // pad to target_size rectangle\n int wpad = target_size - w;\n int hpad = target_size - h;\n ncnn::Mat in_pad;\n // different from yolov5, yolox only pad on bottom and right side,\n // which means users don't need to extra padding info to decode boxes coordinate.\n ncnn::copy_make_border(in, in_pad, 0, hpad, 0, wpad, ncnn::BORDER_CONSTANT, 114.f);\n\n // python 0-1 input tensor with rgb_means = (0.485, 0.456, 0.406), std = (0.229, 0.224, 0.225)\n // so for 0-255 input image, rgb_mean should multiply 255 and norm should div by std.\n const float mean_vals[3] = {255.f * 0.485f, 255.f * 0.456, 255.f * 0.406f};\n const float norm_vals[3] = {1 / (255.f * 0.229f), 1 / (255.f * 0.224f), 1 / (255.f * 0.225f)};\n\n in_pad.substract_mean_normalize(mean_vals, norm_vals);\n\n ncnn::Extractor ex = yolox.create_extractor();\n\n ex.input(\"images\", in_pad);\n\n std::vector proposals;\n\n {\n ncnn::Mat out;\n ex.extract(\"output\", out);\n\n std::vector strides = {8, 16, 32}; // might have stride=64\n std::vector grid_strides;\n generate_grids_and_stride(target_size, strides, grid_strides);\n generate_yolox_proposals(grid_strides, out, prob_threshold, proposals);\n }\n\n // sort all proposals by score from highest to lowest\n qsort_descent_inplace(proposals);\n\n // apply nms with nms_threshold\n std::vector picked;\n nms_sorted_bboxes(proposals, picked, nms_threshold);\n\n int count = picked.size();\n\n objects.resize(count);\n for (int i = 0; i < count; i++)\n {\n objects[i] = proposals[picked[i]];\n\n // adjust offset to original unpadded\n float x0 = (objects[i].rect.x) / scale;\n float y0 = (objects[i].rect.y) / scale;\n float x1 = (objects[i].rect.x + objects[i].rect.width) / scale;\n float y1 = (objects[i].rect.y + objects[i].rect.height) / scale;\n\n // clip\n x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f);\n y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f);\n x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f);\n y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f);\n\n objects[i].rect.x = x0;\n objects[i].rect.y = y0;\n objects[i].rect.width = x1 - x0;\n objects[i].rect.height = y1 - y0;\n }\n\n return 0;\n}\n\nstatic void draw_objects(const cv::Mat& bgr, const std::vector& objects)\n{\n static const char* class_names[] = {\n \"person\", \"bicycle\", \"car\", \"motorcycle\", \"airplane\", \"bus\", \"train\", \"truck\", \"boat\", \"traffic light\",\n \"fire hydrant\", \"stop sign\", \"parking meter\", \"bench\", \"bird\", \"cat\", \"dog\", \"horse\", \"sheep\", \"cow\",\n \"elephant\", \"bear\", \"zebra\", \"giraffe\", \"backpack\", \"umbrella\", \"handbag\", \"tie\", \"suitcase\", \"frisbee\",\n \"skis\", \"snowboard\", \"sports ball\", \"kite\", \"baseball bat\", \"baseball glove\", \"skateboard\", \"surfboard\",\n \"tennis racket\", \"bottle\", \"wine glass\", \"cup\", \"fork\", \"knife\", \"spoon\", \"bowl\", \"banana\", \"apple\",\n \"sandwich\", \"orange\", \"broccoli\", \"carrot\", \"hot dog\", \"pizza\", \"donut\", \"cake\", \"chair\", \"couch\",\n \"potted plant\", \"bed\", \"dining table\", \"toilet\", \"tv\", \"laptop\", \"mouse\", \"remote\", \"keyboard\", \"cell phone\",\n \"microwave\", \"oven\", \"toaster\", \"sink\", \"refrigerator\", \"book\", \"clock\", \"vase\", \"scissors\", \"teddy bear\",\n \"hair drier\", \"toothbrush\"\n };\n\n cv::Mat image = bgr.clone();\n\n for (size_t i = 0; i < objects.size(); i++)\n {\n const Object& obj = objects[i];\n\n fprintf(stderr, \"%d = %.5f at %.2f %.2f %.2f x %.2f\\n\", obj.label, obj.prob,\n obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height);\n\n cv::rectangle(image, obj.rect, cv::Scalar(255, 0, 0));\n\n char text[256];\n sprintf(text, \"%s %.1f%%\", class_names[obj.label], obj.prob * 100);\n\n int baseLine = 0;\n cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);\n\n int x = obj.rect.x;\n int y = obj.rect.y - label_size.height - baseLine;\n if (y < 0)\n y = 0;\n if (x + label_size.width > image.cols)\n x = image.cols - label_size.width;\n\n cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),\n cv::Scalar(255, 255, 255), -1);\n\n cv::putText(image, text, cv::Point(x, y + label_size.height),\n cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));\n }\n\n cv::imshow(\"image\", image);\n cv::waitKey(0);\n}\n\nint main(int argc, char** argv)\n{\n if (argc != 2)\n {\n fprintf(stderr, \"Usage: %s [imagepath]\\n\", argv[0]);\n return -1;\n }\n\n const char* imagepath = argv[1];\n\n cv::Mat m = cv::imread(imagepath, 1);\n if (m.empty())\n {\n fprintf(stderr, \"cv::imread %s failed\\n\", imagepath);\n return -1;\n }\n\n std::vector objects;\n detect_yolox(m, objects);\n\n draw_objects(m, objects);\n\n return 0;\n}\n"
},
{
"path": "detector/YOLOX/demo.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport os\nimport pdb\nimport time\nfrom loguru import logger\n\nimport cv2\n\nimport torch\n\nfrom yolox.data.data_augment import preproc\nfrom yolox.data.datasets import COCO_CLASSES\nfrom yolox.exp import get_exp\nfrom yolox.utils import fuse_model, get_model_info, postprocess, vis\n\nIMAGE_EXT = ['.jpg', '.jpeg', '.webp', '.bmp', '.png']\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"YOLOX Demo!\")\n parser.add_argument('demo', default='image', help='demo type, eg. image, video and webcam')\n parser.add_argument(\"-expn\", \"--experiment-name\", type=str, default=None)\n parser.add_argument(\"-n\", \"--name\", type=str, default=None, help=\"model name\")\n\n parser.add_argument('--path', default='./assets/dog.jpg', help='path to images or video')\n parser.add_argument('--camid', type=int, default=0, help='webcam demo camera id')\n parser.add_argument(\n '--save_result', action='store_true',\n help='whether to save the inference result of image/video'\n )\n\n # exp file\n parser.add_argument(\n \"-f\",\n \"--exp_file\",\n default=None,\n type=str,\n help=\"pls input your expriment description file\",\n )\n parser.add_argument(\"-c\", \"--ckpt\", default=None, type=str, help=\"ckpt for eval\")\n parser.add_argument(\"--device\", default=\"cpu\", type=str, help=\"device to run our model, can either be cpu or gpu\")\n parser.add_argument(\"--conf\", default=None, type=float, help=\"test conf\")\n parser.add_argument(\"--nms\", default=None, type=float, help=\"test nms threshold\")\n parser.add_argument(\"--tsize\", default=None, type=int, help=\"test img size\")\n parser.add_argument(\n \"--fp16\",\n dest=\"fp16\",\n default=False,\n action=\"store_true\",\n help=\"Adopting mix precision evaluating.\",\n )\n parser.add_argument(\n \"--fuse\",\n dest=\"fuse\",\n default=False,\n action=\"store_true\",\n help=\"Fuse conv and bn for testing.\",\n )\n parser.add_argument(\n \"--trt\",\n dest=\"trt\",\n default=False,\n action=\"store_true\",\n help=\"Using TensorRT model for testing.\",\n )\n return parser\n\n\ndef get_image_list(path):\n image_names = []\n for maindir, subdir, file_name_list in os.walk(path):\n for filename in file_name_list:\n apath = os.path.join(maindir, filename)\n ext = os.path.splitext(apath)[1]\n if ext in IMAGE_EXT:\n image_names.append(apath)\n return image_names\n\n\nclass Predictor(object):\n def __init__(self, model, exp, cls_names=COCO_CLASSES, trt_file=None, decoder=None, device=\"cpu\"):\n self.model = model\n self.cls_names = cls_names\n self.decoder = decoder\n self.num_classes = exp.num_classes\n self.confthre = exp.test_conf\n self.nmsthre = exp.nmsthre\n self.test_size = exp.test_size\n self.device = device\n if trt_file is not None:\n from torch2trt import TRTModule\n model_trt = TRTModule()\n model_trt.load_state_dict(torch.load(trt_file))\n\n x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda()\n self.model(x)\n self.model = model_trt\n self.rgb_means = (0.485, 0.456, 0.406)\n self.std = (0.229, 0.224, 0.225)\n\n def inference(self, img):\n img_info = {'id': 0}\n if isinstance(img, str):\n img_info['file_name'] = os.path.basename(img)\n img = cv2.imread(img)\n else:\n img_info['file_name'] = None\n\n height, width = img.shape[:2]\n img_info['height'] = height\n img_info['width'] = width\n img_info['raw_img'] = img\n\n img, ratio = preproc(img, self.test_size, self.rgb_means, self.std)\n img_info['ratio'] = ratio\n img = torch.from_numpy(img).unsqueeze(0)\n if self.device == \"gpu\":\n img = img.cuda()\n\n with torch.no_grad():\n t0 = time.time()\n outputs = self.model(img)\n if self.decoder is not None:\n outputs = self.decoder(outputs, dtype=outputs.type())\n outputs = postprocess(\n outputs, self.num_classes, self.confthre, self.nmsthre\n )\n logger.info('Infer time: {:.4f}s'.format(time.time()-t0))\n return outputs, img_info\n\n def visual(self, output, img_info, cls_conf=0.35):\n ratio = img_info['ratio']\n img = img_info['raw_img']\n if output is None:\n return img\n output = output.cpu()\n\n bboxes = output[:, 0:4]\n\n # preprocessing: resize\n bboxes /= ratio\n\n cls = output[:, 6]\n scores = output[:, 4] * output[:, 5]\n\n vis_res = vis(img, bboxes, scores, cls, cls_conf, self.cls_names)\n return vis_res\n\n\ndef image_demo(predictor, vis_folder, path, current_time, save_result):\n if os.path.isdir(path):\n files = get_image_list(path)\n else:\n files = [path]\n files.sort()\n for image_name in files:\n outputs, img_info = predictor.inference(image_name)\n result_image = predictor.visual(outputs[0], img_info)\n if save_result:\n save_folder = os.path.join(\n vis_folder, time.strftime(\"%Y_%m_%d_%H_%M_%S\", current_time)\n )\n os.makedirs(save_folder, exist_ok=True)\n save_file_name = os.path.join(save_folder, os.path.basename(image_name))\n logger.info(\"Saving detection result in {}\".format(save_file_name))\n cv2.imwrite(save_file_name, result_image)\n ch = cv2.waitKey(0)\n if ch == 27 or ch == ord('q') or ch == ord('Q'):\n break\n\n\ndef imageflow_demo(predictor, vis_folder, current_time, args):\n cap = cv2.VideoCapture(args.path if args.demo == 'video' else args.camid)\n width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float\n height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float\n fps = cap.get(cv2.CAP_PROP_FPS)\n save_folder = os.path.join(vis_folder, time.strftime(\"%Y_%m_%d_%H_%M_%S\", current_time))\n os.makedirs(save_folder, exist_ok=True)\n if args.demo == \"video\":\n save_path = os.path.join(save_folder, args.path.split('/')[-1])\n else:\n save_path = os.path.join(save_folder, 'camera.mp4')\n logger.info(f'video save_path is {save_path}')\n vid_writer = cv2.VideoWriter(\n save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (int(width), int(height))\n )\n while True:\n ret_val, frame = cap.read()\n if ret_val:\n outputs, img_info = predictor.inference(frame)\n result_frame = predictor.visual(outputs[0], img_info)\n if args.save_result:\n vid_writer.write(result_frame)\n ch = cv2.waitKey(1)\n if ch == 27 or ch == ord('q') or ch == ord('Q'):\n break\n else:\n break\n\n\ndef main(exp, args):\n if not args.experiment_name:\n args.experiment_name = exp.exp_name\n\n file_name = os.path.join(exp.output_dir, args.experiment_name)\n os.makedirs(file_name, exist_ok=True)\n\n if args.save_result:\n vis_folder = os.path.join(file_name, 'vis_res')\n os.makedirs(vis_folder, exist_ok=True)\n\n if args.trt:\n args.device=\"gpu\"\n\n logger.info(\"Args: {}\".format(args))\n\n if args.conf is not None:\n exp.test_conf = args.conf\n if args.nms is not None:\n exp.nmsthre = args.nms\n if args.tsize is not None:\n exp.test_size = (args.tsize, args.tsize)\n\n model = exp.get_model()\n logger.info(\"Model Summary: {}\".format(get_model_info(model, exp.test_size)))\n\n if args.device == \"gpu\":\n model.cuda()\n model.eval()\n\n if not args.trt:\n if args.ckpt is None:\n ckpt_file = os.path.join(file_name, \"best_ckpt.pth.tar\")\n else:\n ckpt_file = args.ckpt\n logger.info(\"loading checkpoint\")\n ckpt = torch.load(ckpt_file, map_location=\"cpu\")\n # load the model state dict\n model.load_state_dict(ckpt[\"model\"])\n logger.info(\"loaded checkpoint done.\")\n\n if args.fuse:\n logger.info(\"\\tFusing model...\")\n model = fuse_model(model)\n\n if args.trt:\n assert (not args.fuse),\\\n \"TensorRT model is not support model fusing!\"\n trt_file = os.path.join(file_name, \"model_trt.pth\")\n assert os.path.exists(trt_file), (\n \"TensorRT model is not found!\\n Run python3 tools/trt.py first!\"\n )\n model.head.decode_in_inference = False\n decoder = model.head.decode_outputs\n logger.info(\"Using TensorRT to inference\")\n else:\n trt_file = None\n decoder = None\n\n predictor = Predictor(model, exp, COCO_CLASSES, trt_file, decoder, args.device)\n current_time = time.localtime()\n if args.demo == 'image':\n image_demo(predictor, vis_folder, args.path, current_time, args.save_result)\n elif args.demo == 'video' or args.demo == 'webcam':\n imageflow_demo(predictor, vis_folder, current_time, args)\n\n\nif __name__ == \"__main__\":\n args = make_parser().parse_args()\n exp = get_exp(args.exp_file, args.name)\n\n main(exp, args)\n"
},
{
"path": "detector/YOLOX/docs/train_custom_data.md",
"content": "# Train Custom Data.\nThis page explains how to train your own custom data with YOLOX.\n\nWe take an example of fine-tuning YOLOX-S model on VOC dataset to give a more clear guide.\n\n## 0. Before you start\nClone this repo and follow the [README](../README.md) to install YOLOX.\n\n## 1. Create your own dataset\n**Step 1** Prepare your own dataset with images and labels first. For labeling images, you may use a tool like [Labelme](https://github.com/wkentaro/labelme) or [CVAT](https://github.com/openvinotoolkit/cvat).\n\n**Step 2** Then, you should write the corresponding Dataset Class which can load images and labels through \"\\_\\_getitem\\_\\_\" method. We currently support COCO format and VOC format.\n\nYou can also write the Dataset by you own. Let's take the [VOC](../yolox/data/datasets/voc.py#L151) Dataset file for example:\n```python\n @Dataset.resize_getitem\n def __getitem__(self, index):\n img, target, img_info, img_id = self.pull_item(index)\n\n if self.preproc is not None:\n img, target = self.preproc(img, target, self.input_dim)\n\n return img, target, img_info, img_id\n```\n\nOne more thing worth noting is that you should also implement \"[pull_item](../yolox/data/datasets/voc.py#L129)\" and \"[load_anno](../yolox/data/datasets/voc.py#L121)\" method for the Mosiac and MixUp augmentation.\n\n**Step 3** Prepare the evaluator. We currently have [COCO evaluator](../yolox/evaluators/coco_evaluator.py) and [VOC evaluator](../yolox/evaluators/voc_evaluator.py).\nIf you have your own format data or evaluation metric, you may write your own evaluator.\n\n**Step 4** Put your dataset under $YOLOX_DIR/datasets$, for VOC:\n```shell\nln -s /path/to/your/VOCdevkit ./datasets/VOCdevkit\n```\n* The path \"VOCdevkit\" will be used in your exp file described in next section.Specifically, in \"get_data_loader\" and \"get_eval_loader\" function.\n\n## 2. Create your Exp file to control everything\nWe put everything involved in a model to one single Exp file, including model setting, training setting, and testing setting.\n\n**A complete Exp file is at [yolox_base.py](../yolox/exp/yolox_base.py).** It may be too long to write for every exp, but you can inherit the base Exp file and only overwrite the changed part.\n\nLet's still take the [VOC Exp file](../exps/example/yolox_voc/yolox_voc_s.py) for an example.\n\nWe select YOLOX-S model here, so we should change the network depth and width. VOC has only 20 classes, so we should also change the num_classes.\n\nThese configs are changed in the init() methd:\n```python\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.num_classes = 20\n self.depth = 0.33\n self.width = 0.50\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n```\n\nBesides, you should also overwrite the dataset and evaluator preprared before to training the model on your own data.\n\nPlease see \"[get_data_loader](../exps/example/yolox_voc/yolox_voc_s.py#L20)\", \"[get_eval_loader](../exps/example/yolox_voc/yolox_voc_s.py#L82)\", and \"[get_evaluator](../exps/example/yolox_voc/yolox_voc_s.py#L113)\" for more details.\n\n## 3. Train\nExcept special cases, we always recommend to use our [COCO pretrained weights](../README.md) for initializing.\n\nOnce you get the Exp file and the COCO pretrained weights we provided, you can train your own model by the following command:\n```bash\npython tools/train.py -f /path/to/your/Exp/file -d 8 -b 64 --fp16 -o -c /path/to/the/pretrained/weights\n```\n\nor take the YOLOX-S VOC training for example:\n```bash\npython tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 8 -b 64 --fp16 -o -c /path/to/yolox_s.pth.tar\n```\n\n(Don't worry for the different shape of detection head between the pretrained weights and your own model, we will handle it)\n\n## 4. Tips for Best Training Results\n\nAs YOLOX is an anchor-free detector with only several hyper-parameters, most of the time good results can be obtained with no changes to the models or training settings.\nWe thus always recommend you first train with all default training settings.\n\nIf at first you don't get good results, there are steps you could consider to take to improve.\n\n**Model Selection** We provide YOLOX-Nano, YOLOX-Tiny, and YOLOX-S for mobile deployments, while YOLOX-M/L/X for cloud or high performance GPU deployments.\n\nIf your deployment meets some trouble of compatibility. we recommand YOLOX-DarkNet53.\n\n**Training Configs** If your training overfits early, then you can reduce max\\_epochs or decrease the base\\_lr and min\\_lr\\_ratio in your Exp file:\n```python\n# -------------- training config --------------------- #\n self.warmup_epochs = 5\n self.max_epoch = 300\n self.warmup_lr = 0\n self.basic_lr_per_img = 0.01 / 64.0\n self.scheduler = \"yoloxwarmcos\"\n self.no_aug_epochs = 15\n self.min_lr_ratio = 0.05\n self.ema = True\n\n self.weight_decay = 5e-4\n self.momentum = 0.9\n```\n\n**Aug Configs** You may also change the degree of the augmentations.\n\nGenerally, for small models, you should weak the aug, while for large models or small size of dataset, you may enchance the aug in your Exp file:\n```python\n# --------------- transform config ----------------- #\n self.degrees = 10.0\n self.translate = 0.1\n self.scale = (0.1, 2)\n self.mscale = (0.8, 1.6)\n self.shear = 2.0\n self.perspective = 0.0\n self.enable_mixup = True\n```\n\n**Design your own detector** You may refer to our [Arxiv](https://arxiv.org/abs/2107.08430) paper for details and suggestions for designing your own detector.\n"
},
{
"path": "detector/YOLOX/exps/default/nano.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\nimport torch.nn as nn\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 0.33\n self.width = 0.25\n self.scale = (0.5, 1.5)\n self.random_size = (10, 20)\n self.test_size = (416, 416)\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n self.enable_mixup = False\n\n def get_model(self, sublinear=False):\n\n def init_yolo(M):\n for m in M.modules():\n if isinstance(m, nn.BatchNorm2d):\n m.eps = 1e-3\n m.momentum = 0.03\n if \"model\" not in self.__dict__:\n from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead\n in_channels = [256, 512, 1024]\n # NANO model use depthwise = True, which is main difference.\n backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels, depthwise=True)\n head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels, depthwise=True)\n self.model = YOLOX(backbone, head)\n\n self.model.apply(init_yolo)\n self.model.head.initialize_biases(1e-2)\n return self.model\n"
},
{
"path": "detector/YOLOX/exps/default/yolov3.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\nimport torch\nimport torch.nn as nn\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 1.0\n self.width = 1.0\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n\n def get_model(self, sublinear=False):\n def init_yolo(M):\n for m in M.modules():\n if isinstance(m, nn.BatchNorm2d):\n m.eps = 1e-3\n m.momentum = 0.03\n if \"model\" not in self.__dict__:\n from yolox.models import YOLOX, YOLOFPN, YOLOXHead\n backbone = YOLOFPN()\n head = YOLOXHead(self.num_classes, self.width, in_channels=[128, 256, 512], act=\"lrelu\")\n self.model = YOLOX(backbone, head)\n self.model.apply(init_yolo)\n self.model.head.initialize_biases(1e-2)\n\n return self.model\n\n def get_data_loader(self, batch_size, is_distributed, no_aug=False):\n from data.datasets.cocodataset import COCODataset\n from data.datasets.mosaicdetection import MosaicDetection\n from data.datasets.data_augment import TrainTransform\n from data.datasets.dataloading import YoloBatchSampler, DataLoader, InfiniteSampler\n import torch.distributed as dist\n\n dataset = COCODataset(\n data_dir='data/COCO/',\n json_file=self.train_ann,\n img_size=self.input_size,\n preproc=TrainTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n max_labels=50\n ),\n )\n\n dataset = MosaicDetection(\n dataset,\n mosaic=not no_aug,\n img_size=self.input_size,\n preproc=TrainTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n max_labels=120\n ),\n degrees=self.degrees,\n translate=self.translate,\n scale=self.scale,\n shear=self.shear,\n perspective=self.perspective,\n )\n\n self.dataset = dataset\n\n if is_distributed:\n batch_size = batch_size // dist.get_world_size()\n sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0)\n else:\n sampler = torch.utils.data.RandomSampler(self.dataset)\n\n batch_sampler = YoloBatchSampler(\n sampler=sampler,\n batch_size=batch_size,\n drop_last=False,\n input_dimension=self.input_size,\n mosaic=not no_aug\n )\n\n dataloader_kwargs = {\"num_workers\": self.data_num_workers, \"pin_memory\": True}\n dataloader_kwargs[\"batch_sampler\"] = batch_sampler\n train_loader = DataLoader(self.dataset, **dataloader_kwargs)\n\n return train_loader\n"
},
{
"path": "detector/YOLOX/exps/default/yolox_l.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 1.0\n self.width = 1.0\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n"
},
{
"path": "detector/YOLOX/exps/default/yolox_m.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 0.67\n self.width = 0.75\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n"
},
{
"path": "detector/YOLOX/exps/default/yolox_s.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 0.33\n self.width = 0.50\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n"
},
{
"path": "detector/YOLOX/exps/default/yolox_tiny.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 0.33\n self.width = 0.375\n self.scale = (0.5, 1.5)\n self.random_size = (10, 20)\n self.test_size = (416, 416)\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n self.enable_mixup = False\n"
},
{
"path": "detector/YOLOX/exps/default/yolox_x.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.depth = 1.33\n self.width = 1.25\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n"
},
{
"path": "detector/YOLOX/exps/example/yolox_voc/yolox_voc_s.py",
"content": "# encoding: utf-8\nimport os\nimport random\nimport torch\nimport torch.nn as nn\nimport torch.distributed as dist\n\nfrom yolox.exp import Exp as MyExp\nfrom yolox.data import get_yolox_datadir\n\nclass Exp(MyExp):\n def __init__(self):\n super(Exp, self).__init__()\n self.num_classes = 20\n self.depth = 0.33\n self.width = 0.50\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n\n def get_data_loader(self, batch_size, is_distributed, no_aug=False):\n from yolox.data import (\n VOCDetection,\n TrainTransform,\n YoloBatchSampler,\n DataLoader,\n InfiniteSampler,\n MosaicDetection,\n )\n\n dataset = VOCDetection(\n data_dir=os.path.join(get_yolox_datadir(), \"VOCdevkit\"),\n image_sets=[('2007', 'trainval'), ('2012', 'trainval')],\n img_size=self.input_size,\n preproc=TrainTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n max_labels=50,\n ),\n )\n\n dataset = MosaicDetection(\n dataset,\n mosaic=not no_aug,\n img_size=self.input_size,\n preproc=TrainTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n max_labels=120,\n ),\n degrees=self.degrees,\n translate=self.translate,\n scale=self.scale,\n shear=self.shear,\n perspective=self.perspective,\n enable_mixup=self.enable_mixup,\n )\n\n self.dataset = dataset\n\n if is_distributed:\n batch_size = batch_size // dist.get_world_size()\n\n sampler = InfiniteSampler(\n len(self.dataset), seed=self.seed if self.seed else 0\n )\n\n batch_sampler = YoloBatchSampler(\n sampler=sampler,\n batch_size=batch_size,\n drop_last=False,\n input_dimension=self.input_size,\n mosaic=not no_aug,\n )\n\n dataloader_kwargs = {\"num_workers\": self.data_num_workers, \"pin_memory\": True}\n dataloader_kwargs[\"batch_sampler\"] = batch_sampler\n train_loader = DataLoader(self.dataset, **dataloader_kwargs)\n\n return train_loader\n\n def get_eval_loader(self, batch_size, is_distributed, testdev=False):\n from yolox.data import VOCDetection, ValTransform\n\n valdataset = VOCDetection(\n data_dir=os.path.join(get_yolox_datadir(), \"VOCdevkit\"),\n image_sets=[('2007', 'test')],\n img_size=self.test_size,\n preproc=ValTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n ),\n )\n\n if is_distributed:\n batch_size = batch_size // dist.get_world_size()\n sampler = torch.utils.data.distributed.DistributedSampler(\n valdataset, shuffle=False\n )\n else:\n sampler = torch.utils.data.SequentialSampler(valdataset)\n\n dataloader_kwargs = {\n \"num_workers\": self.data_num_workers,\n \"pin_memory\": True,\n \"sampler\": sampler,\n }\n dataloader_kwargs[\"batch_size\"] = batch_size\n val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)\n\n return val_loader\n\n def get_evaluator(self, batch_size, is_distributed, testdev=False):\n from yolox.evaluators import VOCEvaluator\n\n val_loader = self.get_eval_loader(batch_size, is_distributed, testdev=testdev)\n evaluator = VOCEvaluator(\n dataloader=val_loader,\n img_size=self.test_size,\n confthre=self.test_conf,\n nmsthre=self.nmsthre,\n num_classes=self.num_classes,\n )\n return evaluator\n"
},
{
"path": "detector/YOLOX/requirements.txt",
"content": "numpy\ntorch>=1.7\nopencv_python\nloguru\nscikit-image\ntqdm\ntorchvision\nPillow\nthop\nninja\ntabulate\ntensorboard\nonnxruntime\n\n\n"
},
{
"path": "detector/YOLOX/setup.cfg",
"content": "[isort]\nline_length = 100\nmulti_line_output = 3\nbalanced_wrapping = True\nknown_standard_library = setuptools\nknown_third_party = tqdm,loguru\nknown_data_processing = cv2,numpy,scipy,PIL,matplotlib,scikit_image\nknown_datasets = pycocotools\nknown_deeplearning = torch,torchvision,caffe2,onnx,apex,timm,thop,torch2trt,tensorrt,openvino,onnxruntime\nknown_myself = yolox\nsections = FUTURE,STDLIB,THIRDPARTY,data_processing,datasets,deeplearning,myself,FIRSTPARTY,LOCALFOLDER\nno_lines_before=STDLIB,THIRDPARTY,datasets\ndefault_section = FIRSTPARTY\n\n[flake8]\nmax-line-length = 100\nmax-complexity = 18\nexclude = __init__.py\n"
},
{
"path": "detector/YOLOX/setup.py",
"content": "#!/usr/bin/env python\n# Copyright (c) Megvii, Inc. and its affiliates. All Rights Reserved\n\nimport re\nimport setuptools\nimport glob\nfrom os import path\nimport torch\nfrom torch.utils.cpp_extension import CppExtension\n\ntorch_ver = [int(x) for x in torch.__version__.split(\".\")[:2]]\nassert torch_ver >= [1, 3], \"Requires PyTorch >= 1.3\"\n\n\ndef get_extensions():\n this_dir = path.dirname(path.abspath(__file__))\n extensions_dir = path.join(this_dir, \"yolox\", \"layers\", \"csrc\")\n\n main_source = path.join(extensions_dir, \"vision.cpp\")\n sources = glob.glob(path.join(extensions_dir, \"**\", \"*.cpp\"))\n\n sources = [main_source] + sources\n extension = CppExtension\n\n extra_compile_args = {\"cxx\": [\"-O3\"]}\n define_macros = []\n\n include_dirs = [extensions_dir]\n\n ext_modules = [\n extension(\n \"yolox._C\",\n sources,\n include_dirs=include_dirs,\n define_macros=define_macros,\n extra_compile_args=extra_compile_args,\n )\n ]\n\n return ext_modules\n\n\nwith open(\"yolox/__init__.py\", \"r\") as f:\n version = re.search(\n r'^__version__\\s*=\\s*[\\'\"]([^\\'\"]*)[\\'\"]',\n f.read(), re.MULTILINE\n ).group(1)\n\n\nwith open(\"README.md\", \"r\") as f:\n long_description = f.read()\n\n\nsetuptools.setup(\n name=\"yolox\",\n version=version,\n author=\"basedet team\",\n python_requires=\">=3.6\",\n long_description=long_description,\n ext_modules=None,\n classifiers=[\"Programming Language :: Python :: 3\", \"Operating System :: OS Independent\"],\n cmdclass={\"build_ext\": torch.utils.cpp_extension.BuildExtension},\n packages=setuptools.find_packages(),\n)\n"
},
{
"path": "detector/YOLOX/tools/__init__.py",
"content": "###################################################################\n# File Name: __init__.py\n# Author: Zhongdao Wang\n# mail: wcd17@mails.tsinghua.edu.cn\n# Created Time: Sun Jul 25 17:14:12 2021\n###################################################################\n\nfrom __future__ import print_function\nfrom __future__ import division\nfrom __future__ import absolute_import\n"
},
{
"path": "detector/YOLOX/tools/demo.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport os\nimport time\nfrom loguru import logger\n\nimport cv2\n\nimport torch\n\nfrom yolox.data.data_augment import preproc\nfrom yolox.data.datasets import COCO_CLASSES\nfrom yolox.exp import get_exp\nfrom yolox.utils import fuse_model, get_model_info, postprocess, vis\n\nIMAGE_EXT = ['.jpg', '.jpeg', '.webp', '.bmp', '.png']\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"YOLOX Demo!\")\n parser.add_argument('demo', default='image', help='demo type, eg. image, video and webcam')\n parser.add_argument(\"-expn\", \"--experiment-name\", type=str, default=None)\n parser.add_argument(\"-n\", \"--name\", type=str, default=None, help=\"model name\")\n\n parser.add_argument('--path', default='./assets/dog.jpg', help='path to images or video')\n parser.add_argument('--camid', type=int, default=0, help='webcam demo camera id')\n parser.add_argument(\n '--save_result', action='store_true',\n help='whether to save the inference result of image/video'\n )\n\n # exp file\n parser.add_argument(\n \"-f\",\n \"--exp_file\",\n default=None,\n type=str,\n help=\"pls input your expriment description file\",\n )\n parser.add_argument(\"-c\", \"--ckpt\", default=None, type=str, help=\"ckpt for eval\")\n parser.add_argument(\"--device\", default=\"cpu\", type=str, help=\"device to run our model, can either be cpu or gpu\")\n parser.add_argument(\"--conf\", default=None, type=float, help=\"test conf\")\n parser.add_argument(\"--nms\", default=None, type=float, help=\"test nms threshold\")\n parser.add_argument(\"--tsize\", default=None, type=int, help=\"test img size\")\n parser.add_argument(\n \"--fp16\",\n dest=\"fp16\",\n default=False,\n action=\"store_true\",\n help=\"Adopting mix precision evaluating.\",\n )\n parser.add_argument(\n \"--fuse\",\n dest=\"fuse\",\n default=False,\n action=\"store_true\",\n help=\"Fuse conv and bn for testing.\",\n )\n parser.add_argument(\n \"--trt\",\n dest=\"trt\",\n default=False,\n action=\"store_true\",\n help=\"Using TensorRT model for testing.\",\n )\n return parser\n\n\ndef get_image_list(path):\n image_names = []\n for maindir, subdir, file_name_list in os.walk(path):\n for filename in file_name_list:\n apath = os.path.join(maindir, filename)\n ext = os.path.splitext(apath)[1]\n if ext in IMAGE_EXT:\n image_names.append(apath)\n return image_names\n\n\nclass Predictor(object):\n def __init__(self, model, exp, cls_names=COCO_CLASSES, trt_file=None, decoder=None, device=\"cpu\"):\n self.model = model\n self.cls_names = cls_names\n self.decoder = decoder\n self.num_classes = exp.num_classes\n self.confthre = exp.test_conf\n self.nmsthre = exp.nmsthre\n self.test_size = exp.test_size\n self.device = device\n if trt_file is not None:\n from torch2trt import TRTModule\n model_trt = TRTModule()\n model_trt.load_state_dict(torch.load(trt_file))\n\n x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda()\n self.model(x)\n self.model = model_trt\n self.rgb_means = (0.485, 0.456, 0.406)\n self.std = (0.229, 0.224, 0.225)\n\n def inference(self, img):\n img_info = {'id': 0}\n if isinstance(img, str):\n img_info['file_name'] = os.path.basename(img)\n img = cv2.imread(img)\n else:\n img_info['file_name'] = None\n\n height, width = img.shape[:2]\n img_info['height'] = height\n img_info['width'] = width\n img_info['raw_img'] = img\n img, ratio = preproc(img, self.test_size, self.rgb_means, self.std)\n img_info['ratio'] = ratio\n img = torch.from_numpy(img).unsqueeze(0)\n if self.device == \"gpu\":\n img = img.cuda()\n\n with torch.no_grad():\n t0 = time.time()\n outputs = self.model(img)\n if self.decoder is not None:\n outputs = self.decoder(outputs, dtype=outputs.type())\n outputs = postprocess(\n outputs, self.num_classes, self.confthre, self.nmsthre\n )\n logger.info('Infer time: {:.4f}s'.format(time.time()-t0))\n return outputs, img_info\n\n def visual(self, output, img_info, cls_conf=0.35):\n ratio = img_info['ratio']\n img = img_info['raw_img']\n if output is None:\n return img\n output = output.cpu()\n\n bboxes = output[:, 0:4]\n\n # preprocessing: resize\n bboxes /= ratio\n\n cls = output[:, 6]\n scores = output[:, 4] * output[:, 5]\n\n vis_res = vis(img, bboxes, scores, cls, cls_conf, self.cls_names)\n return vis_res\n\n\ndef image_demo(predictor, vis_folder, path, current_time, save_result):\n if os.path.isdir(path):\n files = get_image_list(path)\n else:\n files = [path]\n files.sort()\n for image_name in files:\n outputs, img_info = predictor.inference(image_name)\n result_image = predictor.visual(outputs[0], img_info)\n if save_result:\n save_folder = os.path.join(\n vis_folder, time.strftime(\"%Y_%m_%d_%H_%M_%S\", current_time)\n )\n os.makedirs(save_folder, exist_ok=True)\n save_file_name = os.path.join(save_folder, os.path.basename(image_name))\n logger.info(\"Saving detection result in {}\".format(save_file_name))\n cv2.imwrite(save_file_name, result_image)\n ch = cv2.waitKey(0)\n if ch == 27 or ch == ord('q') or ch == ord('Q'):\n break\n\n\ndef imageflow_demo(predictor, vis_folder, current_time, args):\n cap = cv2.VideoCapture(args.path if args.demo == 'video' else args.camid)\n width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float\n height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float\n fps = cap.get(cv2.CAP_PROP_FPS)\n save_folder = os.path.join(vis_folder, time.strftime(\"%Y_%m_%d_%H_%M_%S\", current_time))\n os.makedirs(save_folder, exist_ok=True)\n if args.demo == \"video\":\n save_path = os.path.join(save_folder, args.path.split('/')[-1])\n else:\n save_path = os.path.join(save_folder, 'camera.mp4')\n logger.info(f'video save_path is {save_path}')\n vid_writer = cv2.VideoWriter(\n save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (int(width), int(height))\n )\n while True:\n ret_val, frame = cap.read()\n if ret_val:\n outputs, img_info = predictor.inference(frame)\n result_frame = predictor.visual(outputs[0], img_info)\n if args.save_result:\n vid_writer.write(result_frame)\n ch = cv2.waitKey(1)\n if ch == 27 or ch == ord('q') or ch == ord('Q'):\n break\n else:\n break\n\n\ndef main(exp, args):\n if not args.experiment_name:\n args.experiment_name = exp.exp_name\n\n file_name = os.path.join(exp.output_dir, args.experiment_name)\n os.makedirs(file_name, exist_ok=True)\n\n if args.save_result:\n vis_folder = os.path.join(file_name, 'vis_res')\n os.makedirs(vis_folder, exist_ok=True)\n\n if args.trt:\n args.device = \"gpu\"\n\n logger.info(\"Args: {}\".format(args))\n\n if args.conf is not None:\n exp.test_conf = args.conf\n if args.nms is not None:\n exp.nmsthre = args.nms\n if args.tsize is not None:\n exp.test_size = (args.tsize, args.tsize)\n\n model = exp.get_model()\n logger.info(\"Model Summary: {}\".format(get_model_info(model, exp.test_size)))\n\n if args.device == \"gpu\":\n model.cuda()\n model.eval()\n\n if not args.trt:\n if args.ckpt is None:\n ckpt_file = os.path.join(file_name, \"best_ckpt.pth.tar\")\n else:\n ckpt_file = args.ckpt\n logger.info(\"loading checkpoint\")\n ckpt = torch.load(ckpt_file, map_location=\"cpu\")\n # load the model state dict\n model.load_state_dict(ckpt[\"model\"])\n logger.info(\"loaded checkpoint done.\")\n\n if args.fuse:\n logger.info(\"\\tFusing model...\")\n model = fuse_model(model)\n\n if args.trt:\n assert (not args.fuse),\\\n \"TensorRT model is not support model fusing!\"\n trt_file = os.path.join(file_name, \"model_trt.pth\")\n assert os.path.exists(trt_file), (\n \"TensorRT model is not found!\\n Run python3 tools/trt.py first!\"\n )\n model.head.decode_in_inference = False\n decoder = model.head.decode_outputs\n logger.info(\"Using TensorRT to inference\")\n else:\n trt_file = None\n decoder = None\n\n predictor = Predictor(model, exp, COCO_CLASSES, trt_file, decoder, args.device)\n current_time = time.localtime()\n if args.demo == 'image':\n image_demo(predictor, vis_folder, args.path, current_time, args.save_result)\n elif args.demo == 'video' or args.demo == 'webcam':\n imageflow_demo(predictor, vis_folder, current_time, args)\n\n\nif __name__ == \"__main__\":\n args = make_parser().parse_args()\n exp = get_exp(args.exp_file, args.name)\n\n main(exp, args)\n"
},
{
"path": "detector/YOLOX/tools/eval.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport os\nimport random\nimport warnings\nfrom loguru import logger\n\nimport torch\nimport torch.backends.cudnn as cudnn\nfrom torch.nn.parallel import DistributedDataParallel as DDP\n\nfrom yolox.core import launch\nfrom yolox.exp import get_exp\nfrom yolox.utils import configure_nccl, fuse_model, get_local_rank, get_model_info, setup_logger\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"YOLOX Eval\")\n parser.add_argument(\"-expn\", \"--experiment-name\", type=str, default=None)\n parser.add_argument(\"-n\", \"--name\", type=str, default=None, help=\"model name\")\n\n # distributed\n parser.add_argument(\n \"--dist-backend\", default=\"nccl\", type=str, help=\"distributed backend\"\n )\n parser.add_argument(\n \"--dist-url\", default=None, type=str, help=\"url used to set up distributed training\"\n )\n parser.add_argument(\"-b\", \"--batch-size\", type=int, default=64, help=\"batch size\")\n parser.add_argument(\n \"-d\", \"--devices\", default=None, type=int, help=\"device for training\"\n )\n parser.add_argument(\n \"--local_rank\", default=0, type=int, help=\"local rank for dist training\"\n )\n parser.add_argument(\n \"--num_machine\", default=1, type=int, help=\"num of node for training\"\n )\n parser.add_argument(\n \"--machine_rank\", default=0, type=int, help=\"node rank for multi-node training\"\n )\n parser.add_argument(\n \"-f\",\n \"--exp_file\",\n default=None,\n type=str,\n help=\"pls input your expriment description file\",\n )\n parser.add_argument(\"-c\", \"--ckpt\", default=None, type=str, help=\"ckpt for eval\")\n parser.add_argument(\"--conf\", default=None, type=float, help=\"test conf\")\n parser.add_argument(\"--nms\", default=None, type=float, help=\"test nms threshold\")\n parser.add_argument(\"--tsize\", default=None, type=int, help=\"test img size\")\n parser.add_argument(\"--seed\", default=None, type=int, help=\"eval seed\")\n parser.add_argument(\n \"--fp16\",\n dest=\"fp16\",\n default=False,\n action=\"store_true\",\n help=\"Adopting mix precision evaluating.\",\n )\n parser.add_argument(\n \"--fuse\",\n dest=\"fuse\",\n default=False,\n action=\"store_true\",\n help=\"Fuse conv and bn for testing.\",\n )\n parser.add_argument(\n \"--trt\",\n dest=\"trt\",\n default=False,\n action=\"store_true\",\n help=\"Using TensorRT model for testing.\",\n )\n parser.add_argument(\n \"--test\",\n dest=\"test\",\n default=False,\n action=\"store_true\",\n help=\"Evaluating on test-dev set.\",\n )\n parser.add_argument(\n \"--speed\", dest=\"speed\", default=False, action=\"store_true\", help=\"speed test only.\"\n )\n parser.add_argument(\n \"opts\",\n help=\"Modify config options using the command-line\",\n default=None,\n nargs=argparse.REMAINDER,\n )\n return parser\n\n\n@logger.catch\ndef main(exp, num_gpu, args):\n if not args.experiment_name:\n args.experiment_name = exp.exp_name\n\n if args.seed is not None:\n random.seed(args.seed)\n torch.manual_seed(args.seed)\n cudnn.deterministic = True\n warnings.warn(\n \"You have chosen to seed testing. This will turn on the CUDNN deterministic setting, \"\n )\n\n is_distributed = num_gpu > 1\n\n # set environment variables for distributed training\n configure_nccl()\n cudnn.benchmark = True\n\n # rank = args.local_rank\n rank = get_local_rank()\n\n if rank == 0:\n if os.path.exists(\"./\" + args.experiment_name + \"ip_add.txt\"):\n os.remove(\"./\" + args.experiment_name + \"ip_add.txt\")\n\n file_name = os.path.join(exp.output_dir, args.experiment_name)\n\n if rank == 0:\n os.makedirs(file_name, exist_ok=True)\n\n setup_logger(\n file_name, distributed_rank=rank, filename=\"val_log.txt\", mode=\"a\"\n )\n logger.info(\"Args: {}\".format(args))\n\n if args.conf is not None:\n exp.test_conf = args.conf\n if args.nms is not None:\n exp.nmsthre = args.nms\n if args.tsize is not None:\n exp.test_size = (args.tsize, args.tsize)\n\n model = exp.get_model()\n logger.info(\"Model Summary: {}\".format(get_model_info(model, exp.test_size)))\n logger.info(\"Model Structure:\\n{}\".format(str(model)))\n\n evaluator = exp.get_evaluator(args.batch_size, is_distributed, args.test)\n\n torch.cuda.set_device(rank)\n model.cuda(rank)\n model.eval()\n\n if not args.speed and not args.trt:\n if args.ckpt is None:\n ckpt_file = os.path.join(file_name, \"best_ckpt.pth.tar\")\n else:\n ckpt_file = args.ckpt\n logger.info(\"loading checkpoint\")\n loc = \"cuda:{}\".format(rank)\n ckpt = torch.load(ckpt_file, map_location=loc)\n # load the model state dict\n model.load_state_dict(ckpt[\"model\"])\n logger.info(\"loaded checkpoint done.\")\n\n if is_distributed:\n model = DDP(model, device_ids=[rank])\n\n if args.fuse:\n logger.info(\"\\tFusing model...\")\n model = fuse_model(model)\n\n if args.trt:\n assert (not args.fuse and not is_distributed and args.batch_size == 1),\\\n \"TensorRT model is not support model fusing and distributed inferencing!\"\n trt_file = os.path.join(file_name, \"model_trt.pth\")\n assert os.path.exists(trt_file), \"TensorRT model is not found!\\n Run tools/trt.py first!\"\n model.head.decode_in_inference = False\n decoder = model.head.decode_outputs\n else:\n trt_file = None\n decoder = None\n\n # start evaluate\n *_, summary = evaluator.evaluate(\n model, is_distributed, args.fp16, trt_file, decoder, exp.test_size\n )\n logger.info(\"\\n\" + summary)\n\n\nif __name__ == \"__main__\":\n args = make_parser().parse_args()\n exp = get_exp(args.exp_file, args.name)\n exp.merge(args.opts)\n\n num_gpu = torch.cuda.device_count() if args.devices is None else args.devices\n assert num_gpu <= torch.cuda.device_count()\n\n dist_url = \"auto\" if args.dist_url is None else args.dist_url\n launch(\n main, num_gpu, args.num_machine, backend=args.dist_backend,\n dist_url=dist_url, args=(exp, num_gpu, args)\n )\n"
},
{
"path": "detector/YOLOX/tools/export_onnx.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport os\nfrom loguru import logger\n\nimport torch\nfrom torch import nn\n\nfrom yolox.exp import get_exp\nfrom yolox.models.network_blocks import SiLU\nfrom yolox.utils import replace_module\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"YOLOX onnx deploy\")\n parser.add_argument(\n \"--output-name\", type=str, default=\"yolox.onnx\", help=\"output name of models\"\n )\n parser.add_argument(\"--input\", default=\"images\", type=str, help=\"input name of onnx model\")\n parser.add_argument(\"--output\", default=\"output\", type=str, help=\"output name of onnx model\")\n parser.add_argument(\"-o\", \"--opset\", default=11, type=int, help=\"onnx opset version\")\n parser.add_argument(\"--no-onnxsim\", action=\"store_true\", help=\"use onnxsim or not\")\n\n parser.add_argument(\n \"-f\",\n \"--exp_file\",\n default=None,\n type=str,\n help=\"expriment description file\",\n )\n parser.add_argument(\"-expn\", \"--experiment-name\", type=str, default=None)\n parser.add_argument(\"-n\", \"--name\", type=str, default=None, help=\"model name\")\n parser.add_argument(\"-c\", \"--ckpt\", default=None, type=str, help=\"ckpt path\")\n parser.add_argument(\n \"opts\",\n help=\"Modify config options using the command-line\",\n default=None,\n nargs=argparse.REMAINDER,\n )\n\n return parser\n\n\n@logger.catch\ndef main():\n args = make_parser().parse_args()\n logger.info(\"args value: {}\".format(args))\n exp = get_exp(args.exp_file, args.name)\n exp.merge(args.opts)\n\n if not args.experiment_name:\n args.experiment_name = exp.exp_name\n\n model = exp.get_model()\n if args.ckpt is None:\n file_name = os.path.join(exp.output_dir, args.experiment_name)\n ckpt_file = os.path.join(file_name, \"best_ckpt.pth.tar\")\n else:\n ckpt_file = args.ckpt\n\n ckpt = torch.load(ckpt_file, map_location=\"cpu\")\n # load the model state dict\n\n model.eval()\n if \"model\" in ckpt:\n ckpt = ckpt[\"model\"]\n model.load_state_dict(ckpt)\n model = replace_module(model, nn.SiLU, SiLU)\n model.head.decode_in_inference = False\n\n logger.info(\"loaded checkpoint done.\")\n dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])\n torch.onnx._export(\n model,\n dummy_input,\n args.output_name,\n input_names=[args.input],\n output_names=[args.output],\n opset_version=args.opset,\n )\n logger.info(\"generate onnx named {}\".format(args.output_name))\n\n if not args.no_onnxsim:\n # use onnxsimplify to reduce reduent model.\n os.system(\"python3 -m onnxsim {} {}\".format(args.output_name, args.output_name))\n logger.info(\"generate simplify onnx named {}\".format(args.output_name))\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "detector/YOLOX/tools/train.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport random\nimport warnings\nfrom loguru import logger\n\nimport torch\nimport torch.backends.cudnn as cudnn\n\nfrom yolox.core import Trainer, launch\nfrom yolox.exp import get_exp\nfrom yolox.utils import configure_nccl\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"YOLOX train parser\")\n parser.add_argument(\"-expn\", \"--experiment-name\", type=str, default=None)\n parser.add_argument(\"-n\", \"--name\", type=str, default=None, help=\"model name\")\n\n # distributed\n parser.add_argument(\n \"--dist-backend\", default=\"nccl\", type=str, help=\"distributed backend\"\n )\n parser.add_argument(\n \"--dist-url\", default=None, type=str, help=\"url used to set up distributed training\"\n )\n parser.add_argument(\"-b\", \"--batch-size\", type=int, default=64, help=\"batch size\")\n parser.add_argument(\n \"-d\", \"--devices\", default=None, type=int, help=\"device for training\"\n )\n parser.add_argument(\n \"--local_rank\", default=0, type=int, help=\"local rank for dist training\"\n )\n parser.add_argument(\n \"-f\",\n \"--exp_file\",\n default=None,\n type=str,\n help=\"plz input your expriment description file\",\n )\n parser.add_argument(\n \"--resume\", default=False, action=\"store_true\", help=\"resume training\"\n )\n parser.add_argument(\"-c\", \"--ckpt\", default=None, type=str, help=\"checkpoint file\")\n parser.add_argument(\n \"-e\", \"--start_epoch\", default=None, type=int, help=\"resume training start epoch\"\n )\n parser.add_argument(\n \"--num_machine\", default=1, type=int, help=\"num of node for training\"\n )\n parser.add_argument(\n \"--machine_rank\", default=0, type=int, help=\"node rank for multi-node training\"\n )\n parser.add_argument(\n \"--fp16\",\n dest=\"fp16\",\n default=True,\n action=\"store_true\",\n help=\"Adopting mix precision training.\",\n )\n parser.add_argument(\n \"-o\",\n \"--occumpy\",\n dest=\"occumpy\",\n default=False,\n action=\"store_true\",\n help=\"occumpy GPU memory first for training.\",\n )\n parser.add_argument(\n \"opts\",\n help=\"Modify config options using the command-line\",\n default=None,\n nargs=argparse.REMAINDER,\n )\n return parser\n\n\n@logger.catch\ndef main(exp, args):\n if not args.experiment_name:\n args.experiment_name = exp.exp_name\n\n if exp.seed is not None:\n random.seed(exp.seed)\n torch.manual_seed(exp.seed)\n cudnn.deterministic = True\n warnings.warn(\n \"You have chosen to seed training. This will turn on the CUDNN deterministic setting, \"\n \"which can slow down your training considerably! You may see unexpected behavior \"\n \"when restarting from checkpoints.\"\n )\n\n # set environment variables for distributed training\n configure_nccl()\n cudnn.benchmark = True\n\n trainer = Trainer(exp, args)\n trainer.train()\n\n\nif __name__ == \"__main__\":\n args = make_parser().parse_args()\n exp = get_exp(args.exp_file, args.name)\n exp.merge(args.opts)\n\n num_gpu = torch.cuda.device_count() if args.devices is None else args.devices\n assert num_gpu <= torch.cuda.device_count()\n\n dist_url = \"auto\" if args.dist_url is None else args.dist_url\n launch(\n main, num_gpu, args.num_machine, backend=args.dist_backend,\n dist_url=dist_url, args=(exp, args)\n )\n"
},
{
"path": "detector/YOLOX/tools/trt.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport argparse\nimport os\nimport shutil\nfrom loguru import logger\n\nimport tensorrt as trt\nimport torch\nfrom torch2trt import torch2trt\n\nfrom yolox.exp import get_exp\n\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"YOLOX ncnn deploy\")\n parser.add_argument(\"-expn\", \"--experiment-name\", type=str, default=None)\n parser.add_argument(\"-n\", \"--name\", type=str, default=None, help=\"model name\")\n\n parser.add_argument(\n \"-f\",\n \"--exp_file\",\n default=None,\n type=str,\n help=\"pls input your expriment description file\",\n )\n parser.add_argument(\"-c\", \"--ckpt\", default=None, type=str, help=\"ckpt path\")\n return parser\n\n\n@logger.catch\ndef main():\n args = make_parser().parse_args()\n exp = get_exp(args.exp_file, args.name)\n if not args.experiment_name:\n args.experiment_name = exp.exp_name\n\n model = exp.get_model()\n file_name = os.path.join(exp.output_dir, args.experiment_name)\n os.makedirs(file_name, exist_ok=True)\n if args.ckpt is None:\n ckpt_file = os.path.join(file_name, \"best_ckpt.pth.tar\")\n else:\n ckpt_file = args.ckpt\n\n ckpt = torch.load(ckpt_file, map_location=\"cpu\")\n # load the model state dict\n\n model.load_state_dict(ckpt[\"model\"])\n logger.info(\"loaded checkpoint done.\")\n model.eval()\n model.cuda()\n model.head.decode_in_inference = False\n x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda()\n model_trt = torch2trt(\n model,\n [x],\n fp16_mode=True,\n log_level=trt.Logger.INFO,\n max_workspace_size=(1 << 32),\n )\n torch.save(model_trt.state_dict(), os.path.join(file_name, 'model_trt.pth'))\n logger.info(\"Converted TensorRT model done.\")\n engine_file = os.path.join(file_name, 'model_trt.engine')\n engine_file_demo = os.path.join('demo', 'TensorRT', 'cpp', 'model_trt.engine')\n with open(engine_file, 'wb') as f:\n f.write(model_trt.engine.serialize())\n\n shutil.copyfile(engine_file, engine_file_demo)\n\n logger.info(\"Converted TensorRT model engine file is saved for C++ inference.\")\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "detector/YOLOX/yolox/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n\nfrom .utils import configure_module\n\nconfigure_module()\n\n__version__ = \"0.1.0\"\n"
},
{
"path": "detector/YOLOX/yolox/core/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nfrom .launch import launch\nfrom .trainer import Trainer\n"
},
{
"path": "detector/YOLOX/yolox/core/launch.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Code are based on\n# https://github.com/facebookresearch/detectron2/blob/master/detectron2/engine/launch.py\n# Copyright (c) Facebook, Inc. and its affiliates.\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nfrom loguru import logger\n\nimport torch\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\n\nimport yolox.utils.dist as comm\n\n__all__ = [\"launch\"]\n\n\ndef _find_free_port():\n \"\"\"\n Find an available port of current machine / node.\n \"\"\"\n import socket\n\n sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n # Binding to port 0 will cause the OS to find an available port for us\n sock.bind((\"\", 0))\n port = sock.getsockname()[1]\n sock.close()\n # NOTE: there is still a chance the port could be taken by other processes.\n return port\n\n\ndef launch(\n main_func, num_gpus_per_machine, num_machines=1, machine_rank=0,\n backend=\"nccl\", dist_url=None, args=()\n):\n \"\"\"\n Args:\n main_func: a function that will be called by `main_func(*args)`\n num_machines (int): the total number of machines\n machine_rank (int): the rank of this machine (one per machine)\n dist_url (str): url to connect to for distributed training, including protocol\n e.g. \"tcp://127.0.0.1:8686\".\n Can be set to auto to automatically select a free port on localhost\n args (tuple): arguments passed to main_func\n \"\"\"\n world_size = num_machines * num_gpus_per_machine\n if world_size > 1:\n # https://github.com/pytorch/pytorch/pull/14391\n # TODO prctl in spawned processes\n\n if dist_url == \"auto\":\n assert num_machines == 1, \"dist_url=auto cannot work with distributed training.\"\n port = _find_free_port()\n dist_url = f\"tcp://127.0.0.1:{port}\"\n\n mp.spawn(\n _distributed_worker,\n nprocs=num_gpus_per_machine,\n args=(\n main_func, world_size, num_gpus_per_machine,\n machine_rank, backend, dist_url, args\n ),\n daemon=False,\n )\n else:\n main_func(*args)\n\n\ndef _distributed_worker(\n local_rank, main_func, world_size, num_gpus_per_machine,\n machine_rank, backend, dist_url, args\n):\n assert torch.cuda.is_available(), \"cuda is not available. Please check your installation.\"\n global_rank = machine_rank * num_gpus_per_machine + local_rank\n logger.info(\"Rank {} initialization finished.\".format(global_rank))\n try:\n dist.init_process_group(\n backend=backend,\n init_method=dist_url,\n world_size=world_size,\n rank=global_rank,\n )\n except Exception:\n logger.error(\"Process group URL: {}\".format(dist_url))\n raise\n # synchronize is needed here to prevent a possible timeout after calling init_process_group\n # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172\n comm.synchronize()\n\n assert num_gpus_per_machine <= torch.cuda.device_count()\n torch.cuda.set_device(local_rank)\n\n # Setup the local process group (which contains ranks within the same machine)\n assert comm._LOCAL_PROCESS_GROUP is None\n num_machines = world_size // num_gpus_per_machine\n for i in range(num_machines):\n ranks_on_i = list(range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine))\n pg = dist.new_group(ranks_on_i)\n if i == machine_rank:\n comm._LOCAL_PROCESS_GROUP = pg\n\n main_func(*args)\n"
},
{
"path": "detector/YOLOX/yolox/core/trainer.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport datetime\nimport os\nimport time\nfrom loguru import logger\n\nimport apex\nimport torch\nfrom apex import amp\nfrom torch.utils.tensorboard import SummaryWriter\n\nfrom yolox.data import DataPrefetcher\nfrom yolox.utils import (\n MeterBuffer,\n ModelEMA,\n all_reduce_norm,\n get_local_rank,\n get_model_info,\n get_rank,\n get_world_size,\n gpu_mem_usage,\n load_ckpt,\n occumpy_mem,\n save_checkpoint,\n setup_logger,\n synchronize\n)\n\n\nclass Trainer:\n\n def __init__(self, exp, args):\n # init function only defines some basic attr, other attrs like model, optimizer are built in\n # before_train methods.\n self.exp = exp\n self.args = args\n\n # training related attr\n self.max_epoch = exp.max_epoch\n self.amp_training = args.fp16\n self.is_distributed = get_world_size() > 1\n self.rank = get_rank()\n self.local_rank = get_local_rank()\n self.device = \"cuda:{}\".format(self.local_rank)\n self.use_model_ema = exp.ema\n\n # data/dataloader related attr\n self.data_type = torch.float16 if args.fp16 else torch.float32\n self.input_size = exp.input_size\n self.best_ap = 0\n\n # metric record\n self.meter = MeterBuffer(window_size=exp.print_interval)\n self.file_name = os.path.join(exp.output_dir, args.experiment_name)\n\n if self.rank == 0 and os.path.exists(\"./\" + args.experiment_name + \"ip_add.txt\"):\n os.remove(\"./\" + args.experiment_name + \"ip_add.txt\")\n\n if self.rank == 0:\n os.makedirs(self.file_name, exist_ok=True)\n\n setup_logger(self.file_name, distributed_rank=self.rank, filename=\"train_log.txt\", mode=\"a\")\n\n def train(self):\n self.before_train()\n try:\n self.train_in_epoch()\n except Exception:\n raise\n finally:\n self.after_train()\n\n def train_in_epoch(self):\n for self.epoch in range(self.start_epoch, self.max_epoch):\n self.before_epoch()\n self.train_in_iter()\n self.after_epoch()\n\n def train_in_iter(self):\n for self.iter in range(self.max_iter):\n self.before_iter()\n self.train_one_iter()\n self.after_iter()\n\n def train_one_iter(self):\n iter_start_time = time.time()\n\n inps, targets = self.prefetcher.next()\n inps = inps.to(self.data_type)\n targets = targets.to(self.data_type)\n targets.requires_grad = False\n data_end_time = time.time()\n\n outputs = self.model(inps, targets)\n loss = outputs[\"total_loss\"]\n\n self.optimizer.zero_grad()\n if self.amp_training:\n with amp.scale_loss(loss, self.optimizer) as scaled_loss:\n scaled_loss.backward()\n else:\n loss.backward()\n self.optimizer.step()\n\n if self.use_model_ema:\n self.ema_model.update(self.model)\n\n lr = self.lr_scheduler.update_lr(self.progress_in_iter + 1)\n for param_group in self.optimizer.param_groups:\n param_group[\"lr\"] = lr\n\n iter_end_time = time.time()\n self.meter.update(\n iter_time=iter_end_time - iter_start_time,\n data_time=data_end_time - iter_start_time,\n lr=lr,\n **outputs,\n )\n\n def before_train(self):\n logger.info(\"args: {}\".format(self.args))\n logger.info(\"exp value:\\n{}\".format(self.exp))\n\n # model related init\n torch.cuda.set_device(self.local_rank)\n model = self.exp.get_model()\n logger.info(\"Model Summary: {}\".format(get_model_info(model, self.exp.test_size)))\n model.to(self.device)\n\n # solver related init\n self.optimizer = self.exp.get_optimizer(self.args.batch_size)\n\n if self.amp_training:\n model, optimizer = amp.initialize(model, self.optimizer, opt_level=\"O1\")\n\n # value of epoch will be set in `resume_train`\n model = self.resume_train(model)\n\n # data related init\n self.no_aug = self.start_epoch >= self.max_epoch - self.exp.no_aug_epochs\n self.train_loader = self.exp.get_data_loader(\n batch_size=self.args.batch_size,\n is_distributed=self.is_distributed,\n no_aug=self.no_aug\n )\n logger.info(\"init prefetcher, this might take one minute or less...\")\n self.prefetcher = DataPrefetcher(self.train_loader)\n # max_iter means iters per epoch\n self.max_iter = len(self.train_loader)\n\n self.lr_scheduler = self.exp.get_lr_scheduler(\n self.exp.basic_lr_per_img * self.args.batch_size, self.max_iter\n )\n if self.args.occumpy:\n occumpy_mem(self.local_rank)\n\n if self.is_distributed:\n model = apex.parallel.DistributedDataParallel(model)\n # from torch.nn.parallel import DistributedDataParallel as DDP\n # model = DDP(model, device_ids=[self.local_rank], broadcast_buffers=False)\n\n if self.use_model_ema:\n self.ema_model = ModelEMA(model, 0.9998)\n self.ema_model.updates = self.max_iter * self.start_epoch\n\n self.model = model\n self.model.train()\n\n self.evaluator = self.exp.get_evaluator(\n batch_size=self.args.batch_size, is_distributed=self.is_distributed\n )\n # Tensorboard logger\n if self.rank == 0:\n self.tblogger = SummaryWriter(self.file_name)\n\n logger.info(\"Training start...\")\n logger.info(\"\\n{}\".format(model))\n\n def after_train(self):\n logger.info(\n \"Training of experiment is done and the best AP is {:.2f}\".format(self.best_ap * 100)\n )\n\n def before_epoch(self):\n logger.info(\"---> start train epoch{}\".format(self.epoch + 1))\n\n if self.epoch + 1 == self.max_epoch - self.exp.no_aug_epochs or self.no_aug:\n logger.info(\"--->No mosaic aug now!\")\n self.train_loader.close_mosaic()\n logger.info(\"--->Add additional L1 loss now!\")\n if self.is_distributed:\n self.model.module.head.use_l1 = True\n else:\n self.model.head.use_l1 = True\n self.exp.eval_interval = 1\n if not self.no_aug:\n self.save_ckpt(ckpt_name=\"last_mosaic_epoch\")\n\n def after_epoch(self):\n if self.use_model_ema:\n self.ema_model.update_attr(self.model)\n\n self.save_ckpt(ckpt_name=\"latest\")\n\n if (self.epoch + 1) % self.exp.eval_interval == 0:\n all_reduce_norm(self.model)\n self.evaluate_and_save_model()\n\n def before_iter(self):\n pass\n\n def after_iter(self):\n \"\"\"\n `after_iter` contains two parts of logic:\n * log information\n * reset setting of resize\n \"\"\"\n # log needed information\n if (self.iter + 1) % self.exp.print_interval == 0:\n # TODO check ETA logic\n left_iters = self.max_iter * self.max_epoch - (self.progress_in_iter + 1)\n eta_seconds = self.meter[\"iter_time\"].global_avg * left_iters\n eta_str = \"ETA: {}\".format(datetime.timedelta(seconds=int(eta_seconds)))\n\n progress_str = \"epoch: {}/{}, iter: {}/{}\".format(\n self.epoch + 1, self.max_epoch, self.iter + 1, self.max_iter\n )\n loss_meter = self.meter.get_filtered_meter(\"loss\")\n loss_str = \", \".join([\"{}: {:.1f}\".format(k, v.latest) for k, v in loss_meter.items()])\n\n time_meter = self.meter.get_filtered_meter(\"time\")\n time_str = \", \".join([\"{}: {:.3f}s\".format(k, v.avg) for k, v in time_meter.items()])\n\n logger.info(\n \"{}, mem: {:.0f}Mb, {}, {}, lr: {:.3e}\".format(\n progress_str,\n gpu_mem_usage(),\n time_str,\n loss_str,\n self.meter[\"lr\"].latest,\n )\n + (\", size: {:d}, {}\".format(self.input_size[0], eta_str))\n )\n self.meter.clear_meters()\n\n # random resizing\n if self.exp.random_size is not None and (self.progress_in_iter + 1) % 10 == 0:\n self.input_size = self.exp.random_resize(\n self.train_loader, self.epoch, self.rank, self.is_distributed\n )\n\n @property\n def progress_in_iter(self):\n return self.epoch * self.max_iter + self.iter\n\n def resume_train(self, model):\n if self.args.resume:\n logger.info(\"resume training\")\n if self.args.ckpt is None:\n ckpt_file = os.path.join(self.file_name, \"latest\" + \"_ckpt.pth.tar\")\n else:\n ckpt_file = self.args.ckpt\n\n ckpt = torch.load(ckpt_file, map_location=self.device)\n # resume the model/optimizer state dict\n model.load_state_dict(ckpt[\"model\"])\n self.optimizer.load_state_dict(ckpt[\"optimizer\"])\n # resume the training states variables\n if self.amp_training and \"amp\" in ckpt:\n amp.load_state_dict(ckpt[\"amp\"])\n start_epoch = (\n self.args.start_epoch - 1\n if self.args.start_epoch is not None\n else ckpt[\"start_epoch\"]\n )\n self.start_epoch = start_epoch\n logger.info(\"loaded checkpoint '{}' (epoch {})\".format(self.args.resume, self.start_epoch)) # noqa\n else:\n if self.args.ckpt is not None:\n logger.info(\"loading checkpoint for fine tuning\")\n ckpt_file = self.args.ckpt\n ckpt = torch.load(ckpt_file, map_location=self.device)[\"model\"]\n model = load_ckpt(model, ckpt)\n self.start_epoch = 0\n\n return model\n\n def evaluate_and_save_model(self):\n evalmodel = self.ema_model.ema if self.use_model_ema else self.model\n ap50_95, ap50, summary = self.exp.eval(evalmodel, self.evaluator, self.is_distributed)\n self.model.train()\n if self.rank == 0:\n self.tblogger.add_scalar(\"val/COCOAP50\", ap50, self.epoch + 1)\n self.tblogger.add_scalar(\"val/COCOAP50_95\", ap50_95, self.epoch + 1)\n logger.info(\"\\n\" + summary)\n synchronize()\n\n self.save_ckpt(\"last_epoch\", ap50_95 > self.best_ap)\n self.best_ap = max(self.best_ap, ap50_95)\n\n def save_ckpt(self, ckpt_name, update_best_ckpt=False):\n if self.rank == 0:\n save_model = self.ema_model.ema if self.use_model_ema else self.model\n logger.info(\"Save weights to {}\".format(self.file_name))\n ckpt_state = {\n \"start_epoch\": self.epoch + 1,\n \"model\": save_model.state_dict(),\n \"optimizer\": self.optimizer.state_dict(),\n }\n if self.amp_training:\n # save amp state according to\n # https://nvidia.github.io/apex/amp.html#checkpointing\n ckpt_state[\"amp\"] = amp.state_dict()\n save_checkpoint(\n ckpt_state,\n update_best_ckpt,\n self.file_name,\n ckpt_name,\n )\n"
},
{
"path": "detector/YOLOX/yolox/data/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nfrom .data_augment import TrainTransform, ValTransform\nfrom .data_prefetcher import DataPrefetcher\nfrom .dataloading import DataLoader, get_yolox_datadir\nfrom .datasets import *\nfrom .samplers import InfiniteSampler, YoloBatchSampler\n"
},
{
"path": "detector/YOLOX/yolox/data/data_augment.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\"\"\"\nData augmentation functionality. Passed as callable transformations to\nDataset classes.\n\nThe data augmentation procedures were interpreted from @weiliu89's SSD paper\nhttp://arxiv.org/abs/1512.02325\n\"\"\"\n\nimport math\nimport random\n\nimport cv2\nimport numpy as np\n\nimport torch\n\n\ndef augment_hsv(img, hgain=0.015, sgain=0.7, vgain=0.4):\n r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains\n hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))\n dtype = img.dtype # uint8\n\n x = np.arange(0, 256, dtype=np.int16)\n lut_hue = ((x * r[0]) % 180).astype(dtype)\n lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)\n lut_val = np.clip(x * r[2], 0, 255).astype(dtype)\n\n img_hsv = cv2.merge(\n (cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))\n ).astype(dtype)\n cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed\n\n\ndef box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.2):\n # box1(4,n), box2(4,n)\n # Compute candidate boxes which include follwing 5 things:\n # box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio\n w1, h1 = box1[2] - box1[0], box1[3] - box1[1]\n w2, h2 = box2[2] - box2[0], box2[3] - box2[1]\n ar = np.maximum(w2 / (h2 + 1e-16), h2 / (w2 + 1e-16)) # aspect ratio\n return (\n (w2 > wh_thr)\n & (h2 > wh_thr)\n & (w2 * h2 / (w1 * h1 + 1e-16) > area_thr)\n & (ar < ar_thr)\n ) # candidates\n\n\ndef random_perspective(\n img, targets=(), degrees=10, translate=0.1, scale=0.1, shear=10, perspective=0.0, border=(0, 0),\n):\n # targets = [cls, xyxy]\n height = img.shape[0] + border[0] * 2 # shape(h,w,c)\n width = img.shape[1] + border[1] * 2\n\n # Center\n C = np.eye(3)\n C[0, 2] = -img.shape[1] / 2 # x translation (pixels)\n C[1, 2] = -img.shape[0] / 2 # y translation (pixels)\n\n # Rotation and Scale\n R = np.eye(3)\n a = random.uniform(-degrees, degrees)\n # a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations\n s = random.uniform(scale[0], scale[1])\n # s = 2 ** random.uniform(-scale, scale)\n R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)\n\n # Shear\n S = np.eye(3)\n S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)\n S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)\n\n # Translation\n T = np.eye(3)\n T[0, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * width) # x translation (pixels)\n T[1, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * height) # y translation (pixels)\n\n # Combined rotation matrix\n M = T @ S @ R @ C # order of operations (right to left) is IMPORTANT\n\n ###########################\n # For Aug out of Mosaic\n # s = 1.\n # M = np.eye(3)\n ###########################\n\n if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed\n if perspective:\n img = cv2.warpPerspective(img, M, dsize=(width, height), borderValue=(114, 114, 114))\n else: # affine\n img = cv2.warpAffine(img, M[:2], dsize=(width, height), borderValue=(114, 114, 114))\n\n # Transform label coordinates\n n = len(targets)\n if n:\n # warp points\n xy = np.ones((n * 4, 3))\n xy[:, :2] = targets[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1\n xy = xy @ M.T # transform\n if perspective:\n xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) # rescale\n else: # affine\n xy = xy[:, :2].reshape(n, 8)\n\n # create new boxes\n x = xy[:, [0, 2, 4, 6]]\n y = xy[:, [1, 3, 5, 7]]\n xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T\n\n # clip boxes\n xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)\n xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)\n\n # filter candidates\n i = box_candidates(box1=targets[:, :4].T * s, box2=xy.T)\n targets = targets[i]\n targets[:, :4] = xy[i]\n\n return img, targets\n\n\ndef _distort(image):\n def _convert(image, alpha=1, beta=0):\n tmp = image.astype(float) * alpha + beta\n tmp[tmp < 0] = 0\n tmp[tmp > 255] = 255\n image[:] = tmp\n\n image = image.copy()\n\n if random.randrange(2):\n _convert(image, beta=random.uniform(-32, 32))\n\n if random.randrange(2):\n _convert(image, alpha=random.uniform(0.5, 1.5))\n\n image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)\n\n if random.randrange(2):\n tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)\n tmp %= 180\n image[:, :, 0] = tmp\n\n if random.randrange(2):\n _convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))\n\n image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)\n\n return image\n\n\ndef _mirror(image, boxes):\n _, width, _ = image.shape\n if random.randrange(2):\n image = image[:, ::-1]\n boxes = boxes.copy()\n boxes[:, 0::2] = width - boxes[:, 2::-2]\n return image, boxes\n\n\ndef preproc(image, input_size, mean, std, swap=(2, 0, 1)):\n if len(image.shape) == 3:\n padded_img = np.ones((input_size[0], input_size[1], 3)) * 114.0\n else:\n padded_img = np.ones(input_size) * 114.0\n img = np.array(image)\n r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])\n resized_img = cv2.resize(\n img, (int(img.shape[1] * r), int(img.shape[0] * r)), interpolation=cv2.INTER_LINEAR\n ).astype(np.float32)\n padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img\n image = padded_img\n\n image = image.astype(np.float32)\n image = image[:, :, ::-1]\n image /= 255.0\n if mean is not None:\n image -= mean\n if std is not None:\n image /= std\n image = image.transpose(swap)\n image = np.ascontiguousarray(image, dtype=np.float32)\n return image, r\n\n\nclass TrainTransform:\n def __init__(self, p=0.5, rgb_means=None, std=None, max_labels=50):\n self.means = rgb_means\n self.std = std\n self.p = p\n self.max_labels = max_labels\n\n def __call__(self, image, targets, input_dim):\n boxes = targets[:, :4].copy()\n labels = targets[:, 4].copy()\n if targets.shape[1] > 5:\n mixup = True\n ratios = targets[:, -1].copy()\n ratios_o = targets[:, -1].copy()\n else:\n mixup = False\n ratios = None\n ratios_o = None\n lshape = 6 if mixup else 5\n if len(boxes) == 0:\n targets = np.zeros((self.max_labels, lshape), dtype=np.float32)\n image, r_o = preproc(image, input_dim, self.means, self.std)\n image = np.ascontiguousarray(image, dtype=np.float32)\n return image, targets\n\n image_o = image.copy()\n targets_o = targets.copy()\n height_o, width_o, _ = image_o.shape\n boxes_o = targets_o[:, :4]\n labels_o = targets_o[:, 4]\n # bbox_o: [xyxy] to [c_x,c_y,w,h]\n b_x_o = (boxes_o[:, 2] + boxes_o[:, 0]) * 0.5\n b_y_o = (boxes_o[:, 3] + boxes_o[:, 1]) * 0.5\n b_w_o = (boxes_o[:, 2] - boxes_o[:, 0]) * 1.0\n b_h_o = (boxes_o[:, 3] - boxes_o[:, 1]) * 1.0\n boxes_o[:, 0] = b_x_o\n boxes_o[:, 1] = b_y_o\n boxes_o[:, 2] = b_w_o\n boxes_o[:, 3] = b_h_o\n\n image_t = _distort(image)\n image_t, boxes = _mirror(image_t, boxes)\n height, width, _ = image_t.shape\n image_t, r_ = preproc(image_t, input_dim, self.means, self.std)\n boxes = boxes.copy()\n # boxes [xyxy] 2 [cx,cy,w,h]\n b_x = (boxes[:, 2] + boxes[:, 0]) * 0.5\n b_y = (boxes[:, 3] + boxes[:, 1]) * 0.5\n b_w = (boxes[:, 2] - boxes[:, 0]) * 1.0\n b_h = (boxes[:, 3] - boxes[:, 1]) * 1.0\n boxes[:, 0] = b_x\n boxes[:, 1] = b_y\n boxes[:, 2] = b_w\n boxes[:, 3] = b_h\n\n boxes *= r_\n\n mask_b = np.minimum(boxes[:, 2], boxes[:, 3]) > 8\n boxes_t = boxes[mask_b]\n labels_t = labels[mask_b].copy()\n if mixup:\n ratios_t = ratios[mask_b].copy()\n\n if len(boxes_t) == 0:\n image_t, r_o = preproc(image_o, input_dim, self.means, self.std)\n boxes_o *= r_o\n boxes_t = boxes_o\n labels_t = labels_o\n ratios_t = ratios_o\n\n labels_t = np.expand_dims(labels_t, 1)\n if mixup:\n ratios_t = np.expand_dims(ratios_t, 1)\n targets_t = np.hstack((labels_t, boxes_t, ratios_t))\n else:\n targets_t = np.hstack((labels_t, boxes_t))\n padded_labels = np.zeros((self.max_labels, lshape))\n padded_labels[range(len(targets_t))[: self.max_labels]] = targets_t[\n : self.max_labels\n ]\n padded_labels = np.ascontiguousarray(padded_labels, dtype=np.float32)\n image_t = np.ascontiguousarray(image_t, dtype=np.float32)\n return image_t, padded_labels\n\n\nclass ValTransform:\n \"\"\"\n Defines the transformations that should be applied to test PIL image\n for input into the network\n\n dimension -> tensorize -> color adj\n\n Arguments:\n resize (int): input dimension to SSD\n rgb_means ((int,int,int)): average RGB of the dataset\n (104,117,123)\n swap ((int,int,int)): final order of channels\n\n Returns:\n transform (transform) : callable transform to be applied to test/val\n data\n \"\"\"\n\n def __init__(self, rgb_means=None, std=None, swap=(2, 0, 1)):\n self.means = rgb_means\n self.swap = swap\n self.std = std\n\n # assume input is cv2 img for now\n def __call__(self, img, res, input_size):\n img, _ = preproc(img, input_size, self.means, self.std, self.swap)\n return torch.from_numpy(img), torch.zeros(1, 5)\n"
},
{
"path": "detector/YOLOX/yolox/data/data_prefetcher.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport random\n\nimport torch\nimport torch.distributed as dist\n\nfrom ..utils import synchronize\n\n\nclass DataPrefetcher:\n \"\"\"\n DataPrefetcher is inspired by code of following file:\n https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py\n It could speedup your pytorch dataloader. For more information, please check\n https://github.com/NVIDIA/apex/issues/304#issuecomment-493562789.\n \"\"\"\n\n def __init__(self, loader):\n self.loader = iter(loader)\n self.stream = torch.cuda.Stream()\n self.input_cuda = self._input_cuda_for_image\n self.record_stream = DataPrefetcher._record_stream_for_image\n self.preload()\n\n def preload(self):\n try:\n self.next_input, self.next_target, _, _ = next(self.loader)\n except StopIteration:\n self.next_input = None\n self.next_target = None\n return\n\n with torch.cuda.stream(self.stream):\n self.input_cuda()\n self.next_target = self.next_target.cuda(non_blocking=True)\n\n def next(self):\n torch.cuda.current_stream().wait_stream(self.stream)\n input = self.next_input\n target = self.next_target\n if input is not None:\n self.record_stream(input)\n if target is not None:\n target.record_stream(torch.cuda.current_stream())\n self.preload()\n return input, target\n\n def _input_cuda_for_image(self):\n self.next_input = self.next_input.cuda(non_blocking=True)\n\n @staticmethod\n def _record_stream_for_image(input):\n input.record_stream(torch.cuda.current_stream())\n\n\ndef random_resize(data_loader, exp, epoch, rank, is_distributed):\n tensor = torch.LongTensor(1).cuda()\n if is_distributed:\n synchronize()\n\n if rank == 0:\n if epoch > exp.max_epoch - 10:\n size = exp.input_size\n else:\n size = random.randint(*exp.random_size)\n size = int(32 * size)\n tensor.fill_(size)\n\n if is_distributed:\n synchronize()\n dist.broadcast(tensor, 0)\n\n input_size = data_loader.change_input_dim(multiple=tensor.item(), random_range=None)\n return input_size\n"
},
{
"path": "detector/YOLOX/yolox/data/dataloading.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\nimport random\n\nimport torch\nfrom torch.utils.data.dataloader import DataLoader as torchDataLoader\nfrom torch.utils.data.dataloader import default_collate\n\nfrom .samplers import YoloBatchSampler\n\n\ndef get_yolox_datadir():\n \"\"\"\n get dataset dir of YOLOX. If environment variable named `YOLOX_DATADIR` is set,\n this function will return value of the environment variable. Otherwise, use data\n \"\"\"\n yolox_datadir = os.getenv(\"YOLOX_DATADIR\", None)\n if yolox_datadir is None:\n import yolox\n yolox_path = os.path.dirname(os.path.dirname(yolox.__file__))\n yolox_datadir = os.path.join(yolox_path, \"datasets\")\n return yolox_datadir\n\n\nclass DataLoader(torchDataLoader):\n \"\"\"\n Lightnet dataloader that enables on the fly resizing of the images.\n See :class:`torch.utils.data.DataLoader` for more information on the arguments.\n Check more on the following website:\n https://gitlab.com/EAVISE/lightnet/-/blob/master/lightnet/data/_dataloading.py\n\n Note:\n This dataloader only works with :class:`lightnet.data.Dataset` based datasets.\n\n Example:\n >>> class CustomSet(ln.data.Dataset):\n ... def __len__(self):\n ... return 4\n ... @ln.data.Dataset.resize_getitem\n ... def __getitem__(self, index):\n ... # Should return (image, anno) but here we return (input_dim,)\n ... return (self.input_dim,)\n >>> dl = ln.data.DataLoader(\n ... CustomSet((200,200)),\n ... batch_size = 2,\n ... collate_fn = ln.data.list_collate # We want the data to be grouped as a list\n ... )\n >>> dl.dataset.input_dim # Default input_dim\n (200, 200)\n >>> for d in dl:\n ... d\n [[(200, 200), (200, 200)]]\n [[(200, 200), (200, 200)]]\n >>> dl.change_input_dim(320, random_range=None)\n (320, 320)\n >>> for d in dl:\n ... d\n [[(320, 320), (320, 320)]]\n [[(320, 320), (320, 320)]]\n >>> dl.change_input_dim((480, 320), random_range=None)\n (480, 320)\n >>> for d in dl:\n ... d\n [[(480, 320), (480, 320)]]\n [[(480, 320), (480, 320)]]\n \"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self.__initialized = False\n shuffle = False\n batch_sampler = None\n if len(args) > 5:\n shuffle = args[2]\n sampler = args[3]\n batch_sampler = args[4]\n elif len(args) > 4:\n shuffle = args[2]\n sampler = args[3]\n if \"batch_sampler\" in kwargs:\n batch_sampler = kwargs[\"batch_sampler\"]\n elif len(args) > 3:\n shuffle = args[2]\n if \"sampler\" in kwargs:\n sampler = kwargs[\"sampler\"]\n if \"batch_sampler\" in kwargs:\n batch_sampler = kwargs[\"batch_sampler\"]\n else:\n if \"shuffle\" in kwargs:\n shuffle = kwargs[\"shuffle\"]\n if \"sampler\" in kwargs:\n sampler = kwargs[\"sampler\"]\n if \"batch_sampler\" in kwargs:\n batch_sampler = kwargs[\"batch_sampler\"]\n\n # Use custom BatchSampler\n if batch_sampler is None:\n if sampler is None:\n if shuffle:\n sampler = torch.utils.data.sampler.RandomSampler(self.dataset)\n # sampler = torch.utils.data.DistributedSampler(self.dataset)\n else:\n sampler = torch.utils.data.sampler.SequentialSampler(self.dataset)\n batch_sampler = YoloBatchSampler(\n sampler,\n self.batch_size,\n self.drop_last,\n input_dimension=self.dataset.input_dim,\n )\n # batch_sampler = IterationBasedBatchSampler(batch_sampler, num_iterations =\n\n self.batch_sampler = batch_sampler\n\n self.__initialized = True\n\n def close_mosaic(self):\n self.batch_sampler.mosaic = False\n\n def change_input_dim(self, multiple=32, random_range=(10, 19)):\n \"\"\" This function will compute a new size and update it on the next mini_batch.\n\n Args:\n multiple (int or tuple, optional): values to multiply the randomly generated range by.\n Default **32**\n random_range (tuple, optional): This (min, max) tuple sets the range\n for the randomisation; Default **(10, 19)**\n\n Return:\n tuple: width, height tuple with new dimension\n\n Note:\n The new size is generated as follows: |br|\n First we compute a random integer inside ``[random_range]``.\n We then multiply that number with the ``multiple`` argument,\n which gives our final new input size. |br|\n If ``multiple`` is an integer we generate a square size. If you give a tuple\n of **(width, height)**, the size is computed\n as :math:`rng * multiple[0], rng * multiple[1]`.\n\n Note:\n You can set the ``random_range`` argument to **None** to set\n an exact size of multiply. |br|\n See the example above for how this works.\n \"\"\"\n if random_range is None:\n size = 1\n else:\n size = random.randint(*random_range)\n\n if isinstance(multiple, int):\n size = (size * multiple, size * multiple)\n else:\n size = (size * multiple[0], size * multiple[1])\n\n self.batch_sampler.new_input_dim = size\n\n return size\n\n\ndef list_collate(batch):\n \"\"\"\n Function that collates lists or tuples together into one list (of lists/tuples).\n Use this as the collate function in a Dataloader, if you want to have a list of\n items as an output, as opposed to tensors (eg. Brambox.boxes).\n \"\"\"\n items = list(zip(*batch))\n\n for i in range(len(items)):\n if isinstance(items[i][0], (list, tuple)):\n items[i] = list(items[i])\n else:\n items[i] = default_collate(items[i])\n\n return items\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nfrom .coco import COCODataset\nfrom .coco_classes import COCO_CLASSES\nfrom .datasets_wrapper import ConcatDataset, Dataset, MixConcatDataset\n#from .mosaicdetection import MosaicDetection\n#from .voc import VOCDetection\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/coco.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nimport cv2\nimport numpy as np\nfrom pycocotools.coco import COCO\n\nfrom ..dataloading import get_yolox_datadir\nfrom .datasets_wrapper import Dataset\n\n\nclass COCODataset(Dataset):\n \"\"\"\n COCO dataset class.\n \"\"\"\n\n def __init__(\n self,\n data_dir=None,\n json_file=\"instances_train2017.json\",\n name=\"train2017\",\n img_size=(416, 416),\n preproc=None,\n ):\n \"\"\"\n COCO dataset initialization. Annotation data are read into memory by COCO API.\n Args:\n data_dir (str): dataset root directory\n json_file (str): COCO json file name\n name (str): COCO data name (e.g. 'train2017' or 'val2017')\n img_size (int): target image size after pre-processing\n preproc: data augmentation strategy\n \"\"\"\n super().__init__(img_size)\n if data_dir is None:\n data_dir = os.path.join(get_yolox_datadir(), \"COCO\")\n self.data_dir = data_dir\n self.json_file = json_file\n\n self.coco = COCO(os.path.join(self.data_dir, \"annotations\", self.json_file))\n self.ids = self.coco.getImgIds()\n self.class_ids = sorted(self.coco.getCatIds())\n cats = self.coco.loadCats(self.coco.getCatIds())\n self._classes = tuple([c[\"name\"] for c in cats])\n self.name = name\n self.img_size = img_size\n self.preproc = preproc\n\n def __len__(self):\n return len(self.ids)\n\n def load_anno(self, index):\n id_ = self.ids[index]\n anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=False)\n annotations = self.coco.loadAnns(anno_ids)\n\n im_ann = self.coco.loadImgs(id_)[0]\n width = im_ann[\"width\"]\n height = im_ann[\"height\"]\n\n # load labels\n valid_objs = []\n for obj in annotations:\n x1 = np.max((0, obj[\"bbox\"][0]))\n y1 = np.max((0, obj[\"bbox\"][1]))\n x2 = np.min((width - 1, x1 + np.max((0, obj[\"bbox\"][2] - 1))))\n y2 = np.min((height - 1, y1 + np.max((0, obj[\"bbox\"][3] - 1))))\n if obj[\"area\"] > 0 and x2 >= x1 and y2 >= y1:\n obj[\"clean_bbox\"] = [x1, y1, x2, y2]\n valid_objs.append(obj)\n objs = valid_objs\n num_objs = len(objs)\n\n res = np.zeros((num_objs, 5))\n\n for ix, obj in enumerate(objs):\n cls = self.class_ids.index(obj[\"category_id\"])\n res[ix, 0:4] = obj[\"clean_bbox\"]\n res[ix, 4] = cls\n\n return res\n\n def pull_item(self, index):\n id_ = self.ids[index]\n\n im_ann = self.coco.loadImgs(id_)[0]\n width = im_ann[\"width\"]\n height = im_ann[\"height\"]\n\n # load image and preprocess\n img_file = os.path.join(\n self.data_dir, self.name, \"{:012}\".format(id_) + \".jpg\"\n )\n\n img = cv2.imread(img_file)\n assert img is not None\n\n # load anno\n res = self.load_anno(index)\n img_info = (height, width)\n\n return img, res, img_info, id_\n\n @Dataset.resize_getitem\n def __getitem__(self, index):\n \"\"\"\n One image / label pair for the given index is picked up and pre-processed.\n\n Args:\n index (int): data index\n\n Returns:\n img (numpy.ndarray): pre-processed image\n padded_labels (torch.Tensor): pre-processed label data.\n The shape is :math:`[max_labels, 5]`.\n each label consists of [class, xc, yc, w, h]:\n class (float): class index.\n xc, yc (float) : center of bbox whose values range from 0 to 1.\n w, h (float) : size of bbox whose values range from 0 to 1.\n info_img : tuple of h, w, nh, nw, dx, dy.\n h, w (int): original shape of the image\n nh, nw (int): shape of the resized image without padding\n dx, dy (int): pad size\n img_id (int): same as the input index. Used for evaluation.\n \"\"\"\n img, res, img_info, img_id = self.pull_item(index)\n\n if self.preproc is not None:\n img, target = self.preproc(img, res, self.input_dim)\n return img, target, img_info, img_id\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/coco_classes.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nCOCO_CLASSES = (\n \"person\",\n \"bicycle\",\n \"car\",\n \"motorcycle\",\n \"airplane\",\n \"bus\",\n \"train\",\n \"truck\",\n \"boat\",\n \"traffic light\",\n \"fire hydrant\",\n \"stop sign\",\n \"parking meter\",\n \"bench\",\n \"bird\",\n \"cat\",\n \"dog\",\n \"horse\",\n \"sheep\",\n \"cow\",\n \"elephant\",\n \"bear\",\n \"zebra\",\n \"giraffe\",\n \"backpack\",\n \"umbrella\",\n \"handbag\",\n \"tie\",\n \"suitcase\",\n \"frisbee\",\n \"skis\",\n \"snowboard\",\n \"sports ball\",\n \"kite\",\n \"baseball bat\",\n \"baseball glove\",\n \"skateboard\",\n \"surfboard\",\n \"tennis racket\",\n \"bottle\",\n \"wine glass\",\n \"cup\",\n \"fork\",\n \"knife\",\n \"spoon\",\n \"bowl\",\n \"banana\",\n \"apple\",\n \"sandwich\",\n \"orange\",\n \"broccoli\",\n \"carrot\",\n \"hot dog\",\n \"pizza\",\n \"donut\",\n \"cake\",\n \"chair\",\n \"couch\",\n \"potted plant\",\n \"bed\",\n \"dining table\",\n \"toilet\",\n \"tv\",\n \"laptop\",\n \"mouse\",\n \"remote\",\n \"keyboard\",\n \"cell phone\",\n \"microwave\",\n \"oven\",\n \"toaster\",\n \"sink\",\n \"refrigerator\",\n \"book\",\n \"clock\",\n \"vase\",\n \"scissors\",\n \"teddy bear\",\n \"hair drier\",\n \"toothbrush\",\n)\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/datasets_wrapper.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport bisect\nfrom functools import wraps\n\nfrom torch.utils.data.dataset import ConcatDataset as torchConcatDataset\nfrom torch.utils.data.dataset import Dataset as torchDataset\n\n\nclass ConcatDataset(torchConcatDataset):\n def __init__(self, datasets):\n super(ConcatDataset, self).__init__(datasets)\n if hasattr(self.datasets[0], \"input_dim\"):\n self._input_dim = self.datasets[0].input_dim\n self.input_dim = self.datasets[0].input_dim\n\n def pull_item(self, idx):\n if idx < 0:\n if -idx > len(self):\n raise ValueError(\n \"absolute value of index should not exceed dataset length\"\n )\n idx = len(self) + idx\n dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)\n if dataset_idx == 0:\n sample_idx = idx\n else:\n sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]\n return self.datasets[dataset_idx].pull_item(sample_idx)\n\n\nclass MixConcatDataset(torchConcatDataset):\n def __init__(self, datasets):\n super(MixConcatDataset, self).__init__(datasets)\n if hasattr(self.datasets[0], \"input_dim\"):\n self._input_dim = self.datasets[0].input_dim\n self.input_dim = self.datasets[0].input_dim\n\n def __getitem__(self, index):\n\n if not isinstance(index, int):\n idx = index[1]\n if idx < 0:\n if -idx > len(self):\n raise ValueError(\n \"absolute value of index should not exceed dataset length\"\n )\n idx = len(self) + idx\n dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)\n if dataset_idx == 0:\n sample_idx = idx\n else:\n sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]\n if not isinstance(index, int):\n index = (index[0], sample_idx, index[2])\n\n return self.datasets[dataset_idx][index]\n\n\nclass Dataset(torchDataset):\n \"\"\" This class is a subclass of the base :class:`torch.utils.data.Dataset`,\n that enables on the fly resizing of the ``input_dim``.\n\n Args:\n input_dimension (tuple): (width,height) tuple with default dimensions of the network\n \"\"\"\n\n def __init__(self, input_dimension, mosaic=True):\n super().__init__()\n self.__input_dim = input_dimension[:2]\n self._mosaic = mosaic\n\n @property\n def input_dim(self):\n \"\"\"\n Dimension that can be used by transforms to set the correct image size, etc.\n This allows transforms to have a single source of truth\n for the input dimension of the network.\n\n Return:\n list: Tuple containing the current width,height\n \"\"\"\n if hasattr(self, \"_input_dim\"):\n return self._input_dim\n return self.__input_dim\n\n @staticmethod\n def resize_getitem(getitem_fn):\n \"\"\"\n Decorator method that needs to be used around the ``__getitem__`` method. |br|\n This decorator enables the on the fly resizing of\n the ``input_dim`` with our :class:`~lightnet.data.DataLoader` class.\n\n Example:\n >>> class CustomSet(ln.data.Dataset):\n ... def __len__(self):\n ... return 10\n ... @ln.data.Dataset.resize_getitem\n ... def __getitem__(self, index):\n ... # Should return (image, anno) but here we return input_dim\n ... return self.input_dim\n >>> data = CustomSet((200,200))\n >>> data[0]\n (200, 200)\n >>> data[(480,320), 0]\n (480, 320)\n \"\"\"\n\n @wraps(getitem_fn)\n def wrapper(self, index):\n if not isinstance(index, int):\n has_dim = True\n self._input_dim = index[0]\n self._mosaic = index[2]\n index = index[1]\n else:\n has_dim = False\n\n ret_val = getitem_fn(self, index)\n\n if has_dim:\n del self._input_dim\n\n return ret_val\n\n return wrapper\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/mosaicdetection.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport random\n\nimport cv2\nimport numpy as np\n\nfrom yolox.utils import adjust_box_anns\n\nfrom ..data_augment import box_candidates, random_perspective\nfrom .datasets_wrapper import Dataset\n\n\ndef get_mosaic_coordinate(mosaic_image, mosaic_index, xc, yc, w, h, input_h, input_w):\n # TODO update doc\n # index0 to top left part of image\n if mosaic_index == 0:\n x1, y1, x2, y2 = max(xc - w, 0), max(yc - h, 0), xc, yc\n small_coord = w - (x2 - x1), h - (y2 - y1), w, h\n # index1 to top right part of image\n elif mosaic_index == 1:\n x1, y1, x2, y2 = xc, max(yc - h, 0), min(xc + w, input_w * 2), yc\n small_coord = 0, h - (y2 - y1), min(w, x2 - x1), h\n # index2 to bottom left part of image\n elif mosaic_index == 2:\n x1, y1, x2, y2 = max(xc - w, 0), yc, xc, min(input_h * 2, yc + h)\n small_coord = w - (x2 - x1), 0, w, min(y2 - y1, h)\n # index2 to bottom right part of image\n elif mosaic_index == 3:\n x1, y1, x2, y2 = xc, yc, min(xc + w, input_w * 2), min(input_h * 2, yc + h) # noqa\n small_coord = 0, 0, min(w, x2 - x1), min(y2 - y1, h)\n return (x1, y1, x2, y2), small_coord\n\n\nclass MosaicDetection(Dataset):\n \"\"\"Detection dataset wrapper that performs mixup for normal dataset.\"\"\"\n\n def __init__(\n self, dataset, img_size, mosaic=True, preproc=None,\n degrees=10.0, translate=0.1, scale=(0.5, 1.5), mscale=(0.5, 1.5),\n shear=2.0, perspective=0.0, enable_mixup=True, *args\n ):\n \"\"\"\n\n Args:\n dataset(Dataset) : Pytorch dataset object.\n img_size (tuple):\n mosaic (bool): enable mosaic augmentation or not.\n preproc (func):\n degrees (float):\n translate (float):\n scale (tuple):\n mscale (tuple):\n shear (float):\n perspective (float):\n enable_mixup (bool):\n *args(tuple) : Additional arguments for mixup random sampler.\n \"\"\"\n super().__init__(img_size, mosaic=mosaic)\n self._dataset = dataset\n self.preproc = preproc\n self.degrees = degrees\n self.translate = translate\n self.scale = scale\n self.shear = shear\n self.perspective = perspective\n self.mixup_scale = mscale\n self.enable_mosaic = mosaic\n self.enable_mixup = enable_mixup\n\n def __len__(self):\n return len(self._dataset)\n\n @Dataset.resize_getitem\n def __getitem__(self, idx):\n if self.enable_mosaic:\n mosaic_labels = []\n input_dim = self._dataset.input_dim\n input_h, input_w = input_dim[0], input_dim[1]\n\n # yc, xc = s, s # mosaic center x, y\n yc = int(random.uniform(0.5 * input_h, 1.5 * input_h))\n xc = int(random.uniform(0.5 * input_w, 1.5 * input_w))\n\n # 3 additional image indices\n indices = [idx] + [random.randint(0, len(self._dataset) - 1) for _ in range(3)]\n\n for i_mosaic, index in enumerate(indices):\n img, _labels, _, _ = self._dataset.pull_item(index)\n h0, w0 = img.shape[:2] # orig hw\n scale = min(1. * input_h / h0, 1. * input_w / w0)\n img = cv2.resize(\n img, (int(w0 * scale), int(h0 * scale)), interpolation=cv2.INTER_LINEAR\n )\n # generate output mosaic image\n (h, w, c) = img.shape[:3]\n if i_mosaic == 0:\n mosaic_img = np.full((input_h * 2, input_w * 2, c), 114, dtype=np.uint8)\n\n # suffix l means large image, while s means small image in mosaic aug.\n (l_x1, l_y1, l_x2, l_y2), (s_x1, s_y1, s_x2, s_y2) = get_mosaic_coordinate(\n mosaic_img, i_mosaic, xc, yc, w, h, input_h, input_w\n )\n\n mosaic_img[l_y1:l_y2, l_x1:l_x2] = img[s_y1:s_y2, s_x1:s_x2]\n padw, padh = l_x1 - s_x1, l_y1 - s_y1\n\n labels = _labels.copy()\n # Normalized xywh to pixel xyxy format\n if _labels.size > 0:\n labels[:, 0] = scale * _labels[:, 0] + padw\n labels[:, 1] = scale * _labels[:, 1] + padh\n labels[:, 2] = scale * _labels[:, 2] + padw\n labels[:, 3] = scale * _labels[:, 3] + padh\n mosaic_labels.append(labels)\n\n if len(mosaic_labels):\n mosaic_labels = np.concatenate(mosaic_labels, 0)\n np.clip(mosaic_labels[:, 0], 0, 2 * input_w, out=mosaic_labels[:, 0])\n np.clip(mosaic_labels[:, 1], 0, 2 * input_h, out=mosaic_labels[:, 1])\n np.clip(mosaic_labels[:, 2], 0, 2 * input_w, out=mosaic_labels[:, 2])\n np.clip(mosaic_labels[:, 3], 0, 2 * input_h, out=mosaic_labels[:, 3])\n\n mosaic_img, mosaic_labels = random_perspective(\n mosaic_img,\n mosaic_labels,\n degrees=self.degrees,\n translate=self.translate,\n scale=self.scale,\n shear=self.shear,\n perspective=self.perspective,\n border=[-input_h // 2, -input_w // 2],\n ) # border to remove\n\n # -----------------------------------------------------------------\n # CopyPaste: https://arxiv.org/abs/2012.07177\n # -----------------------------------------------------------------\n if self.enable_mixup and not len(mosaic_labels) == 0:\n mosaic_img, mosaic_labels = self.mixup(mosaic_img, mosaic_labels, self.input_dim)\n mix_img, padded_labels = self.preproc(mosaic_img, mosaic_labels, self.input_dim)\n img_info = (mix_img.shape[1], mix_img.shape[0])\n\n return mix_img, padded_labels, img_info, int(idx)\n\n else:\n self._dataset._input_dim = self.input_dim\n img, label, img_info, idx = self._dataset.pull_item(idx)\n img, label = self.preproc(img, label, self.input_dim)\n return img, label, img_info, int(idx)\n\n def mixup(self, origin_img, origin_labels, input_dim):\n jit_factor = random.uniform(*self.mixup_scale)\n FLIP = random.uniform(0, 1) > 0.5\n cp_labels = []\n while len(cp_labels) == 0:\n cp_index = random.randint(0, self.__len__() - 1)\n cp_labels = self._dataset.load_anno(cp_index)\n img, cp_labels, _, _ = self._dataset.pull_item(cp_index)\n\n if len(img.shape) == 3:\n cp_img = np.ones((input_dim[0], input_dim[1], 3)) * 114.0\n else:\n cp_img = np.ones(input_dim) * 114.0\n cp_scale_ratio = min(input_dim[0] / img.shape[0], input_dim[1] / img.shape[1])\n resized_img = cv2.resize(\n img,\n (int(img.shape[1] * cp_scale_ratio), int(img.shape[0] * cp_scale_ratio)),\n interpolation=cv2.INTER_LINEAR,\n ).astype(np.float32)\n cp_img[\n : int(img.shape[0] * cp_scale_ratio), : int(img.shape[1] * cp_scale_ratio)\n ] = resized_img\n cp_img = cv2.resize(\n cp_img,\n (int(cp_img.shape[1] * jit_factor), int(cp_img.shape[0] * jit_factor)),\n )\n cp_scale_ratio *= jit_factor\n if FLIP:\n cp_img = cp_img[:, ::-1, :]\n\n origin_h, origin_w = cp_img.shape[:2]\n target_h, target_w = origin_img.shape[:2]\n padded_img = np.zeros(\n (max(origin_h, target_h), max(origin_w, target_w), 3)\n ).astype(np.uint8)\n padded_img[:origin_h, :origin_w] = cp_img\n\n x_offset, y_offset = 0, 0\n if padded_img.shape[0] > target_h:\n y_offset = random.randint(0, padded_img.shape[0] - target_h - 1)\n if padded_img.shape[1] > target_w:\n x_offset = random.randint(0, padded_img.shape[1] - target_w - 1)\n padded_cropped_img = padded_img[\n y_offset: y_offset + target_h, x_offset: x_offset + target_w\n ]\n\n cp_bboxes_origin_np = adjust_box_anns(\n cp_labels[:, :4], cp_scale_ratio, 0, 0, origin_w, origin_h\n )\n if FLIP:\n cp_bboxes_origin_np[:, 0::2] = (\n origin_w - cp_bboxes_origin_np[:, 0::2][:, ::-1]\n )\n cp_bboxes_transformed_np = cp_bboxes_origin_np.copy()\n cp_bboxes_transformed_np[:, 0::2] = np.clip(\n cp_bboxes_transformed_np[:, 0::2] - x_offset, 0, target_w\n )\n cp_bboxes_transformed_np[:, 1::2] = np.clip(\n cp_bboxes_transformed_np[:, 1::2] - y_offset, 0, target_h\n )\n keep_list = box_candidates(cp_bboxes_origin_np.T, cp_bboxes_transformed_np.T, 5)\n\n if keep_list.sum() >= 1.0:\n cls_labels = cp_labels[keep_list, 4:5]\n box_labels = cp_bboxes_transformed_np[keep_list]\n labels = np.hstack((box_labels, cls_labels))\n origin_labels = np.vstack((origin_labels, labels))\n origin_img = origin_img.astype(np.float32)\n origin_img = 0.5 * origin_img + 0.5 * padded_cropped_img.astype(np.float32)\n\n return origin_img.astype(np.uint8), origin_labels\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/voc.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Code are based on\n# https://github.com/fmassa/vision/blob/voc_dataset/torchvision/datasets/voc.py\n# Copyright (c) Francisco Massa.\n# Copyright (c) Ellis Brown, Max deGroot.\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\nimport os.path\nimport pickle\nimport xml.etree.ElementTree as ET\n\nimport cv2\nimport numpy as np\n\nfrom yolox.evaluators.voc_eval import voc_eval\n\nfrom .datasets_wrapper import Dataset\nfrom .voc_classes import VOC_CLASSES\n\n\nclass AnnotationTransform(object):\n\n \"\"\"Transforms a VOC annotation into a Tensor of bbox coords and label index\n Initilized with a dictionary lookup of classnames to indexes\n\n Arguments:\n class_to_ind (dict, optional): dictionary lookup of classnames -> indexes\n (default: alphabetic indexing of VOC's 20 classes)\n keep_difficult (bool, optional): keep difficult instances or not\n (default: False)\n height (int): height\n width (int): width\n \"\"\"\n\n def __init__(self, class_to_ind=None, keep_difficult=True):\n self.class_to_ind = class_to_ind or dict(zip(VOC_CLASSES, range(len(VOC_CLASSES))))\n self.keep_difficult = keep_difficult\n\n def __call__(self, target):\n \"\"\"\n Arguments:\n target (annotation) : the target annotation to be made usable\n will be an ET.Element\n Returns:\n a list containing lists of bounding boxes [bbox coords, class name]\n \"\"\"\n res = np.empty((0, 5))\n for obj in target.iter(\"object\"):\n difficult = int(obj.find(\"difficult\").text) == 1\n if not self.keep_difficult and difficult:\n continue\n name = obj.find(\"name\").text.lower().strip()\n bbox = obj.find(\"bndbox\")\n\n pts = [\"xmin\", \"ymin\", \"xmax\", \"ymax\"]\n bndbox = []\n for i, pt in enumerate(pts):\n cur_pt = int(bbox.find(pt).text) - 1\n # scale height or width\n # cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height\n bndbox.append(cur_pt)\n label_idx = self.class_to_ind[name]\n bndbox.append(label_idx)\n res = np.vstack((res, bndbox)) # [xmin, ymin, xmax, ymax, label_ind]\n # img_id = target.find('filename').text[:-4]\n\n return res # [[xmin, ymin, xmax, ymax, label_ind], ... ]\n\n\nclass VOCDetection(Dataset):\n\n \"\"\"\n VOC Detection Dataset Object\n\n input is image, target is annotation\n\n Args:\n root (string): filepath to VOCdevkit folder.\n image_set (string): imageset to use (eg. 'train', 'val', 'test')\n transform (callable, optional): transformation to perform on the\n input image\n target_transform (callable, optional): transformation to perform on the\n target `annotation`\n (eg: take in caption string, return tensor of word indices)\n dataset_name (string, optional): which dataset to load\n (default: 'VOC2007')\n \"\"\"\n\n def __init__(\n self,\n data_dir,\n image_sets=[('2007', 'trainval'), ('2012', 'trainval')],\n img_size=(416, 416),\n preproc=None,\n target_transform=AnnotationTransform(),\n dataset_name=\"VOC0712\",\n ):\n super().__init__(img_size)\n self.root = data_dir\n self.image_set = image_sets\n self.img_size = img_size\n self.preproc = preproc\n self.target_transform = target_transform\n self.name = dataset_name\n self._annopath = os.path.join(\"%s\", \"Annotations\", \"%s.xml\")\n self._imgpath = os.path.join(\"%s\", \"JPEGImages\", \"%s.jpg\")\n self._classes = VOC_CLASSES\n self.ids = list()\n for (year, name) in image_sets:\n self._year = year\n rootpath = os.path.join(self.root, \"VOC\" + year)\n for line in open(\n os.path.join(rootpath, \"ImageSets\", \"Main\", name + \".txt\")\n ):\n self.ids.append((rootpath, line.strip()))\n\n def __len__(self):\n return len(self.ids)\n\n def load_anno(self, index):\n img_id = self.ids[index]\n target = ET.parse(self._annopath % img_id).getroot()\n if self.target_transform is not None:\n target = self.target_transform(target)\n\n return target\n\n def pull_item(self, index):\n \"\"\"Returns the original image and target at an index for mixup\n\n Note: not using self.__getitem__(), as any transformations passed in\n could mess up this functionality.\n\n Argument:\n index (int): index of img to show\n Return:\n img, target\n \"\"\"\n img_id = self.ids[index]\n img = cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)\n height, width, _ = img.shape\n\n target = self.load_anno(index)\n\n img_info = (width, height)\n\n return img, target, img_info, index\n\n @Dataset.resize_getitem\n def __getitem__(self, index):\n img, target, img_info, img_id = self.pull_item(index)\n\n if self.preproc is not None:\n img, target = self.preproc(img, target, self.input_dim)\n\n return img, target, img_info, img_id\n\n def evaluate_detections(self, all_boxes, output_dir=None):\n \"\"\"\n all_boxes is a list of length number-of-classes.\n Each list element is a list of length number-of-images.\n Each of those list elements is either an empty list []\n or a numpy array of detection.\n\n all_boxes[class][image] = [] or np.array of shape #dets x 5\n \"\"\"\n self._write_voc_results_file(all_boxes)\n IouTh = np.linspace(0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True)\n mAPs = []\n for iou in IouTh:\n mAP = self._do_python_eval(output_dir, iou)\n mAPs.append(mAP)\n\n print(\"--------------------------------------------------------------\")\n print(\"map_5095:\", np.mean(mAPs))\n print(\"map_50:\", mAPs[0])\n print(\"--------------------------------------------------------------\")\n return np.mean(mAPs), mAPs[0]\n\n def _get_voc_results_file_template(self):\n filename = \"comp4_det_test\" + \"_{:s}.txt\"\n filedir = os.path.join(self.root, \"results\", \"VOC\" + self._year, \"Main\")\n if not os.path.exists(filedir):\n os.makedirs(filedir)\n path = os.path.join(filedir, filename)\n return path\n\n def _write_voc_results_file(self, all_boxes):\n for cls_ind, cls in enumerate(VOC_CLASSES):\n cls_ind = cls_ind\n if cls == \"__background__\":\n continue\n print(\"Writing {} VOC results file\".format(cls))\n filename = self._get_voc_results_file_template().format(cls)\n with open(filename, \"wt\") as f:\n for im_ind, index in enumerate(self.ids):\n index = index[1]\n dets = all_boxes[cls_ind][im_ind]\n if dets == []:\n continue\n for k in range(dets.shape[0]):\n f.write(\n \"{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\\n\".format(\n index,\n dets[k, -1],\n dets[k, 0] + 1,\n dets[k, 1] + 1,\n dets[k, 2] + 1,\n dets[k, 3] + 1,\n )\n )\n\n def _do_python_eval(self, output_dir=\"output\", iou=0.5):\n rootpath = os.path.join(self.root, \"VOC\" + self._year)\n name = self.image_set[0][1]\n annopath = os.path.join(rootpath, \"Annotations\", \"{:s}.xml\")\n imagesetfile = os.path.join(rootpath, \"ImageSets\", \"Main\", name + \".txt\")\n cachedir = os.path.join(\n self.root, \"annotations_cache\", \"VOC\" + self._year, name\n )\n if not os.path.exists(cachedir):\n os.makedirs(cachedir)\n aps = []\n # The PASCAL VOC metric changed in 2010\n use_07_metric = True if int(self._year) < 2010 else False\n print(\"Eval IoU : {:.2f}\".format(iou))\n if output_dir is not None and not os.path.isdir(output_dir):\n os.mkdir(output_dir)\n for i, cls in enumerate(VOC_CLASSES):\n\n if cls == \"__background__\":\n continue\n\n filename = self._get_voc_results_file_template().format(cls)\n rec, prec, ap = voc_eval(\n filename,\n annopath,\n imagesetfile,\n cls,\n cachedir,\n ovthresh=iou,\n use_07_metric=use_07_metric,\n )\n aps += [ap]\n if iou == 0.5:\n print(\"AP for {} = {:.4f}\".format(cls, ap))\n if output_dir is not None:\n with open(os.path.join(output_dir, cls + \"_pr.pkl\"), \"wb\") as f:\n pickle.dump({\"rec\": rec, \"prec\": prec, \"ap\": ap}, f)\n if iou == 0.5:\n print(\"Mean AP = {:.4f}\".format(np.mean(aps)))\n print(\"~~~~~~~~\")\n print(\"Results:\")\n for ap in aps:\n print(\"{:.3f}\".format(ap))\n print(\"{:.3f}\".format(np.mean(aps)))\n print(\"~~~~~~~~\")\n print(\"\")\n print(\"--------------------------------------------------------------\")\n print(\"Results computed with the **unofficial** Python eval code.\")\n print(\"Results should be very close to the official MATLAB eval code.\")\n print(\"Recompute with `./tools/reval.py --matlab ...` for your paper.\")\n print(\"-- Thanks, The Management\")\n print(\"--------------------------------------------------------------\")\n\n return np.mean(aps)\n"
},
{
"path": "detector/YOLOX/yolox/data/datasets/voc_classes.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\n# VOC_CLASSES = ( '__background__', # always index 0\nVOC_CLASSES = (\n \"aeroplane\",\n \"bicycle\",\n \"bird\",\n \"boat\",\n \"bottle\",\n \"bus\",\n \"car\",\n \"cat\",\n \"chair\",\n \"cow\",\n \"diningtable\",\n \"dog\",\n \"horse\",\n \"motorbike\",\n \"person\",\n \"pottedplant\",\n \"sheep\",\n \"sofa\",\n \"train\",\n \"tvmonitor\",\n)\n"
},
{
"path": "detector/YOLOX/yolox/data/samplers.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport itertools\nfrom typing import Optional\n\nimport torch\nimport torch.distributed as dist\nfrom torch.utils.data.sampler import BatchSampler as torchBatchSampler\nfrom torch.utils.data.sampler import Sampler\n\n\nclass YoloBatchSampler(torchBatchSampler):\n \"\"\"\n This batch sampler will generate mini-batches of (dim, index) tuples from another sampler.\n It works just like the :class:`torch.utils.data.sampler.BatchSampler`,\n but it will prepend a dimension, whilst ensuring it stays the same across one mini-batch.\n \"\"\"\n\n def __init__(self, *args, input_dimension=None, mosaic=True, **kwargs):\n super().__init__(*args, **kwargs)\n self.input_dim = input_dimension\n self.new_input_dim = None\n self.mosaic = mosaic\n\n def __iter__(self):\n self.__set_input_dim()\n for batch in super().__iter__():\n yield [(self.input_dim, idx, self.mosaic) for idx in batch]\n self.__set_input_dim()\n\n def __set_input_dim(self):\n \"\"\" This function randomly changes the the input dimension of the dataset. \"\"\"\n if self.new_input_dim is not None:\n self.input_dim = (self.new_input_dim[0], self.new_input_dim[1])\n self.new_input_dim = None\n\n\nclass InfiniteSampler(Sampler):\n \"\"\"\n In training, we only care about the \"infinite stream\" of training data.\n So this sampler produces an infinite stream of indices and\n all workers cooperate to correctly shuffle the indices and sample different indices.\n The samplers in each worker effectively produces `indices[worker_id::num_workers]`\n where `indices` is an infinite stream of indices consisting of\n `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)\n or `range(size) + range(size) + ...` (if shuffle is False)\n \"\"\"\n\n def __init__(\n self,\n size: int,\n shuffle: bool = True,\n seed: Optional[int] = 0,\n rank=0,\n world_size=1,\n ):\n \"\"\"\n Args:\n size (int): the total number of data of the underlying dataset to sample from\n shuffle (bool): whether to shuffle the indices or not\n seed (int): the initial seed of the shuffle. Must be the same\n across all workers. If None, will use a random seed shared\n among workers (require synchronization among all workers).\n \"\"\"\n self._size = size\n assert size > 0\n self._shuffle = shuffle\n self._seed = int(seed)\n\n if dist.is_available() and dist.is_initialized():\n self._rank = dist.get_rank()\n self._world_size = dist.get_world_size()\n else:\n self._rank = rank\n self._world_size = world_size\n\n def __iter__(self):\n start = self._rank\n yield from itertools.islice(\n self._infinite_indices(), start, None, self._world_size\n )\n\n def _infinite_indices(self):\n g = torch.Generator()\n g.manual_seed(self._seed)\n while True:\n if self._shuffle:\n yield from torch.randperm(self._size, generator=g)\n else:\n yield from torch.arange(self._size)\n\n def __len__(self):\n return self._size // self._world_size\n"
},
{
"path": "detector/YOLOX/yolox/evaluators/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nfrom .coco_evaluator import COCOEvaluator\nfrom .voc_evaluator import VOCEvaluator\n"
},
{
"path": "detector/YOLOX/yolox/evaluators/coco_evaluator.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport contextlib\nimport io\nimport itertools\nimport json\nimport tempfile\nimport time\nfrom loguru import logger\nfrom tqdm import tqdm\n\nimport torch\n\nfrom yolox.utils import (\n gather,\n is_main_process,\n postprocess,\n synchronize,\n time_synchronized,\n xyxy2xywh\n)\n\n\nclass COCOEvaluator:\n \"\"\"\n COCO AP Evaluation class. All the data in the val2017 dataset are processed\n and evaluated by COCO API.\n \"\"\"\n\n def __init__(\n self, dataloader, img_size, confthre, nmsthre, num_classes, testdev=False\n ):\n \"\"\"\n Args:\n dataloader (Dataloader): evaluate dataloader.\n img_size (int): image size after preprocess. images are resized\n to squares whose shape is (img_size, img_size).\n confthre (float): confidence threshold ranging from 0 to 1, which\n is defined in the config file.\n nmsthre (float): IoU threshold of non-max supression ranging from 0 to 1.\n \"\"\"\n self.dataloader = dataloader\n self.img_size = img_size\n self.confthre = confthre\n self.nmsthre = nmsthre\n self.num_classes = num_classes\n self.testdev = testdev\n\n def evaluate(\n self,\n model,\n distributed=False,\n half=False,\n trt_file=None,\n decoder=None,\n test_size=None,\n ):\n \"\"\"\n COCO average precision (AP) Evaluation. Iterate inference on the test dataset\n and the results are evaluated by COCO API.\n\n NOTE: This function will change training mode to False, please save states if needed.\n\n Args:\n model : model to evaluate.\n\n Returns:\n ap50_95 (float) : COCO AP of IoU=50:95\n ap50 (float) : COCO AP of IoU=50\n summary (sr): summary info of evaluation.\n \"\"\"\n # TODO half to amp_test\n tensor_type = torch.cuda.HalfTensor if half else torch.cuda.FloatTensor\n model = model.eval()\n if half:\n model = model.half()\n ids = []\n data_list = []\n progress_bar = tqdm if is_main_process() else iter\n\n inference_time = 0\n nms_time = 0\n n_samples = len(self.dataloader) - 1\n\n if trt_file is not None:\n from torch2trt import TRTModule\n\n model_trt = TRTModule()\n model_trt.load_state_dict(torch.load(trt_file))\n\n x = torch.ones(1, 3, test_size[0], test_size[1]).cuda()\n model(x)\n model = model_trt\n\n for cur_iter, (imgs, _, info_imgs, ids) in enumerate(\n progress_bar(self.dataloader)\n ):\n with torch.no_grad():\n imgs = imgs.type(tensor_type)\n\n # skip the the last iters since batchsize might be not enough for batch inference\n is_time_record = cur_iter < len(self.dataloader) - 1\n if is_time_record:\n start = time.time()\n\n outputs = model(imgs)\n if decoder is not None:\n outputs = decoder(outputs, dtype=outputs.type())\n\n if is_time_record:\n infer_end = time_synchronized()\n inference_time += infer_end - start\n\n outputs = postprocess(\n outputs, self.num_classes, self.confthre, self.nmsthre\n )\n if is_time_record:\n nms_end = time_synchronized()\n nms_time += nms_end - infer_end\n\n data_list.extend(self.convert_to_coco_format(outputs, info_imgs, ids))\n\n statistics = torch.cuda.FloatTensor([inference_time, nms_time, n_samples])\n if distributed:\n data_list = gather(data_list, dst=0)\n data_list = list(itertools.chain(*data_list))\n torch.distributed.reduce(statistics, dst=0)\n\n eval_results = self.evaluate_prediction(data_list, statistics)\n synchronize()\n return eval_results\n\n def convert_to_coco_format(self, outputs, info_imgs, ids):\n data_list = []\n for (output, img_h, img_w, img_id) in zip(\n outputs, info_imgs[0], info_imgs[1], ids\n ):\n if output is None:\n continue\n output = output.cpu()\n\n bboxes = output[:, 0:4]\n\n # preprocessing: resize\n scale = min(\n self.img_size[0] / float(img_h), self.img_size[1] / float(img_w)\n )\n bboxes /= scale\n bboxes = xyxy2xywh(bboxes)\n\n cls = output[:, 6]\n scores = output[:, 4] * output[:, 5]\n for ind in range(bboxes.shape[0]):\n label = self.dataloader.dataset.class_ids[int(cls[ind])]\n pred_data = {\n \"image_id\": int(img_id),\n \"category_id\": label,\n \"bbox\": bboxes[ind].numpy().tolist(),\n \"score\": scores[ind].numpy().item(),\n \"segmentation\": [],\n } # COCO json format\n data_list.append(pred_data)\n return data_list\n\n def evaluate_prediction(self, data_dict, statistics):\n if not is_main_process():\n return 0, 0, None\n\n logger.info(\"Evaluate in main process...\")\n\n annType = [\"segm\", \"bbox\", \"keypoints\"]\n\n inference_time = statistics[0].item()\n nms_time = statistics[1].item()\n n_samples = statistics[2].item()\n\n a_infer_time = 1000 * inference_time / (n_samples * self.dataloader.batch_size)\n a_nms_time = 1000 * nms_time / (n_samples * self.dataloader.batch_size)\n\n time_info = \", \".join(\n [\n \"Average {} time: {:.2f} ms\".format(k, v)\n for k, v in zip(\n [\"forward\", \"NMS\", \"inference\"],\n [a_infer_time, a_nms_time, (a_infer_time + a_nms_time)],\n )\n ]\n )\n\n info = time_info + \"\\n\"\n\n # Evaluate the Dt (detection) json comparing with the ground truth\n if len(data_dict) > 0:\n cocoGt = self.dataloader.dataset.coco\n # TODO: since pycocotools can't process dict in py36, write data to json file.\n if self.testdev:\n json.dump(data_dict, open(\"./yolox_testdev_2017.json\", \"w\"))\n cocoDt = cocoGt.loadRes(\"./yolox_testdev_2017.json\")\n else:\n _, tmp = tempfile.mkstemp()\n json.dump(data_dict, open(tmp, \"w\"))\n cocoDt = cocoGt.loadRes(tmp)\n try:\n from yolox.layers import COCOeval_opt as COCOeval\n except ImportError:\n from .cocoeval_mr import COCOeval\n\n logger.warning(\"Use standard COCOeval.\")\n\n cocoEval = COCOeval(cocoGt, cocoDt, annType[1])\n cocoEval.evaluate()\n cocoEval.accumulate()\n redirect_string = io.StringIO()\n with contextlib.redirect_stdout(redirect_string):\n cocoEval.summarize()\n info += redirect_string.getvalue()\n return cocoEval.stats[0], cocoEval.stats[1], info\n else:\n return 0, 0, info\n"
},
{
"path": "detector/YOLOX/yolox/evaluators/voc_eval.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Code are based on\n# https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/datasets/voc_eval.py\n# Copyright (c) Bharath Hariharan.\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\nimport pickle\nimport xml.etree.ElementTree as ET\n\nimport numpy as np\n\n\ndef parse_rec(filename):\n \"\"\" Parse a PASCAL VOC xml file \"\"\"\n tree = ET.parse(filename)\n objects = []\n for obj in tree.findall(\"object\"):\n obj_struct = {}\n obj_struct[\"name\"] = obj.find(\"name\").text\n obj_struct[\"pose\"] = obj.find(\"pose\").text\n obj_struct[\"truncated\"] = int(obj.find(\"truncated\").text)\n obj_struct[\"difficult\"] = int(obj.find(\"difficult\").text)\n bbox = obj.find(\"bndbox\")\n obj_struct[\"bbox\"] = [\n int(bbox.find(\"xmin\").text),\n int(bbox.find(\"ymin\").text),\n int(bbox.find(\"xmax\").text),\n int(bbox.find(\"ymax\").text),\n ]\n objects.append(obj_struct)\n\n return objects\n\n\ndef voc_ap(rec, prec, use_07_metric=False):\n \"\"\" ap = voc_ap(rec, prec, [use_07_metric])\n Compute VOC AP given precision and recall.\n If use_07_metric is true, uses the\n VOC 07 11 point method (default:False).\n \"\"\"\n if use_07_metric:\n # 11 point metric\n ap = 0.0\n for t in np.arange(0.0, 1.1, 0.1):\n if np.sum(rec >= t) == 0:\n p = 0\n else:\n p = np.max(prec[rec >= t])\n ap = ap + p / 11.0\n else:\n # correct AP calculation\n # first append sentinel values at the end\n mrec = np.concatenate(([0.0], rec, [1.0]))\n mpre = np.concatenate(([0.0], prec, [0.0]))\n\n # compute the precision envelope\n for i in range(mpre.size - 1, 0, -1):\n mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])\n\n # to calculate area under PR curve, look for points\n # where X axis (recall) changes value\n i = np.where(mrec[1:] != mrec[:-1])[0]\n\n # and sum (\\Delta recall) * prec\n ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])\n return ap\n\n\ndef voc_eval(\n detpath,\n annopath,\n imagesetfile,\n classname,\n cachedir,\n ovthresh=0.5,\n use_07_metric=False,\n):\n # first load gt\n if not os.path.isdir(cachedir):\n os.mkdir(cachedir)\n cachefile = os.path.join(cachedir, \"annots.pkl\")\n # read list of images\n with open(imagesetfile, \"r\") as f:\n lines = f.readlines()\n imagenames = [x.strip() for x in lines]\n\n if not os.path.isfile(cachefile):\n # load annots\n recs = {}\n for i, imagename in enumerate(imagenames):\n recs[imagename] = parse_rec(annopath.format(imagename))\n if i % 100 == 0:\n print(\"Reading annotation for {:d}/{:d}\".format(i + 1, len(imagenames)))\n # save\n print(\"Saving cached annotations to {:s}\".format(cachefile))\n with open(cachefile, \"wb\") as f:\n pickle.dump(recs, f)\n else:\n # load\n with open(cachefile, \"rb\") as f:\n recs = pickle.load(f)\n\n # extract gt objects for this class\n class_recs = {}\n npos = 0\n for imagename in imagenames:\n R = [obj for obj in recs[imagename] if obj[\"name\"] == classname]\n bbox = np.array([x[\"bbox\"] for x in R])\n difficult = np.array([x[\"difficult\"] for x in R]).astype(np.bool)\n det = [False] * len(R)\n npos = npos + sum(~difficult)\n class_recs[imagename] = {\"bbox\": bbox, \"difficult\": difficult, \"det\": det}\n\n # read dets\n detfile = detpath.format(classname)\n with open(detfile, \"r\") as f:\n lines = f.readlines()\n\n if len(lines) == 0:\n return 0, 0, 0\n\n splitlines = [x.strip().split(\" \") for x in lines]\n image_ids = [x[0] for x in splitlines]\n confidence = np.array([float(x[1]) for x in splitlines])\n BB = np.array([[float(z) for z in x[2:]] for x in splitlines])\n\n # sort by confidence\n sorted_ind = np.argsort(-confidence)\n BB = BB[sorted_ind, :]\n image_ids = [image_ids[x] for x in sorted_ind]\n\n # go down dets and mark TPs and FPs\n nd = len(image_ids)\n tp = np.zeros(nd)\n fp = np.zeros(nd)\n for d in range(nd):\n R = class_recs[image_ids[d]]\n bb = BB[d, :].astype(float)\n ovmax = -np.inf\n BBGT = R[\"bbox\"].astype(float)\n\n if BBGT.size > 0:\n # compute overlaps\n # intersection\n ixmin = np.maximum(BBGT[:, 0], bb[0])\n iymin = np.maximum(BBGT[:, 1], bb[1])\n ixmax = np.minimum(BBGT[:, 2], bb[2])\n iymax = np.minimum(BBGT[:, 3], bb[3])\n iw = np.maximum(ixmax - ixmin + 1.0, 0.0)\n ih = np.maximum(iymax - iymin + 1.0, 0.0)\n inters = iw * ih\n\n # union\n uni = (\n (bb[2] - bb[0] + 1.0) * (bb[3] - bb[1] + 1.0)\n + (BBGT[:, 2] - BBGT[:, 0] + 1.0) * (BBGT[:, 3] - BBGT[:, 1] + 1.0)\n - inters\n )\n\n overlaps = inters / uni\n ovmax = np.max(overlaps)\n jmax = np.argmax(overlaps)\n\n if ovmax > ovthresh:\n if not R[\"difficult\"][jmax]:\n if not R[\"det\"][jmax]:\n tp[d] = 1.0\n R[\"det\"][jmax] = 1\n else:\n fp[d] = 1.0\n else:\n fp[d] = 1.0\n\n # compute precision recall\n fp = np.cumsum(fp)\n tp = np.cumsum(tp)\n rec = tp / float(npos)\n # avoid divide by zero in case the first detection matches a difficult\n # ground truth\n prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)\n ap = voc_ap(rec, prec, use_07_metric)\n\n return rec, prec, ap\n"
},
{
"path": "detector/YOLOX/yolox/evaluators/voc_evaluator.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport sys\nimport tempfile\nimport time\nfrom collections import ChainMap\nfrom loguru import logger\nfrom tqdm import tqdm\n\nimport numpy as np\n\nimport torch\n\nfrom yolox.utils import gather, is_main_process, postprocess, synchronize, time_synchronized\n\n\nclass VOCEvaluator:\n \"\"\"\n VOC AP Evaluation class.\n \"\"\"\n\n def __init__(\n self, dataloader, img_size, confthre, nmsthre, num_classes,\n ):\n \"\"\"\n Args:\n dataloader (Dataloader): evaluate dataloader.\n img_size (int): image size after preprocess. images are resized\n to squares whose shape is (img_size, img_size).\n confthre (float): confidence threshold ranging from 0 to 1, which\n is defined in the config file.\n nmsthre (float): IoU threshold of non-max supression ranging from 0 to 1.\n \"\"\"\n self.dataloader = dataloader\n self.img_size = img_size\n self.confthre = confthre\n self.nmsthre = nmsthre\n self.num_classes = num_classes\n self.num_images = len(dataloader.dataset)\n\n def evaluate(\n self, model, distributed=False, half=False, trt_file=None, decoder=None, test_size=None\n ):\n \"\"\"\n VOC average precision (AP) Evaluation. Iterate inference on the test dataset\n and the results are evaluated by COCO API.\n\n NOTE: This function will change training mode to False, please save states if needed.\n\n Args:\n model : model to evaluate.\n\n Returns:\n ap50_95 (float) : COCO style AP of IoU=50:95\n ap50 (float) : VOC 2007 metric AP of IoU=50\n summary (sr): summary info of evaluation.\n \"\"\"\n # TODO half to amp_test\n tensor_type = torch.cuda.HalfTensor if half else torch.cuda.FloatTensor\n model = model.eval()\n if half:\n model = model.half()\n ids = []\n data_list = {}\n progress_bar = tqdm if is_main_process() else iter\n\n inference_time = 0\n nms_time = 0\n n_samples = len(self.dataloader) - 1\n\n if trt_file is not None:\n from torch2trt import TRTModule\n model_trt = TRTModule()\n model_trt.load_state_dict(torch.load(trt_file))\n\n x = torch.ones(1, 3, test_size[0], test_size[1]).cuda()\n model(x)\n model = model_trt\n\n for cur_iter, (imgs, _, info_imgs, ids) in enumerate(progress_bar(self.dataloader)):\n with torch.no_grad():\n imgs = imgs.type(tensor_type)\n\n # skip the the last iters since batchsize might be not enough for batch inference\n is_time_record = cur_iter < len(self.dataloader) - 1\n if is_time_record:\n start = time.time()\n\n outputs = model(imgs)\n if decoder is not None:\n outputs = decoder(outputs, dtype=outputs.type())\n\n if is_time_record:\n infer_end = time_synchronized()\n inference_time += infer_end - start\n\n outputs = postprocess(\n outputs, self.num_classes, self.confthre, self.nmsthre\n )\n if is_time_record:\n nms_end = time_synchronized()\n nms_time += nms_end - infer_end\n\n data_list.update(self.convert_to_voc_format(outputs, info_imgs, ids))\n\n statistics = torch.cuda.FloatTensor([inference_time, nms_time, n_samples])\n if distributed:\n data_list = gather(data_list, dst=0)\n data_list = ChainMap(*data_list)\n torch.distributed.reduce(statistics, dst=0)\n\n eval_results = self.evaluate_prediction(data_list, statistics)\n synchronize()\n return eval_results\n\n def convert_to_voc_format(self, outputs, info_imgs, ids):\n predictions = {}\n for (output, img_h, img_w, img_id) in zip(outputs, info_imgs[0], info_imgs[1], ids):\n if output is None:\n predictions[int(img_id)] = (None, None, None)\n continue\n output = output.cpu()\n\n bboxes = output[:, 0:4]\n\n # preprocessing: resize\n scale = min(self.img_size[0] / float(img_h), self.img_size[1] / float(img_w))\n bboxes /= scale\n\n cls = output[:, 6]\n scores = output[:, 4] * output[:, 5]\n\n predictions[int(img_id)] = (bboxes, cls, scores)\n return predictions\n\n def evaluate_prediction(self, data_dict, statistics):\n if not is_main_process():\n return 0, 0, None\n\n logger.info(\"Evaluate in main process...\")\n\n inference_time = statistics[0].item()\n nms_time = statistics[1].item()\n n_samples = statistics[2].item()\n\n a_infer_time = 1000 * inference_time / (n_samples * self.dataloader.batch_size)\n a_nms_time = 1000 * nms_time / (n_samples * self.dataloader.batch_size)\n\n time_info = \", \".join(\n [\"Average {} time: {:.2f} ms\".format(k, v) for k, v in zip(\n [\"forward\", \"NMS\", \"inference\"],\n [a_infer_time, a_nms_time, (a_infer_time + a_nms_time)]\n )]\n )\n\n info = time_info + \"\\n\"\n\n all_boxes = [[[] for _ in range(self.num_images)] for _ in range(self.num_classes)]\n for img_num in range(self.num_images):\n bboxes, cls, scores = data_dict[img_num]\n if bboxes is None:\n for j in range(self.num_classes):\n all_boxes[j][img_num] = np.empty([0, 5], dtype=np.float32)\n continue\n for j in range(self.num_classes):\n mask_c = cls == j\n if sum(mask_c) == 0:\n all_boxes[j][img_num] = np.empty([0, 5], dtype=np.float32)\n continue\n\n c_dets = torch.cat((bboxes, scores.unsqueeze(1)), dim=1)\n all_boxes[j][img_num] = c_dets[mask_c].numpy()\n\n sys.stdout.write(\n \"im_eval: {:d}/{:d} \\r\".format(img_num + 1, self.num_images)\n )\n sys.stdout.flush()\n\n with tempfile.TemporaryDirectory() as tempdir:\n mAP50, mAP70 = self.dataloader.dataset.evaluate_detections(all_boxes, tempdir)\n return mAP50, mAP70, info\n"
},
{
"path": "detector/YOLOX/yolox/exp/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nfrom .base_exp import BaseExp\nfrom .build import get_exp\nfrom .yolox_base import Exp\n"
},
{
"path": "detector/YOLOX/yolox/exp/base_exp.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport ast\nimport pprint\nfrom abc import ABCMeta, abstractmethod\nfrom typing import Dict\nfrom tabulate import tabulate\n\nimport torch\nfrom torch.nn import Module\n\nfrom ..utils import LRScheduler\n\n\nclass BaseExp(metaclass=ABCMeta):\n \"\"\"Basic class for any experiment.\n \"\"\"\n\n def __init__(self):\n self.seed = None\n self.output_dir = \"./YOLOX_outputs\"\n self.print_interval = 100\n self.eval_interval = 10\n\n @abstractmethod\n def get_model(self) -> Module:\n pass\n\n @abstractmethod\n def get_data_loader(\n self, batch_size: int, is_distributed: bool\n ) -> Dict[str, torch.utils.data.DataLoader]:\n pass\n\n @abstractmethod\n def get_optimizer(self, batch_size: int) -> torch.optim.Optimizer:\n pass\n\n @abstractmethod\n def get_lr_scheduler(\n self, lr: float, iters_per_epoch: int, **kwargs\n ) -> LRScheduler:\n pass\n\n @abstractmethod\n def get_evaluator(self):\n pass\n\n @abstractmethod\n def eval(self, model, evaluator, weights):\n pass\n\n def __repr__(self):\n table_header = [\"keys\", \"values\"]\n exp_table = [\n (str(k), pprint.pformat(v)) for k, v in vars(self).items() if not k.startswith(\"_\")\n ]\n return tabulate(exp_table, headers=table_header, tablefmt=\"fancy_grid\")\n\n def merge(self, cfg_list):\n assert len(cfg_list) % 2 == 0\n for k, v in zip(cfg_list[0::2], cfg_list[1::2]):\n # only update value with same key\n if hasattr(self, k):\n src_value = getattr(self, k)\n src_type = type(src_value)\n if src_value is not None and src_type != type(v):\n try:\n v = src_type(v)\n except Exception:\n v = ast.literal_eval(v)\n setattr(self, k, v)\n"
},
{
"path": "detector/YOLOX/yolox/exp/build.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport importlib\nimport os\nimport sys\n\n\ndef get_exp_by_file(exp_file):\n sys.path.append(os.path.dirname(exp_file))\n current_exp = importlib.import_module(os.path.basename(exp_file).split(\".\")[0])\n exp = current_exp.Exp()\n return exp\n\n\ndef get_exp_by_name(exp_name):\n import yolox\n yolox_path = os.path.dirname(os.path.dirname(yolox.__file__))\n filedict = {\n \"yolox-s\": \"yolox_s.py\",\n \"yolox-m\": \"yolox_m.py\",\n \"yolox-l\": \"yolox_l.py\",\n \"yolox-x\": \"yolox_x.py\",\n \"yolox-tiny\": \"yolox_tiny.py\",\n \"yolox-nano\": \"nano.py\",\n \"yolov3\": \"yolov3.py\",\n }\n filename = filedict[exp_name]\n exp_path = os.path.join(yolox_path, \"exps\", \"default\", filename)\n return get_exp_by_file(exp_path)\n\n\ndef get_exp(exp_file, exp_name):\n \"\"\"\n get Exp object by file or name. If exp_file and exp_name\n are both provided, get Exp by exp_file.\n\n Args:\n exp_file (str): file path of experiment.\n exp_name (str): name of experiment. \"yolo-s\",\n \"\"\"\n assert exp_file is not None or exp_name is not None, \"plz provide exp file or exp name.\"\n if exp_file is not None:\n return get_exp_by_file(exp_file)\n else:\n return get_exp_by_name(exp_name)\n"
},
{
"path": "detector/YOLOX/yolox/exp/yolox_base.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport os\nimport random\n\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\n\nfrom .base_exp import BaseExp\n\n\nclass Exp(BaseExp):\n\n def __init__(self):\n super().__init__()\n\n # ---------------- model config ---------------- #\n self.num_classes = 80\n self.depth = 1.00\n self.width = 1.00\n\n # ---------------- dataloader config ---------------- #\n # set worker to 4 for shorter dataloader init time\n self.data_num_workers = 4\n self.input_size = (640, 640)\n self.random_size = (14, 26)\n self.train_ann = \"instances_train2017.json\"\n self.val_ann = \"instances_val2017.json\"\n\n # --------------- transform config ----------------- #\n self.degrees = 10.0\n self.translate = 0.1\n self.scale = (0.1, 2)\n self.mscale = (0.8, 1.6)\n self.shear = 2.0\n self.perspective = 0.0\n self.enable_mixup = True\n\n # -------------- training config --------------------- #\n self.warmup_epochs = 5\n self.max_epoch = 300\n self.warmup_lr = 0\n self.basic_lr_per_img = 0.01 / 64.0\n self.scheduler = \"yoloxwarmcos\"\n self.no_aug_epochs = 15\n self.min_lr_ratio = 0.05\n self.ema = True\n\n self.weight_decay = 5e-4\n self.momentum = 0.9\n self.print_interval = 10\n self.eval_interval = 10\n self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(\".\")[0]\n\n # ----------------- testing config ------------------ #\n self.test_size = (640, 640)\n self.test_conf = 0.01\n self.nmsthre = 0.65\n\n def get_model(self):\n from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead\n\n def init_yolo(M):\n for m in M.modules():\n if isinstance(m, nn.BatchNorm2d):\n m.eps = 1e-3\n m.momentum = 0.03\n\n if getattr(self, \"model\", None) is None:\n in_channels = [256, 512, 1024]\n backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels)\n head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels)\n self.model = YOLOX(backbone, head)\n\n self.model.apply(init_yolo)\n self.model.head.initialize_biases(1e-2)\n return self.model\n\n def get_data_loader(self, batch_size, is_distributed, no_aug=False):\n from yolox.data import (\n COCODataset,\n TrainTransform,\n YoloBatchSampler,\n DataLoader,\n InfiniteSampler,\n MosaicDetection,\n )\n\n dataset = COCODataset(\n data_dir=None,\n json_file=self.train_ann,\n img_size=self.input_size,\n preproc=TrainTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n max_labels=50,\n ),\n )\n\n dataset = MosaicDetection(\n dataset,\n mosaic=not no_aug,\n img_size=self.input_size,\n preproc=TrainTransform(\n rgb_means=(0.485, 0.456, 0.406),\n std=(0.229, 0.224, 0.225),\n max_labels=120,\n ),\n degrees=self.degrees,\n translate=self.translate,\n scale=self.scale,\n shear=self.shear,\n perspective=self.perspective,\n enable_mixup=self.enable_mixup,\n )\n\n self.dataset = dataset\n\n if is_distributed:\n batch_size = batch_size // dist.get_world_size()\n\n sampler = InfiniteSampler(\n len(self.dataset), seed=self.seed if self.seed else 0\n )\n\n batch_sampler = YoloBatchSampler(\n sampler=sampler,\n batch_size=batch_size,\n drop_last=False,\n input_dimension=self.input_size,\n mosaic=not no_aug,\n )\n\n dataloader_kwargs = {\"num_workers\": self.data_num_workers, \"pin_memory\": True}\n dataloader_kwargs[\"batch_sampler\"] = batch_sampler\n train_loader = DataLoader(self.dataset, **dataloader_kwargs)\n\n return train_loader\n\n def random_resize(self, data_loader, epoch, rank, is_distributed):\n tensor = torch.LongTensor(2).cuda()\n\n if rank == 0:\n size_factor = self.input_size[1] * 1. / self.input_size[0]\n size = random.randint(*self.random_size)\n size = (int(32 * size), 32 * int(size * size_factor))\n tensor[0] = size[0]\n tensor[1] = size[1]\n\n if is_distributed:\n dist.barrier()\n dist.broadcast(tensor, 0)\n\n input_size = data_loader.change_input_dim(\n multiple=(tensor[0].item(), tensor[1].item()), random_range=None\n )\n return input_size\n\n def get_optimizer(self, batch_size):\n if \"optimizer\" not in self.__dict__:\n if self.warmup_epochs > 0:\n lr = self.warmup_lr\n else:\n lr = self.basic_lr_per_img * batch_size\n\n pg0, pg1, pg2 = [], [], [] # optimizer parameter groups\n\n for k, v in self.model.named_modules():\n if hasattr(v, \"bias\") and isinstance(v.bias, nn.Parameter):\n pg2.append(v.bias) # biases\n if isinstance(v, nn.BatchNorm2d) or \"bn\" in k:\n pg0.append(v.weight) # no decay\n elif hasattr(v, \"weight\") and isinstance(v.weight, nn.Parameter):\n pg1.append(v.weight) # apply decay\n\n optimizer = torch.optim.SGD(\n pg0, lr=lr, momentum=self.momentum, nesterov=True\n )\n optimizer.add_param_group(\n {\"params\": pg1, \"weight_decay\": self.weight_decay}\n ) # add pg1 with weight_decay\n optimizer.add_param_group({\"params\": pg2})\n self.optimizer = optimizer\n\n return self.optimizer\n\n def get_lr_scheduler(self, lr, iters_per_epoch):\n from yolox.utils import LRScheduler\n scheduler = LRScheduler(\n self.scheduler,\n lr,\n iters_per_epoch,\n self.max_epoch,\n warmup_epochs=self.warmup_epochs,\n warmup_lr_start=self.warmup_lr,\n no_aug_epochs=self.no_aug_epochs,\n min_lr_ratio=self.min_lr_ratio,\n )\n return scheduler\n\n def get_eval_loader(self, batch_size, is_distributed, testdev=False):\n from yolox.data import COCODataset, ValTransform\n\n valdataset = COCODataset(\n data_dir=None,\n json_file=self.val_ann if not testdev else \"image_info_test-dev2017.json\",\n name=\"val2017\" if not testdev else \"test2017\",\n img_size=self.test_size,\n preproc=ValTransform(\n rgb_means=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)\n ),\n )\n\n if is_distributed:\n batch_size = batch_size // dist.get_world_size()\n sampler = torch.utils.data.distributed.DistributedSampler(\n valdataset, shuffle=False\n )\n else:\n sampler = torch.utils.data.SequentialSampler(valdataset)\n\n dataloader_kwargs = {\n \"num_workers\": self.data_num_workers,\n \"pin_memory\": True,\n \"sampler\": sampler,\n }\n dataloader_kwargs[\"batch_size\"] = batch_size\n val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)\n\n return val_loader\n\n def get_evaluator(self, batch_size, is_distributed, testdev=False):\n from yolox.evaluators import COCOEvaluator\n\n val_loader = self.get_eval_loader(batch_size, is_distributed, testdev=testdev)\n evaluator = COCOEvaluator(\n dataloader=val_loader,\n img_size=self.test_size,\n confthre=self.test_conf,\n nmsthre=self.nmsthre,\n num_classes=self.num_classes,\n testdev=testdev,\n )\n return evaluator\n\n def eval(self, model, evaluator, is_distributed, half=False):\n return evaluator.evaluate(model, is_distributed, half)\n"
},
{
"path": "detector/YOLOX/yolox/layers/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nfrom .fast_coco_eval_api import COCOeval_opt\n"
},
{
"path": "detector/YOLOX/yolox/layers/csrc/cocoeval/cocoeval.cpp",
"content": "// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n#include \"cocoeval.h\"\n#include \n#include \n#include \n#include \n\nusing namespace pybind11::literals;\n\nnamespace COCOeval {\n\n// Sort detections from highest score to lowest, such that\n// detection_instances[detection_sorted_indices[t]] >=\n// detection_instances[detection_sorted_indices[t+1]]. Use stable_sort to match\n// original COCO API\nvoid SortInstancesByDetectionScore(\n const std::vector& detection_instances,\n std::vector* detection_sorted_indices) {\n detection_sorted_indices->resize(detection_instances.size());\n std::iota(\n detection_sorted_indices->begin(), detection_sorted_indices->end(), 0);\n std::stable_sort(\n detection_sorted_indices->begin(),\n detection_sorted_indices->end(),\n [&detection_instances](size_t j1, size_t j2) {\n return detection_instances[j1].score > detection_instances[j2].score;\n });\n}\n\n// Partition the ground truth objects based on whether or not to ignore them\n// based on area\nvoid SortInstancesByIgnore(\n const std::array& area_range,\n const std::vector& ground_truth_instances,\n std::vector* ground_truth_sorted_indices,\n std::vector* ignores) {\n ignores->clear();\n ignores->reserve(ground_truth_instances.size());\n for (auto o : ground_truth_instances) {\n ignores->push_back(\n o.ignore || o.area < area_range[0] || o.area > area_range[1]);\n }\n\n ground_truth_sorted_indices->resize(ground_truth_instances.size());\n std::iota(\n ground_truth_sorted_indices->begin(),\n ground_truth_sorted_indices->end(),\n 0);\n std::stable_sort(\n ground_truth_sorted_indices->begin(),\n ground_truth_sorted_indices->end(),\n [&ignores](size_t j1, size_t j2) {\n return (int)(*ignores)[j1] < (int)(*ignores)[j2];\n });\n}\n\n// For each IOU threshold, greedily match each detected instance to a ground\n// truth instance (if possible) and store the results\nvoid MatchDetectionsToGroundTruth(\n const std::vector& detection_instances,\n const std::vector& detection_sorted_indices,\n const std::vector& ground_truth_instances,\n const std::vector& ground_truth_sorted_indices,\n const std::vector& ignores,\n const std::vector>& ious,\n const std::vector& iou_thresholds,\n const std::array& area_range,\n ImageEvaluation* results) {\n // Initialize memory to store return data matches and ignore\n const int num_iou_thresholds = iou_thresholds.size();\n const int num_ground_truth = ground_truth_sorted_indices.size();\n const int num_detections = detection_sorted_indices.size();\n std::vector ground_truth_matches(\n num_iou_thresholds * num_ground_truth, 0);\n std::vector& detection_matches = results->detection_matches;\n std::vector& detection_ignores = results->detection_ignores;\n std::vector& ground_truth_ignores = results->ground_truth_ignores;\n detection_matches.resize(num_iou_thresholds * num_detections, 0);\n detection_ignores.resize(num_iou_thresholds * num_detections, false);\n ground_truth_ignores.resize(num_ground_truth);\n for (auto g = 0; g < num_ground_truth; ++g) {\n ground_truth_ignores[g] = ignores[ground_truth_sorted_indices[g]];\n }\n\n for (auto t = 0; t < num_iou_thresholds; ++t) {\n for (auto d = 0; d < num_detections; ++d) {\n // information about best match so far (match=-1 -> unmatched)\n double best_iou = std::min(iou_thresholds[t], 1 - 1e-10);\n int match = -1;\n for (auto g = 0; g < num_ground_truth; ++g) {\n // if this ground truth instance is already matched and not a\n // crowd, it cannot be matched to another detection\n if (ground_truth_matches[t * num_ground_truth + g] > 0 &&\n !ground_truth_instances[ground_truth_sorted_indices[g]].is_crowd) {\n continue;\n }\n\n // if detected instance matched to a regular ground truth\n // instance, we can break on the first ground truth instance\n // tagged as ignore (because they are sorted by the ignore tag)\n if (match >= 0 && !ground_truth_ignores[match] &&\n ground_truth_ignores[g]) {\n break;\n }\n\n // if IOU overlap is the best so far, store the match appropriately\n if (ious[d][ground_truth_sorted_indices[g]] >= best_iou) {\n best_iou = ious[d][ground_truth_sorted_indices[g]];\n match = g;\n }\n }\n // if match was made, store id of match for both detection and\n // ground truth\n if (match >= 0) {\n detection_ignores[t * num_detections + d] = ground_truth_ignores[match];\n detection_matches[t * num_detections + d] =\n ground_truth_instances[ground_truth_sorted_indices[match]].id;\n ground_truth_matches[t * num_ground_truth + match] =\n detection_instances[detection_sorted_indices[d]].id;\n }\n\n // set unmatched detections outside of area range to ignore\n const InstanceAnnotation& detection =\n detection_instances[detection_sorted_indices[d]];\n detection_ignores[t * num_detections + d] =\n detection_ignores[t * num_detections + d] ||\n (detection_matches[t * num_detections + d] == 0 &&\n (detection.area < area_range[0] || detection.area > area_range[1]));\n }\n }\n\n // store detection score results\n results->detection_scores.resize(detection_sorted_indices.size());\n for (size_t d = 0; d < detection_sorted_indices.size(); ++d) {\n results->detection_scores[d] =\n detection_instances[detection_sorted_indices[d]].score;\n }\n}\n\nstd::vector EvaluateImages(\n const std::vector>& area_ranges,\n int max_detections,\n const std::vector& iou_thresholds,\n const ImageCategoryInstances>& image_category_ious,\n const ImageCategoryInstances&\n image_category_ground_truth_instances,\n const ImageCategoryInstances&\n image_category_detection_instances) {\n const int num_area_ranges = area_ranges.size();\n const int num_images = image_category_ground_truth_instances.size();\n const int num_categories =\n image_category_ious.size() > 0 ? image_category_ious[0].size() : 0;\n std::vector detection_sorted_indices;\n std::vector ground_truth_sorted_indices;\n std::vector ignores;\n std::vector results_all(\n num_images * num_area_ranges * num_categories);\n\n // Store results for each image, category, and area range combination. Results\n // for each IOU threshold are packed into the same ImageEvaluation object\n for (auto i = 0; i < num_images; ++i) {\n for (auto c = 0; c < num_categories; ++c) {\n const std::vector& ground_truth_instances =\n image_category_ground_truth_instances[i][c];\n const std::vector& detection_instances =\n image_category_detection_instances[i][c];\n\n SortInstancesByDetectionScore(\n detection_instances, &detection_sorted_indices);\n if ((int)detection_sorted_indices.size() > max_detections) {\n detection_sorted_indices.resize(max_detections);\n }\n\n for (size_t a = 0; a < area_ranges.size(); ++a) {\n SortInstancesByIgnore(\n area_ranges[a],\n ground_truth_instances,\n &ground_truth_sorted_indices,\n &ignores);\n\n MatchDetectionsToGroundTruth(\n detection_instances,\n detection_sorted_indices,\n ground_truth_instances,\n ground_truth_sorted_indices,\n ignores,\n image_category_ious[i][c],\n iou_thresholds,\n area_ranges[a],\n &results_all\n [c * num_area_ranges * num_images + a * num_images + i]);\n }\n }\n }\n\n return results_all;\n}\n\n// Convert a python list to a vector\ntemplate \nstd::vector list_to_vec(const py::list& l) {\n std::vector v(py::len(l));\n for (int i = 0; i < (int)py::len(l); ++i) {\n v[i] = l[i].cast();\n }\n return v;\n}\n\n// Helper function to Accumulate()\n// Considers the evaluation results applicable to a particular category, area\n// range, and max_detections parameter setting, which begin at\n// evaluations[evaluation_index]. Extracts a sorted list of length n of all\n// applicable detection instances concatenated across all images in the dataset,\n// which are represented by the outputs evaluation_indices, detection_scores,\n// image_detection_indices, and detection_sorted_indices--all of which are\n// length n. evaluation_indices[i] stores the applicable index into\n// evaluations[] for instance i, which has detection score detection_score[i],\n// and is the image_detection_indices[i]'th of the list of detections\n// for the image containing i. detection_sorted_indices[] defines a sorted\n// permutation of the 3 other outputs\nint BuildSortedDetectionList(\n const std::vector& evaluations,\n const int64_t evaluation_index,\n const int64_t num_images,\n const int max_detections,\n std::vector* evaluation_indices,\n std::vector* detection_scores,\n std::vector* detection_sorted_indices,\n std::vector* image_detection_indices) {\n assert(evaluations.size() >= evaluation_index + num_images);\n\n // Extract a list of object instances of the applicable category, area\n // range, and max detections requirements such that they can be sorted\n image_detection_indices->clear();\n evaluation_indices->clear();\n detection_scores->clear();\n image_detection_indices->reserve(num_images * max_detections);\n evaluation_indices->reserve(num_images * max_detections);\n detection_scores->reserve(num_images * max_detections);\n int num_valid_ground_truth = 0;\n for (auto i = 0; i < num_images; ++i) {\n const ImageEvaluation& evaluation = evaluations[evaluation_index + i];\n\n for (int d = 0;\n d < (int)evaluation.detection_scores.size() && d < max_detections;\n ++d) { // detected instances\n evaluation_indices->push_back(evaluation_index + i);\n image_detection_indices->push_back(d);\n detection_scores->push_back(evaluation.detection_scores[d]);\n }\n for (auto ground_truth_ignore : evaluation.ground_truth_ignores) {\n if (!ground_truth_ignore) {\n ++num_valid_ground_truth;\n }\n }\n }\n\n // Sort detections by decreasing score, using stable sort to match\n // python implementation\n detection_sorted_indices->resize(detection_scores->size());\n std::iota(\n detection_sorted_indices->begin(), detection_sorted_indices->end(), 0);\n std::stable_sort(\n detection_sorted_indices->begin(),\n detection_sorted_indices->end(),\n [&detection_scores](size_t j1, size_t j2) {\n return (*detection_scores)[j1] > (*detection_scores)[j2];\n });\n\n return num_valid_ground_truth;\n}\n\n// Helper function to Accumulate()\n// Compute a precision recall curve given a sorted list of detected instances\n// encoded in evaluations, evaluation_indices, detection_scores,\n// detection_sorted_indices, image_detection_indices (see\n// BuildSortedDetectionList()). Using vectors precisions and recalls\n// and temporary storage, output the results into precisions_out, recalls_out,\n// and scores_out, which are large buffers containing many precion/recall curves\n// for all possible parameter settings, with precisions_out_index and\n// recalls_out_index defining the applicable indices to store results.\nvoid ComputePrecisionRecallCurve(\n const int64_t precisions_out_index,\n const int64_t precisions_out_stride,\n const int64_t recalls_out_index,\n const std::vector& recall_thresholds,\n const int iou_threshold_index,\n const int num_iou_thresholds,\n const int num_valid_ground_truth,\n const std::vector& evaluations,\n const std::vector& evaluation_indices,\n const std::vector& detection_scores,\n const std::vector& detection_sorted_indices,\n const std::vector& image_detection_indices,\n std::vector* precisions,\n std::vector* recalls,\n std::vector* precisions_out,\n std::vector* scores_out,\n std::vector* recalls_out) {\n assert(recalls_out->size() > recalls_out_index);\n\n // Compute precision/recall for each instance in the sorted list of detections\n int64_t true_positives_sum = 0, false_positives_sum = 0;\n precisions->clear();\n recalls->clear();\n precisions->reserve(detection_sorted_indices.size());\n recalls->reserve(detection_sorted_indices.size());\n assert(!evaluations.empty() || detection_sorted_indices.empty());\n for (auto detection_sorted_index : detection_sorted_indices) {\n const ImageEvaluation& evaluation =\n evaluations[evaluation_indices[detection_sorted_index]];\n const auto num_detections =\n evaluation.detection_matches.size() / num_iou_thresholds;\n const auto detection_index = iou_threshold_index * num_detections +\n image_detection_indices[detection_sorted_index];\n assert(evaluation.detection_matches.size() > detection_index);\n assert(evaluation.detection_ignores.size() > detection_index);\n const int64_t detection_match =\n evaluation.detection_matches[detection_index];\n const bool detection_ignores =\n evaluation.detection_ignores[detection_index];\n const auto true_positive = detection_match > 0 && !detection_ignores;\n const auto false_positive = detection_match == 0 && !detection_ignores;\n if (true_positive) {\n ++true_positives_sum;\n }\n if (false_positive) {\n ++false_positives_sum;\n }\n\n const double recall =\n static_cast(true_positives_sum) / num_valid_ground_truth;\n recalls->push_back(recall);\n const int64_t num_valid_detections =\n true_positives_sum + false_positives_sum;\n const double precision = num_valid_detections > 0\n ? static_cast(true_positives_sum) / num_valid_detections\n : 0.0;\n precisions->push_back(precision);\n }\n\n (*recalls_out)[recalls_out_index] = !recalls->empty() ? recalls->back() : 0;\n\n for (int64_t i = static_cast(precisions->size()) - 1; i > 0; --i) {\n if ((*precisions)[i] > (*precisions)[i - 1]) {\n (*precisions)[i - 1] = (*precisions)[i];\n }\n }\n\n // Sample the per instance precision/recall list at each recall threshold\n for (size_t r = 0; r < recall_thresholds.size(); ++r) {\n // first index in recalls >= recall_thresholds[r]\n std::vector::iterator low = std::lower_bound(\n recalls->begin(), recalls->end(), recall_thresholds[r]);\n size_t precisions_index = low - recalls->begin();\n\n const auto results_ind = precisions_out_index + r * precisions_out_stride;\n assert(results_ind < precisions_out->size());\n assert(results_ind < scores_out->size());\n if (precisions_index < precisions->size()) {\n (*precisions_out)[results_ind] = (*precisions)[precisions_index];\n (*scores_out)[results_ind] =\n detection_scores[detection_sorted_indices[precisions_index]];\n } else {\n (*precisions_out)[results_ind] = 0;\n (*scores_out)[results_ind] = 0;\n }\n }\n}\npy::dict Accumulate(\n const py::object& params,\n const std::vector& evaluations) {\n const std::vector recall_thresholds =\n list_to_vec(params.attr(\"recThrs\"));\n const std::vector max_detections =\n list_to_vec(params.attr(\"maxDets\"));\n const int num_iou_thresholds = py::len(params.attr(\"iouThrs\"));\n const int num_recall_thresholds = py::len(params.attr(\"recThrs\"));\n const int num_categories = params.attr(\"useCats\").cast() == 1\n ? py::len(params.attr(\"catIds\"))\n : 1;\n const int num_area_ranges = py::len(params.attr(\"areaRng\"));\n const int num_max_detections = py::len(params.attr(\"maxDets\"));\n const int num_images = py::len(params.attr(\"imgIds\"));\n\n std::vector precisions_out(\n num_iou_thresholds * num_recall_thresholds * num_categories *\n num_area_ranges * num_max_detections,\n -1);\n std::vector recalls_out(\n num_iou_thresholds * num_categories * num_area_ranges *\n num_max_detections,\n -1);\n std::vector scores_out(\n num_iou_thresholds * num_recall_thresholds * num_categories *\n num_area_ranges * num_max_detections,\n -1);\n\n // Consider the list of all detected instances in the entire dataset in one\n // large list. evaluation_indices, detection_scores,\n // image_detection_indices, and detection_sorted_indices all have the same\n // length as this list, such that each entry corresponds to one detected\n // instance\n std::vector evaluation_indices; // indices into evaluations[]\n std::vector detection_scores; // detection scores of each instance\n std::vector detection_sorted_indices; // sorted indices of all\n // instances in the dataset\n std::vector\n image_detection_indices; // indices into the list of detected instances in\n // the same image as each instance\n std::vector precisions, recalls;\n\n for (auto c = 0; c < num_categories; ++c) {\n for (auto a = 0; a < num_area_ranges; ++a) {\n for (auto m = 0; m < num_max_detections; ++m) {\n // The COCO PythonAPI assumes evaluations[] (the return value of\n // COCOeval::EvaluateImages() is one long list storing results for each\n // combination of category, area range, and image id, with categories in\n // the outermost loop and images in the innermost loop.\n const int64_t evaluations_index =\n c * num_area_ranges * num_images + a * num_images;\n int num_valid_ground_truth = BuildSortedDetectionList(\n evaluations,\n evaluations_index,\n num_images,\n max_detections[m],\n &evaluation_indices,\n &detection_scores,\n &detection_sorted_indices,\n &image_detection_indices);\n\n if (num_valid_ground_truth == 0) {\n continue;\n }\n\n for (auto t = 0; t < num_iou_thresholds; ++t) {\n // recalls_out is a flattened vectors representing a\n // num_iou_thresholds X num_categories X num_area_ranges X\n // num_max_detections matrix\n const int64_t recalls_out_index =\n t * num_categories * num_area_ranges * num_max_detections +\n c * num_area_ranges * num_max_detections +\n a * num_max_detections + m;\n\n // precisions_out and scores_out are flattened vectors\n // representing a num_iou_thresholds X num_recall_thresholds X\n // num_categories X num_area_ranges X num_max_detections matrix\n const int64_t precisions_out_stride =\n num_categories * num_area_ranges * num_max_detections;\n const int64_t precisions_out_index = t * num_recall_thresholds *\n num_categories * num_area_ranges * num_max_detections +\n c * num_area_ranges * num_max_detections +\n a * num_max_detections + m;\n\n ComputePrecisionRecallCurve(\n precisions_out_index,\n precisions_out_stride,\n recalls_out_index,\n recall_thresholds,\n t,\n num_iou_thresholds,\n num_valid_ground_truth,\n evaluations,\n evaluation_indices,\n detection_scores,\n detection_sorted_indices,\n image_detection_indices,\n &precisions,\n &recalls,\n &precisions_out,\n &scores_out,\n &recalls_out);\n }\n }\n }\n }\n\n time_t rawtime;\n struct tm local_time;\n std::array buffer;\n time(&rawtime);\n#ifdef _WIN32\n localtime_s(&local_time, &rawtime);\n#else\n localtime_r(&rawtime, &local_time);\n#endif\n strftime(\n buffer.data(), 200, \"%Y-%m-%d %H:%num_max_detections:%S\", &local_time);\n return py::dict(\n \"params\"_a = params,\n \"counts\"_a = std::vector({num_iou_thresholds,\n num_recall_thresholds,\n num_categories,\n num_area_ranges,\n num_max_detections}),\n \"date\"_a = buffer,\n \"precision\"_a = precisions_out,\n \"recall\"_a = recalls_out,\n \"scores\"_a = scores_out);\n}\n\n} // namespace COCOeval\n"
},
{
"path": "detector/YOLOX/yolox/layers/csrc/cocoeval/cocoeval.h",
"content": "// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n\nnamespace py = pybind11;\n\nnamespace COCOeval {\n\n// Annotation data for a single object instance in an image\nstruct InstanceAnnotation {\n InstanceAnnotation(\n uint64_t id,\n double score,\n double area,\n bool is_crowd,\n bool ignore)\n : id{id}, score{score}, area{area}, is_crowd{is_crowd}, ignore{ignore} {}\n uint64_t id;\n double score = 0.;\n double area = 0.;\n bool is_crowd = false;\n bool ignore = false;\n};\n\n// Stores intermediate results for evaluating detection results for a single\n// image that has D detected instances and G ground truth instances. This stores\n// matches between detected and ground truth instances\nstruct ImageEvaluation {\n // For each of the D detected instances, the id of the matched ground truth\n // instance, or 0 if unmatched\n std::vector detection_matches;\n\n // The detection score of each of the D detected instances\n std::vector detection_scores;\n\n // Marks whether or not each of G instances was ignored from evaluation (e.g.,\n // because it's outside area_range)\n std::vector ground_truth_ignores;\n\n // Marks whether or not each of D instances was ignored from evaluation (e.g.,\n // because it's outside aRng)\n std::vector detection_ignores;\n};\n\ntemplate \nusing ImageCategoryInstances = std::vector>>;\n\n// C++ implementation of COCO API cocoeval.py::COCOeval.evaluateImg(). For each\n// combination of image, category, area range settings, and IOU thresholds to\n// evaluate, it matches detected instances to ground truth instances and stores\n// the results into a vector of ImageEvaluation results, which will be\n// interpreted by the COCOeval::Accumulate() function to produce precion-recall\n// curves. The parameters of nested vectors have the following semantics:\n// image_category_ious[i][c][d][g] is the intersection over union of the d'th\n// detected instance and g'th ground truth instance of\n// category category_ids[c] in image image_ids[i]\n// image_category_ground_truth_instances[i][c] is a vector of ground truth\n// instances in image image_ids[i] of category category_ids[c]\n// image_category_detection_instances[i][c] is a vector of detected\n// instances in image image_ids[i] of category category_ids[c]\nstd::vector EvaluateImages(\n const std::vector>& area_ranges, // vector of 2-tuples\n int max_detections,\n const std::vector& iou_thresholds,\n const ImageCategoryInstances>& image_category_ious,\n const ImageCategoryInstances&\n image_category_ground_truth_instances,\n const ImageCategoryInstances&\n image_category_detection_instances);\n\n// C++ implementation of COCOeval.accumulate(), which generates precision\n// recall curves for each set of category, IOU threshold, detection area range,\n// and max number of detections parameters. It is assumed that the parameter\n// evaluations is the return value of the functon COCOeval::EvaluateImages(),\n// which was called with the same parameter settings params\npy::dict Accumulate(\n const py::object& params,\n const std::vector& evalutations);\n\n} // namespace COCOeval\n"
},
{
"path": "detector/YOLOX/yolox/layers/csrc/vision.cpp",
"content": "#include \"cocoeval/cocoeval.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"COCOevalAccumulate\", &COCOeval::Accumulate, \"COCOeval::Accumulate\");\n m.def(\n \"COCOevalEvaluateImages\",\n &COCOeval::EvaluateImages,\n \"COCOeval::EvaluateImages\");\n pybind11::class_(m, \"InstanceAnnotation\")\n .def(pybind11::init());\n pybind11::class_(m, \"ImageEvaluation\")\n .def(pybind11::init<>());\n}\n"
},
{
"path": "detector/YOLOX/yolox/layers/fast_coco_eval_api.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# This file comes from\n# https://github.com/facebookresearch/detectron2/blob/master/detectron2/evaluation/fast_eval_api.py\n# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport copy\nimport time\n\nimport numpy as np\nfrom pycocotools.cocoeval import COCOeval\n\n# import torch first to make yolox._C work without ImportError of libc10.so\n# in YOLOX, env is already set in __init__.py.\nfrom yolox import _C\n\n\nclass COCOeval_opt(COCOeval):\n \"\"\"\n This is a slightly modified version of the original COCO API, where the functions evaluateImg()\n and accumulate() are implemented in C++ to speedup evaluation\n \"\"\"\n\n def evaluate(self):\n \"\"\"\n Run per image evaluation on given images and store results in self.evalImgs_cpp, a\n datastructure that isn't readable from Python but is used by a c++ implementation of\n accumulate(). Unlike the original COCO PythonAPI, we don't populate the datastructure\n self.evalImgs because this datastructure is a computational bottleneck.\n :return: None\n \"\"\"\n tic = time.time()\n\n print(\"Running per image evaluation...\")\n p = self.params\n # add backward compatibility if useSegm is specified in params\n if p.useSegm is not None:\n p.iouType = \"segm\" if p.useSegm == 1 else \"bbox\"\n print(\n \"useSegm (deprecated) is not None. Running {} evaluation\".format(p.iouType)\n )\n print(\"Evaluate annotation type *{}*\".format(p.iouType))\n p.imgIds = list(np.unique(p.imgIds))\n if p.useCats:\n p.catIds = list(np.unique(p.catIds))\n p.maxDets = sorted(p.maxDets)\n self.params = p\n\n self._prepare()\n\n # loop through images, area range, max detection number\n catIds = p.catIds if p.useCats else [-1]\n\n if p.iouType == \"segm\" or p.iouType == \"bbox\":\n computeIoU = self.computeIoU\n elif p.iouType == \"keypoints\":\n computeIoU = self.computeOks\n self.ious = {\n (imgId, catId): computeIoU(imgId, catId)\n for imgId in p.imgIds\n for catId in catIds\n }\n\n maxDet = p.maxDets[-1]\n\n # <<<< Beginning of code differences with original COCO API\n def convert_instances_to_cpp(instances, is_det=False):\n # Convert annotations for a list of instances in an image to a format that's fast\n # to access in C++\n instances_cpp = []\n for instance in instances:\n instance_cpp = _C.InstanceAnnotation(\n int(instance[\"id\"]),\n instance[\"score\"] if is_det else instance.get(\"score\", 0.0),\n instance[\"area\"],\n bool(instance.get(\"iscrowd\", 0)),\n bool(instance.get(\"ignore\", 0)),\n )\n instances_cpp.append(instance_cpp)\n return instances_cpp\n\n # Convert GT annotations, detections, and IOUs to a format that's fast to access in C++\n ground_truth_instances = [\n [convert_instances_to_cpp(self._gts[imgId, catId]) for catId in p.catIds]\n for imgId in p.imgIds\n ]\n detected_instances = [\n [\n convert_instances_to_cpp(self._dts[imgId, catId], is_det=True)\n for catId in p.catIds\n ]\n for imgId in p.imgIds\n ]\n ious = [[self.ious[imgId, catId] for catId in catIds] for imgId in p.imgIds]\n\n if not p.useCats:\n # For each image, flatten per-category lists into a single list\n ground_truth_instances = [\n [[o for c in i for o in c]] for i in ground_truth_instances\n ]\n detected_instances = [\n [[o for c in i for o in c]] for i in detected_instances\n ]\n\n # Call C++ implementation of self.evaluateImgs()\n self._evalImgs_cpp = _C.COCOevalEvaluateImages(\n p.areaRng,\n maxDet,\n p.iouThrs,\n ious,\n ground_truth_instances,\n detected_instances,\n )\n self._evalImgs = None\n\n self._paramsEval = copy.deepcopy(self.params)\n toc = time.time()\n print(\"COCOeval_opt.evaluate() finished in {:0.2f} seconds.\".format(toc - tic))\n # >>>> End of code differences with original COCO API\n\n def accumulate(self):\n \"\"\"\n Accumulate per image evaluation results and store the result in self.eval. Does not\n support changing parameter settings from those used by self.evaluate()\n \"\"\"\n print(\"Accumulating evaluation results...\")\n tic = time.time()\n if not hasattr(self, \"_evalImgs_cpp\"):\n print(\"Please run evaluate() first\")\n\n self.eval = _C.COCOevalAccumulate(self._paramsEval, self._evalImgs_cpp)\n\n # recall is num_iou_thresholds X num_categories X num_area_ranges X num_max_detections\n self.eval[\"recall\"] = np.array(self.eval[\"recall\"]).reshape(\n self.eval[\"counts\"][:1] + self.eval[\"counts\"][2:]\n )\n\n # precision and scores are num_iou_thresholds X num_recall_thresholds X num_categories X\n # num_area_ranges X num_max_detections\n self.eval[\"precision\"] = np.array(self.eval[\"precision\"]).reshape(\n self.eval[\"counts\"]\n )\n self.eval[\"scores\"] = np.array(self.eval[\"scores\"]).reshape(self.eval[\"counts\"])\n toc = time.time()\n print(\n \"COCOeval_opt.accumulate() finished in {:0.2f} seconds.\".format(toc - tic)\n )\n"
},
{
"path": "detector/YOLOX/yolox/models/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nfrom .darknet import CSPDarknet, Darknet\nfrom .losses import IOUloss\nfrom .yolo_fpn import YOLOFPN\nfrom .yolo_head import YOLOXHead\nfrom .yolo_pafpn import YOLOPAFPN\nfrom .yolox import YOLOX\n"
},
{
"path": "detector/YOLOX/yolox/models/darknet.py",
"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nfrom torch import nn\n\nfrom .network_blocks import BaseConv, CSPLayer, DWConv, Focus, ResLayer, SPPBottleneck\n\n\nclass Darknet(nn.Module):\n # number of blocks from dark2 to dark5.\n depth2blocks = {21: [1, 2, 2, 1], 53: [2, 8, 8, 4]}\n\n def __init__(\n self, depth, in_channels=3, stem_out_channels=32, out_features=(\"dark3\", \"dark4\", \"dark5\"),\n ):\n \"\"\"\n Args:\n depth (int): depth of darknet used in model, usually use [21, 53] for this param.\n in_channels (int): number of input channels, for example, use 3 for RGB image.\n stem_out_channels (int): number of output chanels of darknet stem.\n It decides channels of darknet layer2 to layer5.\n out_features (Tuple[str]): desired output layer name.\n \"\"\"\n super().__init__()\n assert out_features, \"please provide output features of Darknet\"\n self.out_features = out_features\n self.stem = nn.Sequential(\n BaseConv(in_channels, stem_out_channels, ksize=3, stride=1, act=\"lrelu\"),\n *self.make_group_layer(stem_out_channels, num_blocks=1, stride=2),\n )\n in_channels = stem_out_channels * 2 # 64\n\n num_blocks = Darknet.depth2blocks[depth]\n # create darknet with `stem_out_channels` and `num_blocks` layers.\n # to make model structure more clear, we don't use `for` statement in python.\n self.dark2 = nn.Sequential(*self.make_group_layer(in_channels, num_blocks[0], stride=2))\n in_channels *= 2 # 128\n self.dark3 = nn.Sequential(*self.make_group_layer(in_channels, num_blocks[1], stride=2))\n in_channels *= 2 # 256\n self.dark4 = nn.Sequential(*self.make_group_layer(in_channels, num_blocks[2], stride=2))\n in_channels *= 2 # 512\n\n self.dark5 = nn.Sequential(\n *self.make_group_layer(in_channels, num_blocks[3], stride=2),\n *self.make_spp_block([in_channels, in_channels * 2], in_channels * 2),\n )\n\n def make_group_layer(self, in_channels: int, num_blocks: int, stride: int = 1):\n \"starts with conv layer then has `num_blocks` `ResLayer`\"\n return [\n BaseConv(in_channels, in_channels * 2, ksize=3, stride=stride, act=\"lrelu\"),\n *[(ResLayer(in_channels * 2)) for _ in range(num_blocks)]\n ]\n\n def make_spp_block(self, filters_list, in_filters):\n m = nn.Sequential(\n *[\n BaseConv(in_filters, filters_list[0], 1, stride=1, act=\"lrelu\"),\n BaseConv(filters_list[0], filters_list[1], 3, stride=1, act=\"lrelu\"),\n SPPBottleneck(\n in_channels=filters_list[1],\n out_channels=filters_list[0],\n activation=\"lrelu\"\n ),\n BaseConv(filters_list[0], filters_list[1], 3, stride=1, act=\"lrelu\"),\n BaseConv(filters_list[1], filters_list[0], 1, stride=1, act=\"lrelu\"),\n ]\n )\n return m\n\n def forward(self, x):\n outputs = {}\n x = self.stem(x)\n outputs[\"stem\"] = x\n x = self.dark2(x)\n outputs[\"dark2\"] = x\n x = self.dark3(x)\n outputs[\"dark3\"] = x\n x = self.dark4(x)\n outputs[\"dark4\"] = x\n x = self.dark5(x)\n outputs[\"dark5\"] = x\n return {k: v for k, v in outputs.items() if k in self.out_features}\n\n\nclass CSPDarknet(nn.Module):\n\n def __init__(\n self, dep_mul, wid_mul,\n out_features=(\"dark3\", \"dark4\", \"dark5\"),\n depthwise=False, act=\"silu\",\n ):\n super().__init__()\n assert out_features, \"please provide output features of Darknet\"\n self.out_features = out_features\n Conv = DWConv if depthwise else BaseConv\n\n base_channels = int(wid_mul * 64) # 64\n base_depth = max(round(dep_mul * 3), 1) # 3\n\n # stem\n self.stem = Focus(3, base_channels, ksize=3, act=act)\n\n # dark2\n self.dark2 = nn.Sequential(\n Conv(base_channels, base_channels * 2, 3, 2, act=act),\n CSPLayer(\n base_channels * 2, base_channels * 2,\n n=base_depth, depthwise=depthwise, act=act\n ),\n )\n\n # dark3\n self.dark3 = nn.Sequential(\n Conv(base_channels * 2, base_channels * 4, 3, 2, act=act),\n CSPLayer(\n base_channels * 4, base_channels * 4,\n n=base_depth * 3, depthwise=depthwise, act=act,\n ),\n )\n\n # dark4\n self.dark4 = nn.Sequential(\n Conv(base_channels * 4, base_channels * 8, 3, 2, act=act),\n CSPLayer(\n base_channels * 8, base_channels * 8,\n n=base_depth * 3, depthwise=depthwise, act=act,\n ),\n )\n\n # dark5\n self.dark5 = nn.Sequential(\n Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),\n SPPBottleneck(base_channels * 16, base_channels * 16, activation=act),\n CSPLayer(\n base_channels * 16, base_channels * 16, n=base_depth,\n shortcut=False, depthwise=depthwise, act=act,\n ),\n )\n\n def forward(self, x):\n outputs = {}\n x = self.stem(x)\n outputs[\"stem\"] = x\n x = self.dark2(x)\n outputs[\"dark2\"] = x\n x = self.dark3(x)\n outputs[\"dark3\"] = x\n x = self.dark4(x)\n outputs[\"dark4\"] = x\n x = self.dark5(x)\n outputs[\"dark5\"] = x\n return {k: v for k, v in outputs.items() if k in self.out_features}\n"
},
{
"path": "detector/YOLOX/yolox/models/losses.py",
"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport torch\nimport torch.nn as nn\n\n\nclass IOUloss(nn.Module):\n def __init__(self, reduction=\"none\", loss_type=\"iou\"):\n super(IOUloss, self).__init__()\n self.reduction = reduction\n self.loss_type = loss_type\n\n def forward(self, pred, target):\n assert pred.shape[0] == target.shape[0]\n\n pred = pred.view(-1, 4)\n target = target.view(-1, 4)\n tl = torch.max(\n (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)\n )\n br = torch.min(\n (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)\n )\n\n area_p = torch.prod(pred[:, 2:], 1)\n area_g = torch.prod(target[:, 2:], 1)\n\n en = (tl < br).type(tl.type()).prod(dim=1)\n area_i = torch.prod(br - tl, 1) * en\n iou = (area_i) / (area_p + area_g - area_i + 1e-16)\n\n if self.loss_type == \"iou\":\n loss = 1 - iou ** 2\n elif self.loss_type == \"giou\":\n c_tl = torch.min(\n (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)\n )\n c_br = torch.max(\n (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)\n )\n area_c = torch.prod(c_br - c_tl, 1)\n giou = iou - (area_c - area_i) / area_c.clamp(1e-16)\n loss = 1 - giou.clamp(min=-1.0, max=1.0)\n\n if self.reduction == \"mean\":\n loss = loss.mean()\n elif self.reduction == \"sum\":\n loss = loss.sum()\n\n return loss\n"
},
{
"path": "detector/YOLOX/yolox/models/network_blocks.py",
"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport torch\nimport torch.nn as nn\n\n\nclass SiLU(nn.Module):\n \"\"\"export-friendly version of nn.SiLU()\"\"\"\n\n @staticmethod\n def forward(x):\n return x * torch.sigmoid(x)\n\n\ndef get_activation(name=\"silu\", inplace=True):\n if name == \"silu\":\n module = nn.SiLU(inplace=inplace)\n elif name == \"relu\":\n module = nn.ReLU(inplace=inplace)\n elif name == \"lrelu\":\n module = nn.LeakyReLU(0.1, inplace=inplace)\n else:\n raise AttributeError(\"Unsupported act type: {}\".format(name))\n return module\n\n\nclass BaseConv(nn.Module):\n \"\"\"A Conv2d -> Batchnorm -> silu/leaky relu block\"\"\"\n\n def __init__(self, in_channels, out_channels, ksize, stride, groups=1, bias=False, act=\"silu\"):\n super().__init__()\n # same padding\n pad = (ksize - 1) // 2\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n kernel_size=ksize,\n stride=stride,\n padding=pad,\n groups=groups,\n bias=bias,\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.act = get_activation(act, inplace=True)\n\n def forward(self, x):\n return self.act(self.bn(self.conv(x)))\n\n def fuseforward(self, x):\n return self.act(self.conv(x))\n\n\nclass DWConv(nn.Module):\n \"\"\"Depthwise Conv + Conv\"\"\"\n def __init__(self, in_channels, out_channels, ksize, stride=1, act=\"silu\"):\n super().__init__()\n self.dconv = BaseConv(\n in_channels, in_channels, ksize=ksize,\n stride=stride, groups=in_channels, act=act\n )\n self.pconv = BaseConv(\n in_channels, out_channels, ksize=1,\n stride=1, groups=1, act=act\n )\n\n def forward(self, x):\n x = self.dconv(x)\n return self.pconv(x)\n\n\nclass Bottleneck(nn.Module):\n # Standard bottleneck\n def __init__(\n self, in_channels, out_channels, shortcut=True,\n expansion=0.5, depthwise=False, act=\"silu\"\n ):\n super().__init__()\n hidden_channels = int(out_channels * expansion)\n Conv = DWConv if depthwise else BaseConv\n self.conv1 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=act)\n self.conv2 = Conv(hidden_channels, out_channels, 3, stride=1, act=act)\n self.use_add = shortcut and in_channels == out_channels\n\n def forward(self, x):\n y = self.conv2(self.conv1(x))\n if self.use_add:\n y = y + x\n return y\n\n\nclass ResLayer(nn.Module):\n \"Residual layer with `in_channels` inputs.\"\n def __init__(self, in_channels: int):\n super().__init__()\n mid_channels = in_channels // 2\n self.layer1 = BaseConv(in_channels, mid_channels, ksize=1, stride=1, act=\"lrelu\")\n self.layer2 = BaseConv(mid_channels, in_channels, ksize=3, stride=1, act=\"lrelu\")\n\n def forward(self, x):\n out = self.layer2(self.layer1(x))\n return x + out\n\n\nclass SPPBottleneck(nn.Module):\n \"\"\"Spatial pyramid pooling layer used in YOLOv3-SPP\"\"\"\n def __init__(self, in_channels, out_channels, kernel_sizes=(5, 9, 13), activation=\"silu\"):\n super().__init__()\n hidden_channels = in_channels // 2\n self.conv1 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=activation)\n self.m = nn.ModuleList(\n [nn.MaxPool2d(kernel_size=ks, stride=1, padding=ks // 2) for ks in kernel_sizes]\n )\n conv2_channels = hidden_channels * (len(kernel_sizes) + 1)\n self.conv2 = BaseConv(conv2_channels, out_channels, 1, stride=1, act=activation)\n\n def forward(self, x):\n x = self.conv1(x)\n x = torch.cat([x] + [m(x) for m in self.m], dim=1)\n x = self.conv2(x)\n return x\n\n\nclass CSPLayer(nn.Module):\n \"\"\"C3 in yolov5, CSP Bottleneck with 3 convolutions\"\"\"\n\n def __init__(\n self, in_channels, out_channels, n=1,\n shortcut=True, expansion=0.5, depthwise=False, act=\"silu\"\n ):\n \"\"\"\n Args:\n in_channels (int): input channels.\n out_channels (int): output channels.\n n (int): number of Bottlenecks. Default value: 1.\n \"\"\"\n # ch_in, ch_out, number, shortcut, groups, expansion\n super().__init__()\n hidden_channels = int(out_channels * expansion) # hidden channels\n self.conv1 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=act)\n self.conv2 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=act)\n self.conv3 = BaseConv(2 * hidden_channels, out_channels, 1, stride=1, act=act)\n module_list = [\n Bottleneck(hidden_channels, hidden_channels, shortcut, 1.0, depthwise, act=act)\n for _ in range(n)\n ]\n self.m = nn.Sequential(*module_list)\n\n def forward(self, x):\n x_1 = self.conv1(x)\n x_2 = self.conv2(x)\n x_1 = self.m(x_1)\n x = torch.cat((x_1, x_2), dim=1)\n return self.conv3(x)\n\n\nclass Focus(nn.Module):\n \"\"\"Focus width and height information into channel space.\"\"\"\n\n def __init__(self, in_channels, out_channels, ksize=1, stride=1, act=\"silu\"):\n super().__init__()\n self.conv = BaseConv(in_channels * 4, out_channels, ksize, stride, act=act)\n\n def forward(self, x):\n # shape of x (b,c,w,h) -> y(b,4c,w/2,h/2)\n patch_top_left = x[..., ::2, ::2]\n patch_top_right = x[..., ::2, 1::2]\n patch_bot_left = x[..., 1::2, ::2]\n patch_bot_right = x[..., 1::2, 1::2]\n x = torch.cat(\n (patch_top_left, patch_bot_left, patch_top_right, patch_bot_right,), dim=1,\n )\n return self.conv(x)\n"
},
{
"path": "detector/YOLOX/yolox/models/yolo_fpn.py",
"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport torch\nimport torch.nn as nn\n\nfrom .darknet import Darknet\nfrom .network_blocks import BaseConv\n\n\nclass YOLOFPN(nn.Module):\n \"\"\"\n YOLOFPN module. Darknet 53 is the default backbone of this model.\n \"\"\"\n\n def __init__(\n self, depth=53, in_features=[\"dark3\", \"dark4\", \"dark5\"],\n ):\n super().__init__()\n\n self.backbone = Darknet(depth)\n self.in_features = in_features\n\n # out 1\n self.out1_cbl = self._make_cbl(512, 256, 1)\n self.out1 = self._make_embedding([256, 512], 512 + 256)\n\n # out 2\n self.out2_cbl = self._make_cbl(256, 128, 1)\n self.out2 = self._make_embedding([128, 256], 256 + 128)\n\n # upsample\n self.upsample = nn.Upsample(scale_factor=2, mode=\"nearest\")\n\n def _make_cbl(self, _in, _out, ks):\n return BaseConv(_in, _out, ks, stride=1, act=\"lrelu\")\n\n def _make_embedding(self, filters_list, in_filters):\n m = nn.Sequential(\n *[\n self._make_cbl(in_filters, filters_list[0], 1),\n self._make_cbl(filters_list[0], filters_list[1], 3),\n\n self._make_cbl(filters_list[1], filters_list[0], 1),\n\n self._make_cbl(filters_list[0], filters_list[1], 3),\n self._make_cbl(filters_list[1], filters_list[0], 1),\n ]\n )\n return m\n\n def load_pretrained_model(self, filename=\"./weights/darknet53.mix.pth\"):\n with open(filename, \"rb\") as f:\n state_dict = torch.load(f, map_location=\"cpu\")\n print(\"loading pretrained weights...\")\n self.backbone.load_state_dict(state_dict)\n\n def forward(self, inputs):\n \"\"\"\n Args:\n inputs (Tensor): input image.\n\n Returns:\n Tuple[Tensor]: FPN output features..\n \"\"\"\n # backbone\n out_features = self.backbone(inputs)\n x2, x1, x0 = [out_features[f] for f in self.in_features]\n\n # yolo branch 1\n x1_in = self.out1_cbl(x0)\n x1_in = self.upsample(x1_in)\n x1_in = torch.cat([x1_in, x1], 1)\n out_dark4 = self.out1(x1_in)\n\n # yolo branch 2\n x2_in = self.out2_cbl(out_dark4)\n x2_in = self.upsample(x2_in)\n x2_in = torch.cat([x2_in, x2], 1)\n out_dark3 = self.out2(x2_in)\n\n outputs = (out_dark3, out_dark4, x0)\n return outputs\n"
},
{
"path": "detector/YOLOX/yolox/models/yolo_head.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport math\nfrom loguru import logger\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom yolox.utils import bboxes_iou\n\nfrom .losses import IOUloss\nfrom .network_blocks import BaseConv, DWConv\n\n\nclass YOLOXHead(nn.Module):\n def __init__(\n self, num_classes, width=1.0, strides=[8, 16, 32],\n in_channels=[256, 512, 1024], act=\"silu\", depthwise=False\n ):\n \"\"\"\n Args:\n act (str): activation type of conv. Defalut value: \"silu\".\n depthwise (bool): wheather apply depthwise conv in conv branch. Defalut value: False.\n \"\"\"\n super().__init__()\n\n self.n_anchors = 1\n self.num_classes = num_classes\n self.decode_in_inference = True # for deploy, set to False\n\n self.cls_convs = nn.ModuleList()\n self.reg_convs = nn.ModuleList()\n self.cls_preds = nn.ModuleList()\n self.reg_preds = nn.ModuleList()\n self.obj_preds = nn.ModuleList()\n self.stems = nn.ModuleList()\n Conv = DWConv if depthwise else BaseConv\n\n for i in range(len(in_channels)):\n self.stems.append(\n BaseConv(\n in_channels=int(in_channels[i] * width),\n out_channels=int(256 * width),\n ksize=1,\n stride=1,\n act=act,\n )\n )\n self.cls_convs.append(\n nn.Sequential(\n *[\n Conv(\n in_channels=int(256 * width),\n out_channels=int(256 * width),\n ksize=3,\n stride=1,\n act=act,\n ),\n Conv(\n in_channels=int(256 * width),\n out_channels=int(256 * width),\n ksize=3,\n stride=1,\n act=act,\n ),\n ]\n )\n )\n self.reg_convs.append(\n nn.Sequential(\n *[\n Conv(\n in_channels=int(256 * width),\n out_channels=int(256 * width),\n ksize=3,\n stride=1,\n act=act,\n ),\n Conv(\n in_channels=int(256 * width),\n out_channels=int(256 * width),\n ksize=3,\n stride=1,\n act=act,\n ),\n ]\n )\n )\n self.cls_preds.append(\n nn.Conv2d(\n in_channels=int(256 * width),\n out_channels=self.n_anchors * self.num_classes,\n kernel_size=1,\n stride=1,\n padding=0,\n )\n )\n self.reg_preds.append(\n nn.Conv2d(\n in_channels=int(256 * width),\n out_channels=4,\n kernel_size=1,\n stride=1,\n padding=0,\n )\n )\n self.obj_preds.append(\n nn.Conv2d(\n in_channels=int(256 * width),\n out_channels=self.n_anchors * 1,\n kernel_size=1,\n stride=1,\n padding=0,\n )\n )\n\n self.use_l1 = False\n self.l1_loss = nn.L1Loss(reduction=\"none\")\n self.bcewithlog_loss = nn.BCEWithLogitsLoss(reduction=\"none\")\n self.iou_loss = IOUloss(reduction=\"none\")\n self.strides = strides\n self.grids = [torch.zeros(1)] * len(in_channels)\n self.expanded_strides = [None] * len(in_channels)\n\n def initialize_biases(self, prior_prob):\n for conv in self.cls_preds:\n b = conv.bias.view(self.n_anchors, -1)\n b.data.fill_(-math.log((1 - prior_prob) / prior_prob))\n conv.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)\n\n for conv in self.obj_preds:\n b = conv.bias.view(self.n_anchors, -1)\n b.data.fill_(-math.log((1 - prior_prob) / prior_prob))\n conv.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)\n\n def forward(self, xin, labels=None, imgs=None):\n outputs = []\n origin_preds = []\n x_shifts = []\n y_shifts = []\n expanded_strides = []\n\n for k, (cls_conv, reg_conv, stride_this_level, x) in enumerate(\n zip(self.cls_convs, self.reg_convs, self.strides, xin)\n ):\n x = self.stems[k](x)\n cls_x = x\n reg_x = x\n\n cls_feat = cls_conv(cls_x)\n cls_output = self.cls_preds[k](cls_feat)\n\n reg_feat = reg_conv(reg_x)\n reg_output = self.reg_preds[k](reg_feat)\n obj_output = self.obj_preds[k](reg_feat)\n\n if self.training:\n output = torch.cat([reg_output, obj_output, cls_output], 1)\n output, grid = self.get_output_and_grid(output, k, stride_this_level, xin[0].type())\n x_shifts.append(grid[:, :, 0])\n y_shifts.append(grid[:, :, 1])\n expanded_strides.append(\n torch.zeros(1, grid.shape[1]).fill_(stride_this_level).type_as(xin[0])\n )\n if self.use_l1:\n batch_size = reg_output.shape[0]\n hsize, wsize = reg_output.shape[-2:]\n reg_output = reg_output.view(batch_size, self.n_anchors, 4, hsize, wsize)\n reg_output = (\n reg_output.permute(0, 1, 3, 4, 2)\n .reshape(batch_size, -1, 4)\n )\n origin_preds.append(reg_output.clone())\n\n else:\n output = torch.cat([reg_output, obj_output.sigmoid(), cls_output.sigmoid()], 1)\n\n outputs.append(output)\n\n if self.training:\n return self.get_losses(\n imgs, x_shifts, y_shifts, expanded_strides, labels,\n torch.cat(outputs, 1), origin_preds, dtype=xin[0].dtype\n )\n else:\n self.hw = [x.shape[-2:] for x in outputs]\n # [batch, n_anchors_all, 85]\n outputs = torch.cat([x.flatten(start_dim=2) for x in outputs], dim=2).permute(0, 2, 1)\n if self.decode_in_inference:\n return self.decode_outputs(outputs, dtype=xin[0].type())\n else:\n return outputs\n\n def get_output_and_grid(self, output, k, stride, dtype):\n grid = self.grids[k]\n\n batch_size = output.shape[0]\n n_ch = 5 + self.num_classes\n hsize, wsize = output.shape[-2:]\n if grid.shape[2:4] != output.shape[2:4]:\n yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])\n grid = torch.stack((xv, yv), 2).view(1, 1, hsize, wsize, 2).type(dtype)\n self.grids[k] = grid\n\n output = output.view(batch_size, self.n_anchors, n_ch, hsize, wsize)\n output = (\n output.permute(0, 1, 3, 4, 2)\n .reshape(batch_size, self.n_anchors * hsize * wsize, -1)\n )\n grid = grid.view(1, -1, 2)\n output[..., :2] = (output[..., :2] + grid) * stride\n output[..., 2:4] = torch.exp(output[..., 2:4]) * stride\n return output, grid\n\n def decode_outputs(self, outputs, dtype):\n grids = []\n strides = []\n for (hsize, wsize), stride in zip(self.hw, self.strides):\n yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])\n grid = torch.stack((xv, yv), 2).view(1, -1, 2)\n grids.append(grid)\n shape = grid.shape[:2]\n strides.append(torch.full((*shape, 1), stride))\n\n grids = torch.cat(grids, dim=1).type(dtype)\n strides = torch.cat(strides, dim=1).type(dtype)\n\n outputs[..., :2] = (outputs[..., :2] + grids) * strides\n outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides\n return outputs\n\n def get_losses(\n self, imgs, x_shifts, y_shifts, expanded_strides, labels, outputs, origin_preds, dtype,\n ):\n bbox_preds = outputs[:, :, :4] # [batch, n_anchors_all, 4]\n obj_preds = outputs[:, :, 4].unsqueeze(-1) # [batch, n_anchors_all, 1]\n cls_preds = outputs[:, :, 5:] # [batch, n_anchors_all, n_cls]\n\n # calculate targets\n mixup = labels.shape[2] > 5\n if mixup:\n label_cut = labels[..., :5]\n else:\n label_cut = labels\n nlabel = (label_cut.sum(dim=2) > 0).sum(dim=1) # number of objects\n\n total_num_anchors = outputs.shape[1]\n x_shifts = torch.cat(x_shifts, 1) # [1, n_anchors_all]\n y_shifts = torch.cat(y_shifts, 1) # [1, n_anchors_all]\n expanded_strides = torch.cat(expanded_strides, 1)\n if self.use_l1:\n origin_preds = torch.cat(origin_preds, 1)\n\n cls_targets = []\n reg_targets = []\n l1_targets = []\n obj_targets = []\n fg_masks = []\n\n num_fg = 0.0\n num_gts = 0.0\n\n for batch_idx in range(outputs.shape[0]):\n num_gt = int(nlabel[batch_idx])\n num_gts += num_gt\n if num_gt == 0:\n cls_target = outputs.new_zeros((0, self.num_classes))\n reg_target = outputs.new_zeros((0, 4))\n l1_target = outputs.new_zeros((0, 4))\n obj_target = outputs.new_zeros((total_num_anchors, 1))\n fg_mask = outputs.new_zeros(total_num_anchors).bool()\n else:\n gt_bboxes_per_image = labels[batch_idx, :num_gt, 1:5]\n gt_classes = labels[batch_idx, :num_gt, 0]\n bboxes_preds_per_image = bbox_preds[batch_idx]\n\n try:\n gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg_img = self.get_assignments( # noqa\n batch_idx, num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes,\n bboxes_preds_per_image, expanded_strides, x_shifts, y_shifts,\n cls_preds, bbox_preds, obj_preds, labels, imgs,\n )\n except RuntimeError:\n logger.error(\n \"OOM RuntimeError is raised due to the huge memory cost during label assignment. \\\n CPU mode is applied in this batch. If you want to avoid this issue, \\\n try to reduce the batch size or image size.\"\n )\n torch.cuda.empty_cache()\n gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg_img = self.get_assignments( # noqa\n batch_idx, num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes,\n bboxes_preds_per_image, expanded_strides, x_shifts, y_shifts,\n cls_preds, bbox_preds, obj_preds, labels, imgs, \"cpu\",\n )\n\n torch.cuda.empty_cache()\n num_fg += num_fg_img\n\n cls_target = F.one_hot(\n gt_matched_classes.to(torch.int64), self.num_classes\n ) * pred_ious_this_matching.unsqueeze(-1)\n obj_target = fg_mask.unsqueeze(-1)\n reg_target = gt_bboxes_per_image[matched_gt_inds]\n if self.use_l1:\n l1_target = self.get_l1_target(\n outputs.new_zeros((num_fg_img, 4)),\n gt_bboxes_per_image[matched_gt_inds],\n expanded_strides[0][fg_mask],\n x_shifts=x_shifts[0][fg_mask],\n y_shifts=y_shifts[0][fg_mask],\n )\n\n cls_targets.append(cls_target)\n reg_targets.append(reg_target)\n obj_targets.append(obj_target.to(dtype))\n fg_masks.append(fg_mask)\n if self.use_l1:\n l1_targets.append(l1_target)\n\n cls_targets = torch.cat(cls_targets, 0)\n reg_targets = torch.cat(reg_targets, 0)\n obj_targets = torch.cat(obj_targets, 0)\n fg_masks = torch.cat(fg_masks, 0)\n if self.use_l1:\n l1_targets = torch.cat(l1_targets, 0)\n\n num_fg = max(num_fg, 1)\n loss_iou = (self.iou_loss(bbox_preds.view(-1, 4)[fg_masks], reg_targets)).sum() / num_fg\n loss_obj = (self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets)).sum() / num_fg\n loss_cls = (\n self.bcewithlog_loss(cls_preds.view(-1, self.num_classes)[fg_masks], cls_targets)\n ).sum() / num_fg\n if self.use_l1:\n loss_l1 = (self.l1_loss(origin_preds.view(-1, 4)[fg_masks], l1_targets)).sum() / num_fg\n else:\n loss_l1 = 0.0\n\n reg_weight = 5.0\n loss = reg_weight * loss_iou + loss_obj + loss_cls + loss_l1\n\n return loss, reg_weight * loss_iou, loss_obj, loss_cls, loss_l1, num_fg / max(num_gts, 1)\n\n def get_l1_target(self, l1_target, gt, stride, x_shifts, y_shifts, eps=1e-8):\n l1_target[:, 0] = gt[:, 0] / stride - x_shifts\n l1_target[:, 1] = gt[:, 1] / stride - y_shifts\n l1_target[:, 2] = torch.log(gt[:, 2] / stride + eps)\n l1_target[:, 3] = torch.log(gt[:, 3] / stride + eps)\n return l1_target\n\n @torch.no_grad()\n def get_assignments(\n self, batch_idx, num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes,\n bboxes_preds_per_image, expanded_strides, x_shifts, y_shifts,\n cls_preds, bbox_preds, obj_preds, labels, imgs, mode=\"gpu\",\n ):\n\n if mode == \"cpu\":\n print(\"------------CPU Mode for This Batch-------------\")\n gt_bboxes_per_image = gt_bboxes_per_image.cpu().float()\n bboxes_preds_per_image = bboxes_preds_per_image.cpu().float()\n gt_classes = gt_classes.cpu().float()\n expanded_strides = expanded_strides.cpu().float()\n x_shifts = x_shifts.cpu()\n y_shifts = y_shifts.cpu()\n\n fg_mask, is_in_boxes_and_center = self.get_in_boxes_info(\n gt_bboxes_per_image, expanded_strides, x_shifts, y_shifts, total_num_anchors, num_gt,\n )\n\n bboxes_preds_per_image = bboxes_preds_per_image[fg_mask]\n cls_preds_ = cls_preds[batch_idx][fg_mask]\n obj_preds_ = obj_preds[batch_idx][fg_mask]\n num_in_boxes_anchor = bboxes_preds_per_image.shape[0]\n\n if mode == \"cpu\":\n gt_bboxes_per_image = gt_bboxes_per_image.cpu()\n bboxes_preds_per_image = bboxes_preds_per_image.cpu()\n\n pair_wise_ious = bboxes_iou(\n gt_bboxes_per_image, bboxes_preds_per_image, False\n )\n\n gt_cls_per_image = (\n F.one_hot(gt_classes.to(torch.int64), self.num_classes).float()\n .unsqueeze(1).repeat(1, num_in_boxes_anchor, 1)\n )\n pair_wise_ious_loss = -torch.log(pair_wise_ious + 1e-8)\n\n if mode == \"cpu\":\n cls_preds_, obj_preds_ = cls_preds_.cpu(), obj_preds_.cpu()\n\n cls_preds_ = (\n cls_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()\n * obj_preds_.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()\n )\n pair_wise_cls_loss = F.binary_cross_entropy(\n cls_preds_.sqrt_(), gt_cls_per_image, reduction=\"none\"\n ).sum(-1)\n del cls_preds_\n\n cost = (\n pair_wise_cls_loss\n + 3.0 * pair_wise_ious_loss\n + 100000.0 * (~is_in_boxes_and_center)\n )\n\n (\n num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds\n ) = self.dynamic_k_matching(cost, pair_wise_ious, gt_classes, num_gt, fg_mask)\n del pair_wise_cls_loss, cost, pair_wise_ious, pair_wise_ious_loss\n\n if mode == \"cpu\":\n gt_matched_classes = gt_matched_classes.cuda()\n fg_mask = fg_mask.cuda()\n pred_ious_this_matching = pred_ious_this_matching.cuda()\n matched_gt_inds = matched_gt_inds.cuda()\n\n return gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg\n\n def get_in_boxes_info(\n self, gt_bboxes_per_image, expanded_strides, x_shifts, y_shifts, total_num_anchors, num_gt,\n ):\n expanded_strides_per_image = expanded_strides[0]\n x_shifts_per_image = x_shifts[0] * expanded_strides_per_image\n y_shifts_per_image = y_shifts[0] * expanded_strides_per_image\n x_centers_per_image = (\n (x_shifts_per_image + 0.5 * expanded_strides_per_image)\n .unsqueeze(0)\n .repeat(num_gt, 1)\n ) # [n_anchor] -> [n_gt, n_anchor]\n y_centers_per_image = (\n (y_shifts_per_image + 0.5 * expanded_strides_per_image)\n .unsqueeze(0)\n .repeat(num_gt, 1)\n )\n\n gt_bboxes_per_image_l = (\n (gt_bboxes_per_image[:, 0] - 0.5 * gt_bboxes_per_image[:, 2])\n .unsqueeze(1)\n .repeat(1, total_num_anchors)\n )\n gt_bboxes_per_image_r = (\n (gt_bboxes_per_image[:, 0] + 0.5 * gt_bboxes_per_image[:, 2])\n .unsqueeze(1)\n .repeat(1, total_num_anchors)\n )\n gt_bboxes_per_image_t = (\n (gt_bboxes_per_image[:, 1] - 0.5 * gt_bboxes_per_image[:, 3])\n .unsqueeze(1)\n .repeat(1, total_num_anchors)\n )\n gt_bboxes_per_image_b = (\n (gt_bboxes_per_image[:, 1] + 0.5 * gt_bboxes_per_image[:, 3])\n .unsqueeze(1)\n .repeat(1, total_num_anchors)\n )\n\n b_l = x_centers_per_image - gt_bboxes_per_image_l\n b_r = gt_bboxes_per_image_r - x_centers_per_image\n b_t = y_centers_per_image - gt_bboxes_per_image_t\n b_b = gt_bboxes_per_image_b - y_centers_per_image\n bbox_deltas = torch.stack([b_l, b_t, b_r, b_b], 2)\n\n is_in_boxes = bbox_deltas.min(dim=-1).values > 0.0\n is_in_boxes_all = is_in_boxes.sum(dim=0) > 0\n # in fixed center\n\n center_radius = 2.5\n\n gt_bboxes_per_image_l = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(\n 1, total_num_anchors\n ) - center_radius * expanded_strides_per_image.unsqueeze(0)\n gt_bboxes_per_image_r = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(\n 1, total_num_anchors\n ) + center_radius * expanded_strides_per_image.unsqueeze(0)\n gt_bboxes_per_image_t = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(\n 1, total_num_anchors\n ) - center_radius * expanded_strides_per_image.unsqueeze(0)\n gt_bboxes_per_image_b = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(\n 1, total_num_anchors\n ) + center_radius * expanded_strides_per_image.unsqueeze(0)\n\n c_l = x_centers_per_image - gt_bboxes_per_image_l\n c_r = gt_bboxes_per_image_r - x_centers_per_image\n c_t = y_centers_per_image - gt_bboxes_per_image_t\n c_b = gt_bboxes_per_image_b - y_centers_per_image\n center_deltas = torch.stack([c_l, c_t, c_r, c_b], 2)\n is_in_centers = center_deltas.min(dim=-1).values > 0.0\n is_in_centers_all = is_in_centers.sum(dim=0) > 0\n\n # in boxes and in centers\n is_in_boxes_anchor = is_in_boxes_all | is_in_centers_all\n\n is_in_boxes_and_center = (\n is_in_boxes[:, is_in_boxes_anchor] & is_in_centers[:, is_in_boxes_anchor]\n )\n return is_in_boxes_anchor, is_in_boxes_and_center\n\n def dynamic_k_matching(self, cost, pair_wise_ious, gt_classes, num_gt, fg_mask):\n # Dynamic K\n # ---------------------------------------------------------------\n matching_matrix = torch.zeros_like(cost)\n\n ious_in_boxes_matrix = pair_wise_ious\n n_candidate_k = 10\n topk_ious, _ = torch.topk(ious_in_boxes_matrix, n_candidate_k, dim=1)\n dynamic_ks = torch.clamp(topk_ious.sum(1).int(), min=1)\n for gt_idx in range(num_gt):\n _, pos_idx = torch.topk(\n cost[gt_idx], k=dynamic_ks[gt_idx].item(), largest=False\n )\n matching_matrix[gt_idx][pos_idx] = 1.0\n\n del topk_ious, dynamic_ks, pos_idx\n\n anchor_matching_gt = matching_matrix.sum(0)\n if (anchor_matching_gt > 1).sum() > 0:\n cost_min, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)\n matching_matrix[:, anchor_matching_gt > 1] *= 0.0\n matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.0\n fg_mask_inboxes = matching_matrix.sum(0) > 0.0\n num_fg = fg_mask_inboxes.sum().item()\n\n fg_mask[fg_mask.clone()] = fg_mask_inboxes\n\n matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)\n gt_matched_classes = gt_classes[matched_gt_inds]\n\n pred_ious_this_matching = (matching_matrix * pair_wise_ious).sum(0)[fg_mask_inboxes]\n return num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds\n"
},
{
"path": "detector/YOLOX/yolox/models/yolo_pafpn.py",
"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport torch\nimport torch.nn as nn\n\nfrom .darknet import CSPDarknet\nfrom .network_blocks import BaseConv, CSPLayer, DWConv\n\n\nclass YOLOPAFPN(nn.Module):\n \"\"\"\n YOLOv3 model. Darknet 53 is the default backbone of this model.\n \"\"\"\n\n def __init__(\n self, depth=1.0, width=1.0, in_features=(\"dark3\", \"dark4\", \"dark5\"),\n in_channels=[256, 512, 1024], depthwise=False, act=\"silu\",\n ):\n super().__init__()\n self.backbone = CSPDarknet(depth, width, depthwise=depthwise, act=act)\n self.in_features = in_features\n self.in_channels = in_channels\n Conv = DWConv if depthwise else BaseConv\n\n self.upsample = nn.Upsample(scale_factor=2, mode=\"nearest\")\n self.lateral_conv0 = BaseConv(\n int(in_channels[2] * width), int(in_channels[1] * width), 1, 1, act=act\n )\n self.C3_p4 = CSPLayer(\n int(2 * in_channels[1] * width),\n int(in_channels[1] * width),\n round(3 * depth),\n False,\n depthwise=depthwise,\n act=act,\n ) # cat\n\n self.reduce_conv1 = BaseConv(\n int(in_channels[1] * width), int(in_channels[0] * width), 1, 1, act=act\n )\n self.C3_p3 = CSPLayer(\n int(2 * in_channels[0] * width),\n int(in_channels[0] * width),\n round(3 * depth),\n False,\n depthwise=depthwise,\n act=act,\n )\n\n # bottom-up conv\n self.bu_conv2 = Conv(\n int(in_channels[0] * width), int(in_channels[0] * width), 3, 2, act=act\n )\n self.C3_n3 = CSPLayer(\n int(2 * in_channels[0] * width),\n int(in_channels[1] * width),\n round(3 * depth),\n False,\n depthwise=depthwise,\n act=act,\n )\n\n # bottom-up conv\n self.bu_conv1 = Conv(\n int(in_channels[1] * width), int(in_channels[1] * width), 3, 2, act=act\n )\n self.C3_n4 = CSPLayer(\n int(2 * in_channels[1] * width),\n int(in_channels[2] * width),\n round(3 * depth),\n False,\n depthwise=depthwise,\n act=act,\n )\n\n def forward(self, input):\n \"\"\"\n Args:\n inputs: input images.\n\n Returns:\n Tuple[Tensor]: FPN feature.\n \"\"\"\n\n # backbone\n out_features = self.backbone(input)\n features = [out_features[f] for f in self.in_features]\n [x2, x1, x0] = features\n\n fpn_out0 = self.lateral_conv0(x0) # 1024->512/32\n f_out0 = self.upsample(fpn_out0) # 512/16\n f_out0 = torch.cat([f_out0, x1], 1) # 512->1024/16\n f_out0 = self.C3_p4(f_out0) # 1024->512/16\n\n fpn_out1 = self.reduce_conv1(f_out0) # 512->256/16\n f_out1 = self.upsample(fpn_out1) # 256/8\n f_out1 = torch.cat([f_out1, x2], 1) # 256->512/8\n pan_out2 = self.C3_p3(f_out1) # 512->256/8\n\n p_out1 = self.bu_conv2(pan_out2) # 256->256/16\n p_out1 = torch.cat([p_out1, fpn_out1], 1) # 256->512/16\n pan_out1 = self.C3_n3(p_out1) # 512->512/16\n\n p_out0 = self.bu_conv1(pan_out1) # 512->512/32\n p_out0 = torch.cat([p_out0, fpn_out0], 1) # 512->1024/32\n pan_out0 = self.C3_n4(p_out0) # 1024->1024/32\n\n outputs = (pan_out2, pan_out1, pan_out0)\n return outputs\n"
},
{
"path": "detector/YOLOX/yolox/models/yolox.py",
"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport torch.nn as nn\n\nfrom .yolo_head import YOLOXHead\nfrom .yolo_pafpn import YOLOPAFPN\n\n\nclass YOLOX(nn.Module):\n \"\"\"\n YOLOX model module. The module list is defined by create_yolov3_modules function.\n The network returns loss values from three YOLO layers during training\n and detection results during test.\n \"\"\"\n\n def __init__(self, backbone=None, head=None):\n super().__init__()\n if backbone is None:\n backbone = YOLOPAFPN()\n if head is None:\n head = YOLOXHead(80)\n\n self.backbone = backbone\n self.head = head\n\n def forward(self, x, targets=None):\n # fpn output content features of [dark3, dark4, dark5]\n fpn_outs = self.backbone(x)\n\n if self.training:\n assert targets is not None\n loss, iou_loss, conf_loss, cls_loss, l1_loss, num_fg = self.head(\n fpn_outs, targets, x\n )\n outputs = {\n \"total_loss\": loss,\n \"iou_loss\": iou_loss,\n \"l1_loss\": l1_loss,\n \"conf_loss\": conf_loss,\n \"cls_loss\": cls_loss,\n \"num_fg\": num_fg,\n }\n else:\n outputs = self.head(fpn_outs)\n\n return outputs\n"
},
{
"path": "detector/YOLOX/yolox/utils/__init__.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nfrom .allreduce_norm import *\nfrom .boxes import *\nfrom .checkpoint import load_ckpt, save_checkpoint\nfrom .demo_utils import *\nfrom .dist import *\nfrom .ema import ModelEMA\nfrom .logger import setup_logger\nfrom .lr_scheduler import LRScheduler\nfrom .metric import *\nfrom .model_utils import *\nfrom .setup_env import *\nfrom .visualize import *\n"
},
{
"path": "detector/YOLOX/yolox/utils/allreduce_norm.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport pickle\nfrom collections import OrderedDict\n\nimport torch\nfrom torch import distributed as dist\nfrom torch import nn\n\nfrom .dist import _get_global_gloo_group, get_world_size\n\nASYNC_NORM = (\n nn.BatchNorm1d,\n nn.BatchNorm2d,\n nn.BatchNorm3d,\n nn.InstanceNorm1d,\n nn.InstanceNorm2d,\n nn.InstanceNorm3d,\n)\n\n__all__ = [\n \"get_async_norm_states\", \"pyobj2tensor\", \"tensor2pyobj\", \"all_reduce\", \"all_reduce_norm\"\n]\n\n\ndef get_async_norm_states(module):\n async_norm_states = OrderedDict()\n for name, child in module.named_modules():\n if isinstance(child, ASYNC_NORM):\n for k, v in child.state_dict().items():\n async_norm_states[\".\".join([name, k])] = v\n return async_norm_states\n\n\ndef pyobj2tensor(pyobj, device=\"cuda\"):\n \"\"\"serialize picklable python object to tensor\"\"\"\n storage = torch.ByteStorage.from_buffer(pickle.dumps(pyobj))\n return torch.ByteTensor(storage).to(device=device)\n\n\ndef tensor2pyobj(tensor):\n \"\"\"deserialize tensor to picklable python object\"\"\"\n return pickle.loads(tensor.cpu().numpy().tobytes())\n\n\ndef _get_reduce_op(op_name):\n return {\n \"sum\": dist.ReduceOp.SUM,\n \"mean\": dist.ReduceOp.SUM,\n }[op_name.lower()]\n\n\ndef all_reduce(py_dict, op=\"sum\", group=None):\n \"\"\"\n Apply all reduce function for python dict object.\n NOTE: make sure that every py_dict has the same keys and values are in the same shape.\n\n Args:\n py_dict (dict): dict to apply all reduce op.\n op (str): operator, could be \"sum\" or \"mean\".\n \"\"\"\n world_size = get_world_size()\n if world_size == 1:\n return py_dict\n if group is None:\n group = _get_global_gloo_group()\n if dist.get_world_size(group) == 1:\n return py_dict\n\n # all reduce logic across different devices.\n py_key = list(py_dict.keys())\n py_key_tensor = pyobj2tensor(py_key)\n dist.broadcast(py_key_tensor, src=0)\n py_key = tensor2pyobj(py_key_tensor)\n\n tensor_shapes = [py_dict[k].shape for k in py_key]\n tensor_numels = [py_dict[k].numel() for k in py_key]\n\n flatten_tensor = torch.cat([py_dict[k].flatten() for k in py_key])\n dist.all_reduce(flatten_tensor, op=_get_reduce_op(op))\n if op == \"mean\":\n flatten_tensor /= world_size\n\n split_tensors = [\n x.reshape(shape) for x, shape in zip(\n torch.split(flatten_tensor, tensor_numels), tensor_shapes\n )\n ]\n return OrderedDict({k: v for k, v in zip(py_key, split_tensors)})\n\n\ndef all_reduce_norm(module):\n \"\"\"\n All reduce norm statistics in different devices.\n \"\"\"\n states = get_async_norm_states(module)\n states = all_reduce(states, op=\"mean\")\n module.load_state_dict(states, strict=False)\n"
},
{
"path": "detector/YOLOX/yolox/utils/boxes.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport numpy as np\n\nimport torch\nimport torchvision\n\n__all__ = [\n \"filter_box\", \"postprocess\", \"bboxes_iou\", \"matrix_iou\",\n \"adjust_box_anns\", \"xyxy2xywh\",\n]\n\n\ndef filter_box(output, scale_range):\n \"\"\"\n output: (N, 5+class) shape\n \"\"\"\n min_scale, max_scale = scale_range\n w = output[:, 2] - output[:, 0]\n h = output[:, 3] - output[:, 1]\n keep = (w * h > min_scale * min_scale) & (w * h < max_scale * max_scale)\n return output[keep]\n\n\ndef postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45):\n box_corner = prediction.new(prediction.shape)\n box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2\n box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2\n box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2\n box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2\n prediction[:, :, :4] = box_corner[:, :, :4]\n\n output = [None for _ in range(len(prediction))]\n for i, image_pred in enumerate(prediction):\n\n # If none are remaining => process next image\n if not image_pred.size(0):\n continue\n # Get score and class with highest confidence\n class_conf, class_pred = torch.max(image_pred[:, 5: 5 + num_classes], 1, keepdim=True)\n\n conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_thre).squeeze()\n # _, conf_mask = torch.topk((image_pred[:, 4] * class_conf.squeeze()), 1000)\n # Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred)\n detections = torch.cat((image_pred[:, :5], class_conf, class_pred.float()), 1)\n detections = detections[conf_mask]\n if not detections.size(0):\n continue\n\n nms_out_index = torchvision.ops.batched_nms(\n detections[:, :4],\n detections[:, 4] * detections[:, 5],\n detections[:, 6],\n nms_thre,\n )\n detections = detections[nms_out_index]\n if output[i] is None:\n output[i] = detections\n else:\n output[i] = torch.cat((output[i], detections))\n\n return output\n\n\ndef bboxes_iou(bboxes_a, bboxes_b, xyxy=True):\n if bboxes_a.shape[1] != 4 or bboxes_b.shape[1] != 4:\n raise IndexError\n\n if xyxy:\n tl = torch.max(bboxes_a[:, None, :2], bboxes_b[:, :2])\n br = torch.min(bboxes_a[:, None, 2:], bboxes_b[:, 2:])\n area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)\n area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)\n else:\n tl = torch.max(\n (bboxes_a[:, None, :2] - bboxes_a[:, None, 2:] / 2),\n (bboxes_b[:, :2] - bboxes_b[:, 2:] / 2),\n )\n br = torch.min(\n (bboxes_a[:, None, :2] + bboxes_a[:, None, 2:] / 2),\n (bboxes_b[:, :2] + bboxes_b[:, 2:] / 2),\n )\n\n area_a = torch.prod(bboxes_a[:, 2:], 1)\n area_b = torch.prod(bboxes_b[:, 2:], 1)\n en = (tl < br).type(tl.type()).prod(dim=2)\n area_i = torch.prod(br - tl, 2) * en # * ((tl < br).all())\n return area_i / (area_a[:, None] + area_b - area_i)\n\n\ndef matrix_iou(a, b):\n \"\"\"\n return iou of a and b, numpy version for data augenmentation\n \"\"\"\n lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])\n rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])\n\n area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)\n area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)\n area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)\n return area_i / (area_a[:, np.newaxis] + area_b - area_i + 1e-12)\n\n\ndef adjust_box_anns(bbox, scale_ratio, padw, padh, w_max, h_max):\n bbox[:, 0::2] = np.clip(bbox[:, 0::2] * scale_ratio + padw, 0, w_max)\n bbox[:, 1::2] = np.clip(bbox[:, 1::2] * scale_ratio + padh, 0, h_max)\n return bbox\n\n\ndef xyxy2xywh(bboxes):\n bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]\n bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]\n return bboxes\n"
},
{
"path": "detector/YOLOX/yolox/utils/checkpoint.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\nimport os\nimport shutil\nfrom loguru import logger\n\nimport torch\n\n\ndef load_ckpt(model, ckpt):\n model_state_dict = model.state_dict()\n load_dict = {}\n for key_model, v in model_state_dict.items():\n if key_model not in ckpt:\n logger.warning(\n \"{} is not in the ckpt. Please double check and see if this is desired.\".format(\n key_model\n )\n )\n continue\n v_ckpt = ckpt[key_model]\n if v.shape != v_ckpt.shape:\n logger.warning(\n \"Shape of {} in checkpoint is {}, while shape of {} in model is {}.\".format(\n key_model, v_ckpt.shape, key_model, v.shape\n )\n )\n continue\n load_dict[key_model] = v_ckpt\n\n model.load_state_dict(load_dict, strict=False)\n return model\n\n\ndef save_checkpoint(state, is_best, save_dir, model_name=\"\"):\n if not os.path.exists(save_dir):\n os.makedirs(save_dir)\n filename = os.path.join(save_dir, model_name + \"_ckpt.pth.tar\")\n torch.save(state, filename)\n if is_best:\n best_filename = os.path.join(save_dir, \"best_ckpt.pth.tar\")\n shutil.copyfile(filename, best_filename)\n"
},
{
"path": "detector/YOLOX/yolox/utils/demo_utils.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport os\n\nimport numpy as np\n\n__all__ = [\"mkdir\", \"nms\", \"multiclass_nms\", \"demo_postprocess\"]\n\n\ndef mkdir(path):\n if not os.path.exists(path):\n os.makedirs(path)\n\n\ndef nms(boxes, scores, nms_thr):\n \"\"\"Single class NMS implemented in Numpy.\"\"\"\n x1 = boxes[:, 0]\n y1 = boxes[:, 1]\n x2 = boxes[:, 2]\n y2 = boxes[:, 3]\n\n areas = (x2 - x1 + 1) * (y2 - y1 + 1)\n order = scores.argsort()[::-1]\n\n keep = []\n while order.size > 0:\n i = order[0]\n keep.append(i)\n xx1 = np.maximum(x1[i], x1[order[1:]])\n yy1 = np.maximum(y1[i], y1[order[1:]])\n xx2 = np.minimum(x2[i], x2[order[1:]])\n yy2 = np.minimum(y2[i], y2[order[1:]])\n\n w = np.maximum(0.0, xx2 - xx1 + 1)\n h = np.maximum(0.0, yy2 - yy1 + 1)\n inter = w * h\n ovr = inter / (areas[i] + areas[order[1:]] - inter)\n\n inds = np.where(ovr <= nms_thr)[0]\n order = order[inds + 1]\n\n return keep\n\n\ndef multiclass_nms(boxes, scores, nms_thr, score_thr):\n \"\"\"Multiclass NMS implemented in Numpy\"\"\"\n final_dets = []\n num_classes = scores.shape[1]\n for cls_ind in range(num_classes):\n cls_scores = scores[:, cls_ind]\n valid_score_mask = cls_scores > score_thr\n if valid_score_mask.sum() == 0:\n continue\n else:\n valid_scores = cls_scores[valid_score_mask]\n valid_boxes = boxes[valid_score_mask]\n keep = nms(valid_boxes, valid_scores, nms_thr)\n if len(keep) > 0:\n cls_inds = np.ones((len(keep), 1)) * cls_ind\n dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)\n final_dets.append(dets)\n if len(final_dets) == 0:\n return None\n return np.concatenate(final_dets, 0)\n\n\ndef demo_postprocess(outputs, img_size, p6=False):\n\n grids = []\n expanded_strides = []\n\n if not p6:\n strides = [8, 16, 32]\n else:\n strides = [8, 16, 32, 64]\n\n hsizes = [img_size[0]//stride for stride in strides]\n wsizes = [img_size[1]//stride for stride in strides]\n\n for hsize, wsize, stride in zip(hsizes, wsizes, strides):\n xv, yv = np.meshgrid(np.arange(hsize), np.arange(wsize))\n grid = np.stack((xv, yv), 2).reshape(1, -1, 2)\n grids.append(grid)\n shape = grid.shape[:2]\n expanded_strides.append(np.full((*shape, 1), stride))\n\n grids = np.concatenate(grids, 1)\n expanded_strides = np.concatenate(expanded_strides, 1)\n outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides\n outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides\n\n return outputs\n"
},
{
"path": "detector/YOLOX/yolox/utils/dist.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# This file mainly comes from\n# https://github.com/facebookresearch/detectron2/blob/master/detectron2/utils/comm.py\n# Copyright (c) Facebook, Inc. and its affiliates.\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\"\"\"\nThis file contains primitives for multi-gpu communication.\nThis is useful when doing distributed training.\n\"\"\"\n\nimport functools\nimport logging\nimport pickle\nimport time\n\nimport numpy as np\n\nimport torch\nfrom torch import distributed as dist\n\n__all__ = [\n \"is_main_process\",\n \"synchronize\",\n \"get_world_size\",\n \"get_rank\",\n \"get_local_rank\",\n \"get_local_size\",\n \"time_synchronized\",\n \"gather\",\n \"all_gather\",\n]\n\n_LOCAL_PROCESS_GROUP = None\n\n\ndef synchronize():\n \"\"\"\n Helper function to synchronize (barrier) among all processes when using distributed training\n \"\"\"\n if not dist.is_available():\n return\n if not dist.is_initialized():\n return\n world_size = dist.get_world_size()\n if world_size == 1:\n return\n dist.barrier()\n\n\ndef get_world_size() -> int:\n if not dist.is_available():\n return 1\n if not dist.is_initialized():\n return 1\n return dist.get_world_size()\n\n\ndef get_rank() -> int:\n if not dist.is_available():\n return 0\n if not dist.is_initialized():\n return 0\n return dist.get_rank()\n\n\ndef get_local_rank() -> int:\n \"\"\"\n Returns:\n The rank of the current process within the local (per-machine) process group.\n \"\"\"\n if not dist.is_available():\n return 0\n if not dist.is_initialized():\n return 0\n assert _LOCAL_PROCESS_GROUP is not None\n return dist.get_rank(group=_LOCAL_PROCESS_GROUP)\n\n\ndef get_local_size() -> int:\n \"\"\"\n Returns:\n The size of the per-machine process group, i.e. the number of processes per machine.\n \"\"\"\n if not dist.is_available():\n return 1\n if not dist.is_initialized():\n return 1\n return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)\n\n\ndef is_main_process() -> bool:\n return get_rank() == 0\n\n\n@functools.lru_cache()\ndef _get_global_gloo_group():\n \"\"\"\n Return a process group based on gloo backend, containing all the ranks\n The result is cached.\n \"\"\"\n if dist.get_backend() == \"nccl\":\n return dist.new_group(backend=\"gloo\")\n else:\n return dist.group.WORLD\n\n\ndef _serialize_to_tensor(data, group):\n backend = dist.get_backend(group)\n assert backend in [\"gloo\", \"nccl\"]\n device = torch.device(\"cpu\" if backend == \"gloo\" else \"cuda\")\n\n buffer = pickle.dumps(data)\n if len(buffer) > 1024 ** 3:\n logger = logging.getLogger(__name__)\n logger.warning(\n \"Rank {} trying to all-gather {:.2f} GB of data on device {}\".format(\n get_rank(), len(buffer) / (1024 ** 3), device\n )\n )\n storage = torch.ByteStorage.from_buffer(buffer)\n tensor = torch.ByteTensor(storage).to(device=device)\n return tensor\n\n\ndef _pad_to_largest_tensor(tensor, group):\n \"\"\"\n Returns:\n list[int]: size of the tensor, on each rank\n Tensor: padded tensor that has the max size\n \"\"\"\n world_size = dist.get_world_size(group=group)\n assert (\n world_size >= 1\n ), \"comm.gather/all_gather must be called from ranks within the given group!\"\n local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device)\n size_list = [\n torch.zeros([1], dtype=torch.int64, device=tensor.device)\n for _ in range(world_size)\n ]\n dist.all_gather(size_list, local_size, group=group)\n size_list = [int(size.item()) for size in size_list]\n\n max_size = max(size_list)\n\n # we pad the tensor because torch all_gather does not support\n # gathering tensors of different shapes\n if local_size != max_size:\n padding = torch.zeros(\n (max_size - local_size,), dtype=torch.uint8, device=tensor.device\n )\n tensor = torch.cat((tensor, padding), dim=0)\n return size_list, tensor\n\n\ndef all_gather(data, group=None):\n \"\"\"\n Run all_gather on arbitrary picklable data (not necessarily tensors).\n\n Args:\n data: any picklable object\n group: a torch process group. By default, will use a group which\n contains all ranks on gloo backend.\n Returns:\n list[data]: list of data gathered from each rank\n \"\"\"\n if get_world_size() == 1:\n return [data]\n if group is None:\n group = _get_global_gloo_group()\n if dist.get_world_size(group) == 1:\n return [data]\n\n tensor = _serialize_to_tensor(data, group)\n\n size_list, tensor = _pad_to_largest_tensor(tensor, group)\n max_size = max(size_list)\n\n # receiving Tensor from all ranks\n tensor_list = [\n torch.empty((max_size,), dtype=torch.uint8, device=tensor.device)\n for _ in size_list\n ]\n dist.all_gather(tensor_list, tensor, group=group)\n\n data_list = []\n for size, tensor in zip(size_list, tensor_list):\n buffer = tensor.cpu().numpy().tobytes()[:size]\n data_list.append(pickle.loads(buffer))\n\n return data_list\n\n\ndef gather(data, dst=0, group=None):\n \"\"\"\n Run gather on arbitrary picklable data (not necessarily tensors).\n\n Args:\n data: any picklable object\n dst (int): destination rank\n group: a torch process group. By default, will use a group which\n contains all ranks on gloo backend.\n\n Returns:\n list[data]: on dst, a list of data gathered from each rank. Otherwise,\n an empty list.\n \"\"\"\n if get_world_size() == 1:\n return [data]\n if group is None:\n group = _get_global_gloo_group()\n if dist.get_world_size(group=group) == 1:\n return [data]\n rank = dist.get_rank(group=group)\n\n tensor = _serialize_to_tensor(data, group)\n size_list, tensor = _pad_to_largest_tensor(tensor, group)\n\n # receiving Tensor from all ranks\n if rank == dst:\n max_size = max(size_list)\n tensor_list = [\n torch.empty((max_size,), dtype=torch.uint8, device=tensor.device)\n for _ in size_list\n ]\n dist.gather(tensor, tensor_list, dst=dst, group=group)\n\n data_list = []\n for size, tensor in zip(size_list, tensor_list):\n buffer = tensor.cpu().numpy().tobytes()[:size]\n data_list.append(pickle.loads(buffer))\n return data_list\n else:\n dist.gather(tensor, [], dst=dst, group=group)\n return []\n\n\ndef shared_random_seed():\n \"\"\"\n Returns:\n int: a random number that is the same across all workers.\n If workers need a shared RNG, they can use this shared seed to\n create one.\n All workers must call this function, otherwise it will deadlock.\n \"\"\"\n ints = np.random.randint(2 ** 31)\n all_ints = all_gather(ints)\n return all_ints[0]\n\n\ndef time_synchronized():\n \"\"\"pytorch-accurate time\"\"\"\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n return time.time()\n"
},
{
"path": "detector/YOLOX/yolox/utils/ema.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\nimport math\nfrom copy import deepcopy\n\nimport torch\nimport torch.nn as nn\n\n\ndef is_parallel(model):\n \"\"\"check if model is in parallel mode.\"\"\"\n import apex\n\n parallel_type = (\n nn.parallel.DataParallel,\n nn.parallel.DistributedDataParallel,\n apex.parallel.distributed.DistributedDataParallel,\n )\n return isinstance(model, parallel_type)\n\n\ndef copy_attr(a, b, include=(), exclude=()):\n # Copy attributes from b to a, options to only include [...] and to exclude [...]\n for k, v in b.__dict__.items():\n if (len(include) and k not in include) or k.startswith(\"_\") or k in exclude:\n continue\n else:\n setattr(a, k, v)\n\n\nclass ModelEMA:\n \"\"\"\n Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models\n Keep a moving average of everything in the model state_dict (parameters and buffers).\n This is intended to allow functionality like\n https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage\n A smoothed version of the weights is necessary for some training schemes to perform well.\n This class is sensitive where it is initialized in the sequence of model init,\n GPU assignment and distributed training wrappers.\n \"\"\"\n def __init__(self, model, decay=0.9999, updates=0):\n \"\"\"\n Args:\n model (nn.Module): model to apply EMA.\n decay (float): ema decay reate.\n updates (int): counter of EMA updates.\n \"\"\"\n # Create EMA(FP32)\n self.ema = deepcopy(model.module if is_parallel(model) else model).eval()\n self.updates = updates\n # decay exponential ramp (to help early epochs)\n self.decay = lambda x: decay * (1 - math.exp(-x / 2000))\n for p in self.ema.parameters():\n p.requires_grad_(False)\n\n def update(self, model):\n # Update EMA parameters\n with torch.no_grad():\n self.updates += 1\n d = self.decay(self.updates)\n\n msd = (\n model.module.state_dict() if is_parallel(model) else model.state_dict()\n ) # model state_dict\n for k, v in self.ema.state_dict().items():\n if v.dtype.is_floating_point:\n v *= d\n v += (1.0 - d) * msd[k].detach()\n\n def update_attr(self, model, include=(), exclude=(\"process_group\", \"reducer\")):\n # Update EMA attributes\n copy_attr(self.ema, model, include, exclude)\n"
},
{
"path": "detector/YOLOX/yolox/utils/logger.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport inspect\nimport os\nimport sys\nfrom loguru import logger\n\n\ndef get_caller_name(depth=0):\n \"\"\"\n Args:\n depth (int): Depth of caller conext, use 0 for caller depth. Default value: 0.\n\n Returns:\n str: module name of the caller\n \"\"\"\n # the following logic is a little bit faster than inspect.stack() logic\n frame = inspect.currentframe().f_back\n for _ in range(depth):\n frame = frame.f_back\n\n return frame.f_globals[\"__name__\"]\n\n\nclass StreamToLoguru:\n \"\"\"\n stream object that redirects writes to a logger instance.\n \"\"\"\n def __init__(self, level=\"INFO\", caller_names=(\"apex\", \"pycocotools\")):\n \"\"\"\n Args:\n level(str): log level string of loguru. Default value: \"INFO\".\n caller_names(tuple): caller names of redirected module.\n Default value: (apex, pycocotools).\n \"\"\"\n self.level = level\n self.linebuf = \"\"\n self.caller_names = caller_names\n\n def write(self, buf):\n full_name = get_caller_name(depth=1)\n module_name = full_name.rsplit(\".\", maxsplit=-1)[0]\n if module_name in self.caller_names:\n for line in buf.rstrip().splitlines():\n # use caller level log\n logger.opt(depth=2).log(self.level, line.rstrip())\n else:\n sys.__stdout__.write(buf)\n\n def flush(self):\n pass\n\n\ndef redirect_sys_output(log_level=\"INFO\"):\n redirect_logger = StreamToLoguru(log_level)\n sys.stderr = redirect_logger\n sys.stdout = redirect_logger\n\n\ndef setup_logger(save_dir, distributed_rank=0, filename=\"log.txt\", mode=\"a\"):\n \"\"\"setup logger for training and testing.\n Args:\n save_dir(str): location to save log file\n distributed_rank(int): device rank when multi-gpu environment\n filename (string): log save name.\n mode(str): log file write mode, `append` or `override`. default is `a`.\n\n Return:\n logger instance.\n \"\"\"\n loguru_format = (\n \"{time:YYYY-MM-DD HH:mm:ss} | \"\n \"{level: <8} | \"\n \"{name} :{line} - {message} \"\n )\n\n logger.remove()\n save_file = os.path.join(save_dir, filename)\n if mode == \"o\" and os.path.exists(save_file):\n os.remove(save_file)\n # only keep logger in rank0 process\n if distributed_rank == 0:\n logger.add(\n sys.stderr,\n format=loguru_format,\n level=\"INFO\",\n enqueue=True,\n )\n logger.add(save_file)\n\n # redirect stdout/stderr to loguru\n redirect_sys_output(\"INFO\")\n"
},
{
"path": "detector/YOLOX/yolox/utils/lr_scheduler.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport math\nfrom functools import partial\n\n\nclass LRScheduler:\n def __init__(self, name, lr, iters_per_epoch, total_epochs, **kwargs):\n \"\"\"\n Supported lr schedulers: [cos, warmcos, multistep]\n\n Args:\n lr (float): learning rate.\n iters_per_peoch (int): number of iterations in one epoch.\n total_epochs (int): number of epochs in training.\n kwargs (dict):\n - cos: None\n - warmcos: [warmup_epochs, warmup_lr_start (default 1e-6)]\n - multistep: [milestones (epochs), gamma (default 0.1)]\n \"\"\"\n\n self.lr = lr\n self.iters_per_epoch = iters_per_epoch\n self.total_epochs = total_epochs\n self.total_iters = iters_per_epoch * total_epochs\n\n self.__dict__.update(kwargs)\n\n self.lr_func = self._get_lr_func(name)\n\n def update_lr(self, iters):\n return self.lr_func(iters)\n\n def _get_lr_func(self, name):\n if name == \"cos\": # cosine lr schedule\n lr_func = partial(cos_lr, self.lr, self.total_iters)\n elif name == \"warmcos\":\n warmup_total_iters = self.iters_per_epoch * self.warmup_epochs\n warmup_lr_start = getattr(self, \"warmup_lr_start\", 1e-6)\n lr_func = partial(\n warm_cos_lr,\n self.lr,\n self.total_iters,\n warmup_total_iters,\n warmup_lr_start,\n )\n elif name == \"yoloxwarmcos\":\n warmup_total_iters = self.iters_per_epoch * self.warmup_epochs\n no_aug_iters = self.iters_per_epoch * self.no_aug_epochs\n warmup_lr_start = getattr(self, \"warmup_lr_start\", 0)\n min_lr_ratio = getattr(self, \"min_lr_ratio\", 0.2)\n lr_func = partial(\n yolox_warm_cos_lr,\n self.lr,\n min_lr_ratio,\n self.total_iters,\n warmup_total_iters,\n warmup_lr_start,\n no_aug_iters,\n )\n elif name == \"yoloxsemiwarmcos\":\n warmup_lr_start = getattr(self, \"warmup_lr_start\", 0)\n min_lr_ratio = getattr(self, \"min_lr_ratio\", 0.2)\n warmup_total_iters = self.iters_per_epoch * self.warmup_epochs\n no_aug_iters = self.iters_per_epoch * self.no_aug_epochs\n normal_iters = self.iters_per_epoch * self.semi_epoch\n semi_iters = self.iters_per_epoch_semi * (\n self.total_epochs - self.semi_epoch - self.no_aug_epochs\n )\n lr_func = partial(\n yolox_semi_warm_cos_lr,\n self.lr,\n min_lr_ratio,\n warmup_lr_start,\n self.total_iters,\n normal_iters,\n no_aug_iters,\n warmup_total_iters,\n semi_iters,\n self.iters_per_epoch,\n self.iters_per_epoch_semi,\n )\n elif name == \"multistep\": # stepwise lr schedule\n milestones = [\n int(self.total_iters * milestone / self.total_epochs)\n for milestone in self.milestones\n ]\n gamma = getattr(self, \"gamma\", 0.1)\n lr_func = partial(multistep_lr, self.lr, milestones, gamma)\n else:\n raise ValueError(\"Scheduler version {} not supported.\".format(name))\n return lr_func\n\n\ndef cos_lr(lr, total_iters, iters):\n \"\"\"Cosine learning rate\"\"\"\n lr *= 0.5 * (1.0 + math.cos(math.pi * iters / total_iters))\n return lr\n\n\ndef warm_cos_lr(lr, total_iters, warmup_total_iters, warmup_lr_start, iters):\n \"\"\"Cosine learning rate with warm up.\"\"\"\n if iters <= warmup_total_iters:\n lr = (lr - warmup_lr_start) * iters / float(\n warmup_total_iters\n ) + warmup_lr_start\n else:\n lr *= 0.5 * (\n 1.0\n + math.cos(\n math.pi\n * (iters - warmup_total_iters)\n / (total_iters - warmup_total_iters)\n )\n )\n return lr\n\n\ndef yolox_warm_cos_lr(\n lr,\n min_lr_ratio,\n total_iters,\n warmup_total_iters,\n warmup_lr_start,\n no_aug_iter,\n iters,\n):\n \"\"\"Cosine learning rate with warm up.\"\"\"\n min_lr = lr * min_lr_ratio\n if iters <= warmup_total_iters:\n # lr = (lr - warmup_lr_start) * iters / float(warmup_total_iters) + warmup_lr_start\n lr = (lr - warmup_lr_start) * pow(\n iters / float(warmup_total_iters), 2\n ) + warmup_lr_start\n elif iters >= total_iters - no_aug_iter:\n lr = min_lr\n else:\n lr = min_lr + 0.5 * (lr - min_lr) * (\n 1.0\n + math.cos(\n math.pi\n * (iters - warmup_total_iters)\n / (total_iters - warmup_total_iters - no_aug_iter)\n )\n )\n return lr\n\n\ndef yolox_semi_warm_cos_lr(\n lr,\n min_lr_ratio,\n warmup_lr_start,\n total_iters,\n normal_iters,\n no_aug_iters,\n warmup_total_iters,\n semi_iters,\n iters_per_epoch,\n iters_per_epoch_semi,\n iters,\n):\n \"\"\"Cosine learning rate with warm up.\"\"\"\n min_lr = lr * min_lr_ratio\n if iters <= warmup_total_iters:\n # lr = (lr - warmup_lr_start) * iters / float(warmup_total_iters) + warmup_lr_start\n lr = (lr - warmup_lr_start) * pow(\n iters / float(warmup_total_iters), 2\n ) + warmup_lr_start\n elif iters >= normal_iters + semi_iters:\n lr = min_lr\n elif iters <= normal_iters:\n lr = min_lr + 0.5 * (lr - min_lr) * (\n 1.0\n + math.cos(\n math.pi\n * (iters - warmup_total_iters)\n / (total_iters - warmup_total_iters - no_aug_iters)\n )\n )\n else:\n lr = min_lr + 0.5 * (lr - min_lr) * (\n 1.0\n + math.cos(\n math.pi\n * (\n normal_iters\n - warmup_total_iters\n + (iters - normal_iters)\n * iters_per_epoch\n * 1.0\n / iters_per_epoch_semi\n )\n / (total_iters - warmup_total_iters - no_aug_iters)\n )\n )\n return lr\n\n\ndef multistep_lr(lr, milestones, gamma, iters):\n \"\"\"MultiStep learning rate\"\"\"\n for milestone in milestones:\n lr *= gamma if iters >= milestone else 1.0\n return lr\n"
},
{
"path": "detector/YOLOX/yolox/utils/metric.py",
"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\nimport functools\nimport os\nimport time\nfrom collections import defaultdict, deque\n\nimport numpy as np\n\nimport torch\n\n__all__ = [\n \"AverageMeter\",\n \"MeterBuffer\",\n \"get_total_and_free_memory_in_Mb\",\n \"occumpy_mem\",\n \"gpu_mem_usage\",\n]\n\n\ndef get_total_and_free_memory_in_Mb(cuda_device):\n devices_info_str = os.popen(\n \"nvidia-smi --query-gpu=memory.total,memory.used --format=csv,nounits,noheader\"\n )\n devices_info = devices_info_str.read().strip().split(\"\\n\")\n total, used = devices_info[int(cuda_device)].split(\",\")\n return int(total), int(used)\n\n\ndef occumpy_mem(cuda_device, mem_ratio=0.9):\n \"\"\"\n pre-allocate gpu memory for training to avoid memory Fragmentation.\n \"\"\"\n total, used = get_total_and_free_memory_in_Mb(cuda_device)\n max_mem = int(total * mem_ratio)\n block_mem = max_mem - used\n x = torch.cuda.FloatTensor(256, 1024, block_mem)\n del x\n time.sleep(5)\n\n\ndef gpu_mem_usage():\n \"\"\"\n Compute the GPU memory usage for the current device (MB).\n \"\"\"\n mem_usage_bytes = torch.cuda.max_memory_allocated()\n return mem_usage_bytes / (1024 * 1024)\n\n\nclass AverageMeter:\n \"\"\"Track a series of values and provide access to smoothed values over a\n window or the global series average.\n \"\"\"\n\n def __init__(self, window_size=50):\n self._deque = deque(maxlen=window_size)\n self._total = 0.0\n self._count = 0\n\n def update(self, value):\n self._deque.append(value)\n self._count += 1\n self._total += value\n\n @property\n def median(self):\n d = np.array(list(self._deque))\n return np.median(d)\n\n @property\n def avg(self):\n # if deque is empty, nan will be returned.\n d = np.array(list(self._deque))\n return d.mean()\n\n @property\n def global_avg(self):\n return self._total / max(self._count, 1e-5)\n\n @property\n def latest(self):\n return self._deque[-1] if len(self._deque) > 0 else None\n\n @property\n def total(self):\n return self._total\n\n def reset(self):\n self._deque.clear()\n self._total = 0.0\n self._count = 0\n\n def clear(self):\n self._deque.clear()\n\n\nclass MeterBuffer(defaultdict):\n \"\"\"Computes and stores the average and current value\"\"\"\n\n def __init__(self, window_size=20):\n factory = functools.partial(AverageMeter, window_size=window_size)\n super().__init__(factory)\n\n def reset(self):\n for v in self.values():\n v.reset()\n\n def get_filtered_meter(self, filter_key=\"time\"):\n return {k: v for k, v in self.items() if filter_key in k}\n\n def update(self, values=None, **kwargs):\n if values is None:\n values = {}\n values.update(kwargs)\n for k, v in values.items():\n self[k].update(v)\n\n def clear_meters(self):\n for v in self.values():\n v.clear()\n"
},
{
"path": "detector/YOLOX/yolox/utils/model_utils.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nfrom copy import deepcopy\n\nimport torch\nimport torch.nn as nn\nfrom thop import profile\n\n__all__ = [\n \"fuse_conv_and_bn\", \"fuse_model\", \"get_model_info\", \"replace_module\",\n]\n\n\ndef get_model_info(model, tsize):\n\n stride = 64\n img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device)\n flops, params = profile(deepcopy(model), inputs=(img,), verbose=False)\n params /= 1e6\n flops /= 1e9\n flops *= tsize[0] * tsize[1] / stride / stride * 2 # Gflops\n info = \"Params: {:.2f}M, Gflops: {:.2f}\".format(params, flops)\n return info\n\n\ndef fuse_conv_and_bn(conv, bn):\n # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/\n fusedconv = (\n nn.Conv2d(\n conv.in_channels,\n conv.out_channels,\n kernel_size=conv.kernel_size,\n stride=conv.stride,\n padding=conv.padding,\n groups=conv.groups,\n bias=True,\n )\n .requires_grad_(False)\n .to(conv.weight.device)\n )\n\n # prepare filters\n w_conv = conv.weight.clone().view(conv.out_channels, -1)\n w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))\n fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))\n\n # prepare spatial bias\n b_conv = (\n torch.zeros(conv.weight.size(0), device=conv.weight.device)\n if conv.bias is None\n else conv.bias\n )\n b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(\n torch.sqrt(bn.running_var + bn.eps)\n )\n fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)\n\n return fusedconv\n\n\ndef fuse_model(model):\n from yolox.models.network_blocks import BaseConv\n\n for m in model.modules():\n if type(m) is BaseConv and hasattr(m, \"bn\"):\n m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv\n delattr(m, \"bn\") # remove batchnorm\n m.forward = m.fuseforward # update forward\n return model\n\n\ndef replace_module(module, replaced_module_type, new_module_type, replace_func=None):\n \"\"\"\n Replace given type in module to a new type. mostly used in deploy.\n\n Args:\n module (nn.Module): model to apply replace operation.\n replaced_module_type (Type): module type to be replaced.\n new_module_type (Type)\n replace_func (function): python function to describe replace logic. Defalut value None.\n\n Returns:\n model (nn.Module): module that already been replaced.\n \"\"\"\n def default_replace_func(replaced_module_type, new_module_type):\n return new_module_type()\n\n if replace_func is None:\n replace_func = default_replace_func\n\n model = module\n if isinstance(module, replaced_module_type):\n model = replace_func(replaced_module_type, new_module_type)\n else: # recurrsively replace\n for name, child in module.named_children():\n new_child = replace_module(child, replaced_module_type, new_module_type)\n if new_child is not child: # child is already replaced\n model.add_module(name, new_child)\n\n return model\n"
},
{
"path": "detector/YOLOX/yolox/utils/setup_env.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport os\nimport subprocess\n\nimport cv2\n\n__all__ = [\"configure_nccl\", \"configure_module\"]\n\n\ndef configure_nccl():\n \"\"\"Configure multi-machine environment variables of NCCL.\"\"\"\n os.environ[\"NCCL_LAUNCH_MODE\"] = \"PARALLEL\"\n os.environ[\"NCCL_IB_HCA\"] = subprocess.getoutput(\n \"pushd /sys/class/infiniband/ > /dev/null; for i in mlx5_*; \"\n \"do cat $i/ports/1/gid_attrs/types/* 2>/dev/null \"\n \"| grep v >/dev/null && echo $i ; done; popd > /dev/null\"\n )\n os.environ[\"NCCL_IB_GID_INDEX\"] = \"3\"\n os.environ[\"NCCL_IB_TC\"] = \"106\"\n\n\ndef configure_module(ulimit_value=8192):\n \"\"\"\n Configure pytorch module environment. setting of ulimit and cv2 will be set.\n\n Args:\n ulimit_value(int): default open file number on linux. Default value: 8192.\n \"\"\"\n # system setting\n try:\n import resource\n rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)\n resource.setrlimit(resource.RLIMIT_NOFILE, (ulimit_value, rlimit[1]))\n except Exception:\n # Exception might be raised in Windows OS or rlimit reaches max limit number.\n # However, set rlimit value might not be necessary.\n pass\n\n # cv2\n # multiprocess might be harmful on performance of torch dataloader\n os.environ[\"OPENCV_OPENCL_RUNTIME\"] = \"disabled\"\n try:\n cv2.setNumThreads(0)\n cv2.ocl.setUseOpenCL(False)\n except Exception:\n # cv2 version mismatch might rasie exceptions.\n pass\n"
},
{
"path": "detector/YOLOX/yolox/utils/visualize.py",
"content": "#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.\n\nimport cv2\nimport numpy as np\n\n__all__ = [\"vis\"]\n\n\ndef vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None):\n\n for i in range(len(boxes)):\n box = boxes[i]\n cls_id = int(cls_ids[i])\n score = scores[i]\n if score < conf:\n continue\n x0 = int(box[0])\n y0 = int(box[1])\n x1 = int(box[2])\n y1 = int(box[3])\n\n color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()\n text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100)\n txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)\n font = cv2.FONT_HERSHEY_SIMPLEX\n\n txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]\n cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)\n\n txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()\n cv2.rectangle(\n img,\n (x0, y0 + 1),\n (x0 + txt_size[0] + 1, y0 + int(1.5*txt_size[1])),\n txt_bk_color,\n -1\n )\n cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)\n\n return img\n\n\n_COLORS = np.array(\n [\n 0.000, 0.447, 0.741,\n 0.850, 0.325, 0.098,\n 0.929, 0.694, 0.125,\n 0.494, 0.184, 0.556,\n 0.466, 0.674, 0.188,\n 0.301, 0.745, 0.933,\n 0.635, 0.078, 0.184,\n 0.300, 0.300, 0.300,\n 0.600, 0.600, 0.600,\n 1.000, 0.000, 0.000,\n 1.000, 0.500, 0.000,\n 0.749, 0.749, 0.000,\n 0.000, 1.000, 0.000,\n 0.000, 0.000, 1.000,\n 0.667, 0.000, 1.000,\n 0.333, 0.333, 0.000,\n 0.333, 0.667, 0.000,\n 0.333, 1.000, 0.000,\n 0.667, 0.333, 0.000,\n 0.667, 0.667, 0.000,\n 0.667, 1.000, 0.000,\n 1.000, 0.333, 0.000,\n 1.000, 0.667, 0.000,\n 1.000, 1.000, 0.000,\n 0.000, 0.333, 0.500,\n 0.000, 0.667, 0.500,\n 0.000, 1.000, 0.500,\n 0.333, 0.000, 0.500,\n 0.333, 0.333, 0.500,\n 0.333, 0.667, 0.500,\n 0.333, 1.000, 0.500,\n 0.667, 0.000, 0.500,\n 0.667, 0.333, 0.500,\n 0.667, 0.667, 0.500,\n 0.667, 1.000, 0.500,\n 1.000, 0.000, 0.500,\n 1.000, 0.333, 0.500,\n 1.000, 0.667, 0.500,\n 1.000, 1.000, 0.500,\n 0.000, 0.333, 1.000,\n 0.000, 0.667, 1.000,\n 0.000, 1.000, 1.000,\n 0.333, 0.000, 1.000,\n 0.333, 0.333, 1.000,\n 0.333, 0.667, 1.000,\n 0.333, 1.000, 1.000,\n 0.667, 0.000, 1.000,\n 0.667, 0.333, 1.000,\n 0.667, 0.667, 1.000,\n 0.667, 1.000, 1.000,\n 1.000, 0.000, 1.000,\n 1.000, 0.333, 1.000,\n 1.000, 0.667, 1.000,\n 0.333, 0.000, 0.000,\n 0.500, 0.000, 0.000,\n 0.667, 0.000, 0.000,\n 0.833, 0.000, 0.000,\n 1.000, 0.000, 0.000,\n 0.000, 0.167, 0.000,\n 0.000, 0.333, 0.000,\n 0.000, 0.500, 0.000,\n 0.000, 0.667, 0.000,\n 0.000, 0.833, 0.000,\n 0.000, 1.000, 0.000,\n 0.000, 0.000, 0.167,\n 0.000, 0.000, 0.333,\n 0.000, 0.000, 0.500,\n 0.000, 0.000, 0.667,\n 0.000, 0.000, 0.833,\n 0.000, 0.000, 1.000,\n 0.000, 0.000, 0.000,\n 0.143, 0.143, 0.143,\n 0.286, 0.286, 0.286,\n 0.429, 0.429, 0.429,\n 0.571, 0.571, 0.571,\n 0.714, 0.714, 0.714,\n 0.857, 0.857, 0.857,\n 0.000, 0.447, 0.741,\n 0.314, 0.717, 0.741,\n 0.50, 0.5, 0\n ]\n).astype(np.float32).reshape(-1, 3)\n"
},
{
"path": "docs/DATA.md",
"content": "# Dataset preparation\n\n### Introduction\n\nIn this documentation we introduce how to prepara standard datasets to benchmark UniTrack on different tasks. We consider five tasks: Single Object Tracking (SOT) on OTB 2015 dataset, Video Object Segmentation (VOS) on DAVIS 2017 dataset, Multiple Object Tracking (MOT) on MOT 16 dataset, Multiple Object Tracking and Segmentation (MOTS) on MOTS dataset, Pose Tracking on PoseTrack 2018 dataset. Among them, SOT and VOS are propagation-type tasks, in which only one observation (usually in the very first frame) is given to indicate the object to be tracked, while others are association-type tasks that support to make use of observations in every timestamp given by an automatic detector . \n\n- **Table of contents**\n - [Prepare OTB 2015 datset for SOT](#OTB-2015-dataset-for-SOT)\n - [Prepare DAVIS 2017 datset for VOS](#DAVIS-2017-dataset-for-VOS)\n - [Prepare MOT 16 datset for MOT](#MOT-16-dataset-for-MOT)\n - [Prepare MOTS dataset for MOTS](#MOTS-dataset-for-MOTS)\n - [Prepare PoseTrack 2018 for Pose Tracking](#PoseTrack-2018-dataset-for-Pose-Tracking)\n\n\n### OTB 2015 dataset for SOT\n\nThe [original source](http://cvlab.hanyang.ac.kr/tracker_benchmark/datasets.html) of OTB benchmark does not provide a convenient way to download the entire dataset. Lukily [Gluon CV](https://cv.gluon.ai/contents.html) provides a [script](https://cv.gluon.ai/_downloads/719c5c0d73fb22deacc84b4557b6fd5f/otb2015.py) for easy downloading all OTB video sequences. This script includes both dataset downloading and data processing，simply run this script:\n\n`python otb2015.py`\n\nand you will get all the 100 sequences of OTB. After this, you need to copy Jogging to Jogging-1 and Jogging-2, and copy Skating2 to Skating2-1 and Skating2-2 or using softlink, following [STVIR](https://github.com/STVIR/pysot/tree/master/testing_dataset). Finally, please download OTB2015.josn \\[[Google Drive](https://drive.google.com/file/d/1jHYta8wsSid9DwcWl5hcNJNPzgQMcI_r/view?usp=sharing)\\]\\[[Baidu NetDisk](https://pan.baidu.com/s/1d9oR7ZEHq4V5i6bLpEllng)\\] (code:k93s) and place it under the OTB-2015 root. The structrue should look like this:\n\n```\n${OTB_ROOT}\n |——— OTB2015.json\n | \n └———Basketball/\n | \n └———Biker/\n | \n ...\n```\n\n\n\n### DAVIS 2017 dataset for VOS\n\nDownload DAVIS 2017 trainval via [this link](https://data.vision.ee.ethz.ch/csergi/share/davis/DAVIS-2017-trainval-480p.zip) and unzip it. No other processing is needed.\n\n### MOT 16 dataset for MOT\n\n1. Download MOT-16 dataset from [this page](https://motchallenge.net/data/MOT16/).\n2. Get detections for MOT-16 sequences. Here we offer three options:\n\n - Using ground-truth detections. This is feasible only in the *train* split (we do not have labels for the *test* split). Run `python tools/gen_mot16_gt.py` to prepare the detections.\n - Using three kind of official detections (DPM/FRCNN/SDP) provided by MOT Challenge. Detections are from MOT-17 dataset (the video sequences in MOT-17 are the same as MOT-16), so you may need download MOT-17 and unzip it under the same root of MOT-16 first. Then run `python tools/gen_mot16_label17.py` to prepare the detections. Can generate detections for both *train* and *test* splits. \n - [Recommended] Using custom detectors to generate detection results. You need to first run the detector on MOT-16 dataset and output a series of `MOT16-XX.txt` files to store the detection results, where XX ranges from 01 to 14. Each line in the `.txt` file represents a bounding box in format of `[frame_index](starts from 1), x, y, w, h, confidence`. We provide an example generated by FairMOT detector \\[[Google Drive](https://drive.google.com/file/d/113xks7UIZ6LeBY_CTlOh5Z_OQ0hiP551/view?usp=sharing)\\]/\\[[Baidu NetDisk](https://pan.baidu.com/s/1-E9SN4rWWpZRT1ermcX0JA)\\] (code:k93s). Finally run `tools/gen_mot16_fairmot.py` to prepare the detections. \n \n A good point is that you can also download `.txt` results of other trackers from the [MOT-16 leaderboard](https://motchallenge.net/results/MOT16/) or of other detectors from the [MOT-17 DET leaderboard](https://motchallenge.net/results/MOT17Det/), and use their detection results with very few modifications on `tools/gen_mot16_fairmot.py`. \n\n### MOTS dataset for MOTS\n\n1. Download MOT-16 dataset from [this page](https://motchallenge.net/data/MOTS/).\n2. Get segmentation masks for MOTS sequences. Here we offer two options:\n - Using ground-truth detections. This is feasible only in the *train* split (we do not have labels for the *test* split). Run `python tools/gen_mots_gt.py` to prepare the detections.\n - [Recommended] Using custom models to generate segmentation masks. You need to first run the model on MOTS dataset and output a series of `MOTS-XX.txt` files to store the mask results. See [here](https://motchallenge.net/instructions/) for the output format. Note that we do not use the track \"id\" field so you can output any number as a place holder. You can download results of off-the-shelf trackers and use their masks, for example, simply download the raw data of results of the COSTA tracker in the bottom of [this page](https://motchallenge.net/method/MOTS=87&chl=17). Finally run `tools/gen_mot16_fairmot.py` to prepare the masks (The script will keep the track \"id\" field. But again, it should be noted that the track \"id\" field is ignored when run tracking with UniTrack). \n \n \n### PoseTrack 2018 dataset for Pose Tracking\n\n1. Register and download [PoseTrack 2018 dataset](https://posetrack.net/).\n2. Get single-frame pose estimation results. Run a single-frame pose estimator, and save results in a `$OBS_NAME.json` file.Results should be formatted as instructed [here](https://github.com/leonid-pishchulin/poseval). The \"track_id\" field is ignored so you can output any number as a place holder. \n3. Put the `.json` file under `$POSETRACK_ROOT/obs/$SPLIT/` folder, where `SPLIT` could be \"train\" or \"val\"."
},
{
"path": "docs/INSTALL.md",
"content": "# Installation\n\n### Requirements\n* Nvidia device with CUDA \n* Python 3.7+\n* PyTorch 1.7.0+\n* torchvision 0.8.0+\n* Other python packages in requirements.txt\n\n### Code installation\n\n#### (Recommended) Install with conda\n\nInstall conda from [here](https://repo.anaconda.com/miniconda/), Miniconda3-latest-(OS)-(platform).\n```shell\n# 1. Create a conda virtual environment.\nconda create -n unitrack python=3.7 -y\nconda activate unitrack\n\n# 2. Install PyTorch\nconda install pytorch==1.7.0 torchvision cudatoolkit\n\n# 3. Get UniTrack\ngit clone https://github.com/Zhongdao/UniTrack.git\ncd UniTrack\n\n# 4. Install ohter dependency\nconda install --file requirements.txt \npip install cython_bbox==0.1.3\npython setup.py\n\n```\n\n"
},
{
"path": "docs/MODELZOO.md",
"content": "# MODEL ZOO\n\n\n### Prepare apperance models\nOne beneficial usage of UniTrack is that it allows easy evaluation of pre-trained models (as appearance models) on diverse tracking tasks. By far we have tested the following models, mostly self-supervised pre-trained:\n\n| Pre-training Method | Architecture |Link | \n| :---: | :---: | :---: |\n| ImageNet classification | ResNet-50 | torchvision |\n| InsDist| ResNet-50 | [Google Drive](https://www.dropbox.com/sh/87d24jqsl6ra7t2/AACcsSIt1_Njv7GsmsuzZ6Sta/InsDis.pth)|\n| MoCo-V1| ResNet-50 |[Google Drive](https://dl.fbaipublicfiles.com/moco/moco_checkpoints/moco_v1_200ep/moco_v1_200ep_pretrain.pth.tar)|\n| PCL-V1| ResNet-50 |[Google Drive](https://storage.googleapis.com/sfr-pcl-data-research/PCL_checkpoint/PCL_v1_epoch200.pth.tar)|\n| PIRL| ResNet-50 | [Google Drive](https://www.dropbox.com/sh/87d24jqsl6ra7t2/AADN4jKnvTI0U5oT6hTmQZz8a/PIRL.pth)|\n| PCL-V2| ResNet-50 | [Google Drive](https://storage.googleapis.com/sfr-pcl-data-research/PCL_checkpoint/PCL_v2_epoch200.pth.tar)|\n| SimCLR-V1| ResNet-50 |[Google Drive](https://drive.google.com/file/d/1RdB2KaaXOtU2_t-Uk_HQbxMZgSGUcy6c/view?usp=sharing)|\n| MoCo-V2| ResNet-50 |[Google Drive](https://dl.fbaipublicfiles.com/moco/moco_checkpoints/moco_v2_800ep/moco_v2_800ep_pretrain.pth.tar)|\n| SimCLR-V2| ResNet-50 |[Google Drive](https://drive.google.com/file/d/1NSCrZ7MaejJaOS7yA3URtbubxLR-fz5X/view?usp=sharing)|\n| SeLa-V2| ResNet-50 |[Google Drive](https://dl.fbaipublicfiles.com/deepcluster/selav2_400ep_pretrain.pth.tar)|\n| InfoMin| ResNet-50 | [Google Drive](https://www.dropbox.com/sh/87d24jqsl6ra7t2/AAAzMTynP3Qc8mIE4XWkgILUa/InfoMin_800.pth)|\n| BarlowTwins| ResNet-50 | [Google Drive](https://drive.google.com/file/d/1iXfAiAZP3Lrc-Hk4QHUzO-mk4M4fElQw/view?usp=sharing)|\n| BYOL| ResNet-50 | [Google Drive](https://storage.googleapis.com/deepmind-byol/checkpoints/pretrain_res50x1.pkl)|\n| DeepCluster-V2| ResNet-50 |[Google Drive](https://dl.fbaipublicfiles.com/deepcluster/deepclusterv2_800ep_pretrain.pth.tar)|\n| SwAV| ResNet-50 |[Google Drive](https://dl.fbaipublicfiles.com/deepcluster/swav_800ep_pretrain.pth.tar)|\n| PixPro| ResNet-50 |[Google Drive](https://drive.google.com/file/d/1u172sUx-kldPvrZzZxijciBHLMiSJp46/view?usp=sharing)|\n| DetCo| ResNet-50 | [Google Drive](https://drive.google.com/file/d/1ahyX8HEbLUZXS-9Jr2GIMWDEZdqWe1GV/view?usp=sharing)|\n| TimeCycle| ResNet-50 |[Google Drive](https://drive.google.com/file/d/1WUYLkfowJ853RG_9OhbrKpb3r-cc-cOA/view?usp=sharing)|\n| ImageNet classification | ResNet-18 |torchvision|\n| Colorization + memory| ResNet-18 | [Google Drive](https://drive.google.com/file/d/1gWPRgYH70t-9uwj0EId826ZxFdosbzQv/view?usp=sharing)|\n| UVC| ResNet-18 |[Google Drive](https://drive.google.com/file/d/1nl0ehS8mvE5PUBOPLQSCWtrmFmS0-dPX/view?usp=sharing)|\n| CRW| ResNet-18 |[Google Drive](https://drive.google.com/file/d/1C1ujnpFRijJqVD3PV7qzyYwGSWoS9fLb/view?usp=sharing)|\n\nAfter downloading an appearance model, please place it under `$UNITRACK_ROOT/weights`. A large part of the model checkpoints are adopted from [ssl-transfer](https://github.com/linusericsson/ssl-transfer), many thanks to [linusericsson](https://github.com/linusericsson)!\n\n### Test your own pre-trained models as appearance models\nIf your model uses the standard ResNet architecture, you can directly test it using UniTrack without additional modifications. If you use ResNet but the parameter names are not consistent with the standard naming, you can simply rename parameter groups and load your weights into the standard ResNet. If you are using other architectures, it is also possible to test it with UniTrack. You may need a little hack: just remember to let the model output 8x down-sampled feature maps. You can check out `models/hrnet.py` for an example. \n"
},
{
"path": "docs/RESULTS.md",
"content": "### Quantitative results\n\n**Single Object Tracking (SOT) on OTB-2015**\n\n| Method | SiamFC | SiamRPN | SiamRPN++ | UDT* | UDT+* | LUDT* | LUDT+* | UniTrack_XCorr* | UniTrack_DCF* |\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| AUC | 58.2 | 63.7 | 69.6 | 59.4 | 63.2 | 60.2 | 63.9 | 55.5 | 61.8|\n\n \\* indicates non-supervised methods\n\n**Video Object Segmentation (VOS) on DAVIS-2017 *val* split**\n\n| Method | SiamMask | FeelVOS | STM | Colorization* | TimeCycle* | UVC* | CRW* | VFS* | UniTrack* |\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| J-mean | 54.3 | 63.7 | 79.2 | 34.6 | 40.1 | 56.7 | 64.8 | 66.5 | 58.4|\n\n \\* indicates non-supervised methods \n\n**Multiple Object Tracking (MOT) on MOT-16 [*test* set *private detector* track](https://motchallenge.net/method/MOT=3856&chl=5)**\n\n| Method | POI | DeepSORT-2 | JDE | CTrack | TubeTK | TraDes | CSTrack | FairMOT* | UniTrack* |\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| IDF-1 | 65.1 | 62.2 | 55.8 | 57.2 | 62.2 | 64.7 | 71.8 | 72.8 | 71.8|\n| IDs | 805 | 781 | 1544 | 1897 | 1236 | 1144 | 1071 | 1074 | 683 |\n| MOTA | 66.1 | 61.4 | 64.4 | 67.6 | 66.9 | 70.1 | 70.7 | 74.9 | 74.7|\n\n \\* indicates methods using the same detections\n\n**Multiple Object Tracking and Segmentation (MOTS) on MOTS challenge [*test* set](https://motchallenge.net/method/MOTS=109&chl=17)**\n\n| Method | TrackRCNN | SORTS | PointTrack | GMPHD | COSTA_st* | UniTrack* |\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | \n| IDF-1 | 42.7 | 57.3 | 42.9 | 65.6 | 70.3 | 67.2 |\n| IDs | 567 | 577 | 868 | 566 | 421 | 622 | \n| sMOTA | 40.6 | 55.0 | 62.3 | 69.0 | 70.2 | 68.9 | \n\n \\* indicates methods using the same detections\n\n**Pose Tracking on PoseTrack-2018 *val* split**\n\n| Method | MDPN | OpenSVAI | Miracle | KeyTrack | LightTrack* | UniTrack* |\n| :---: | :---: | :---: | :---: | :---: | :---: | :---: | \n| IDF-1 | - | - | - | - | 52.2 | 73.2 |\n| IDs | - | - | - | - | 3024 | 6760 | \n| sMOTA | 50.6 | 62.4 | 64.0 | 66.6 | 64.8 | 63.5 | \n\n \\* indicates methods using the same detections\n"
},
{
"path": "docs/RUN.md",
"content": "# Run evaluation on multiple tasks\n\n### Prepare config file\n\nTo evaluate an apperance model on multiple tasks, first you need to prepare a config file `${EXP_NAME}.yaml` and place it under the `config/` folder. We provide several example config files: \n 1. `crw_resnet18_s3.yaml` : Self-supervised model trained with Contrastive Random Walk [1], ResNet-18 stage-3 features.\n 2. `imagenet_resnet18_s3.yaml`: ImageNet pre-trained model, ResNet-18 stage-3 features.\n 3. `crw_resnet18_s3_womotion.yaml` : Model same as 1 but motion cues are discarded in association type tasks. This way, distinctions between different representations are better highlighted and potential confounding factors are avoided.\n 4. `imagenet_resnet18_s3_womotion.yaml`: Model same as 2, motion cues are discared in association type tasks.\n \n\n### Note for the config file\n\nWhen you are testing a new model, please take care to make sure the following fields in the config file are correct:\n\n```yaml\ncommon:\n # Experiment name, an identifier.\n exp_name: crw_resnet18_s3 \n \n # Model type, currently support:\n # ['imagenet18', 'imagenet50', 'imagenet101', 'random18', 'random50',\n # 'imagenet_resnext50', 'imagenet_resnext101'\n # 'byol', 'deepcluster-v2', 'infomin', 'insdis', 'moco-v1', 'moco-v2',\n # 'pcl-v1', 'pcl-v2','pirl', 'sela-v2', 'swav', 'simclr-v1', 'simclr-v2',\n # 'pixpro', 'detco', 'barlowtwins', 'crw', 'uvc', 'timecycle']\n model_type: crw \n\n # For ResNet architecture, remove layer4 means output layer3 features\n remove_layers: ['layer4'] \n \n # Be careful about this\n im_mean: [0.4914, 0.4822, 0.4465] \n im_std: [0.2023, 0.1994, 0.2010]\n \n # Path to the model weights.\n resume: 'weights/crw.pth'\n \nmot:\n # The single-frame observations. should correspond to a folder ${mot_root}/obs/${obid}\n obid: 'FairMOT'\n # Dataset root\n mot_root: '/home/wangzd/datasets/MOT/MOT16'\n # There is no validation set, so by default we test on the train split. \n \nmots:\n # The single-frame observations. should correspond to a folder ${mots_root}/obs/${obid}\n obid: 'COSTA'\n # Dataset root\n mots_root: '/home/wangzd/datasets/GOT/MOTS'\n # There is no validation set, so by default we test on the train split. \n \nposetrack:\n # The single-frame observations. should correspond to a folder ${mots_root}/obs/val/${obid}\n obid: 'lighttrack_MSRA152\n # Dataset root\n data_root: '/home/wangzd/datasets/GOT/Posetrack2018'\n # There is a validation set, by default we test on the val split.\n split: 'val'\n \n```\n\nFor other arguments, just refer to `crw_resnet18_s3.yaml` or `crw_resnet18_s3_womotion.yaml`.\n\n\n### Run\n\nSuppose the current path is `$UNITRACK_ROOT`, you can run multiple tasks with a single command:\n\n```shell\n./eval.sh $EXP_NAME $GPU_ID\n```\n\nYou will obtain a set of summaries of quantitative results under `results/summary`, and also visualizations of all results under `results`\n\n\n[1]. Jabri, Allan, Andrew Owens, and Alexei A. Efros. \"Space-time correspondence as a contrastive random walk.\" In NeurIPS, 2020."
},
{
"path": "eval/convert_davis.py",
"content": "import os\nimport numpy as np\nimport cv2\nimport os.path as osp\nimport pdb\n\nfrom PIL import Image\n\njpglist = []\n\nimport palette\nimport argparse\n\nparser = argparse.ArgumentParser()\nparser.add_argument('-o', '--out_folder', default='/scratch/ajabri/davis_results/', type=str)\nparser.add_argument('-i', '--in_folder', default='/scratch/ajabri/davis_results_masks/', type=str)\nparser.add_argument('-d', '--dataset', default='/scratch/ajabri/data/davis/', type=str)\n\nargs = parser.parse_args()\n\nannotations_folder = args.dataset + '/Annotations/480p/'\nf1 = open(args.dataset + '/ImageSets/2017/val.txt', 'r')\nfor line in f1:\n line = line[:-1]\n jpglist.append(line)\nf1.close()\n\nout_folder = args.out_folder\ncurrent_folder = args.in_folder\n\nif not os.path.exists(out_folder):\n os.makedirs(out_folder)\n\npalette = palette.tensor.astype(np.uint8)\ndef color2id(c):\n return np.arange(0, palette.shape[0])[np.all(palette == c, axis=-1)]\n\ndef convert_dir(i):\n fname = jpglist[i]\n gtfolder = osp.join(annotations_folder,fname)\n outfolder = osp.join(out_folder,fname)\n\n if not os.path.exists(outfolder):\n os.mkdir(outfolder)\n\n files = [_ for _ in os.listdir(gtfolder) if _[-4:] == '.png']\n lblimg = cv2.imread(osp.join(gtfolder,\"{:05d}.png\".format(0)))\n height = lblimg.shape[0]\n width = lblimg.shape[1]\n\n for j in range(len(files)):\n outname = osp.join(outfolder, \"{:05d}.png\".format(j))\n inname = osp.join(current_folder, str(i) + '_' + str(j) + '_mask.png')\n\n lblimg = cv2.imread(inname)[:,:,::-1]\n flat_lblimg = lblimg.reshape(-1, 3)\n lblidx = np.zeros((lblimg.shape[0], lblimg.shape[1]))\n lblidx2 = np.zeros((lblimg.shape[0], lblimg.shape[1]))\n\n colors = np.unique(flat_lblimg, axis=0)\n\n for c in colors:\n cid = color2id(c)\n if len(cid) > 0:\n lblidx2[np.all(lblimg == c, axis=-1)] = cid\n\n lblidx = lblidx2\n\n lblidx = lblidx.astype(np.uint8)\n lblidx = cv2.resize(lblidx, (width, height), interpolation=cv2.INTER_NEAREST)\n lblidx = lblidx.astype(np.uint8)\n\n im = Image.fromarray(lblidx)\n im.putpalette(palette.ravel())\n im.save(outname, format='PNG')\n\nimport multiprocessing as mp\npool = mp.Pool(10)\nresults = pool.map(convert_dir, range(len(jpglist)))\n"
},
{
"path": "eval/davis_dummy.txt",
"content": "/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/dance-twirl /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/dance-twirl\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/dog /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/dog\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/dogs-jump /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/dogs-jump\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/drift-chicane /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/drift-chicane\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/drift-straight /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/drift-straight\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/goat /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/goat\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/gold-fish /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/gold-fish\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/horsejump-high /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/horsejump-high\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/india /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/india\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/judo /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/judo\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/kite-surf /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/kite-surf\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/lab-coat /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/lab-coat\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/libby /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/libby\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/loading /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/loading\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/mbike-trick /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/mbike-trick\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/motocross-jump /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/motocross-jump\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/paragliding-launch /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/paragliding-launch\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/parkour /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/parkour\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/pigs /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/pigs\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/scooter-black /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/scooter-black\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/shooting /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/shooting\n/home/wangzd/datasets/uvc/DAVIS/JPEGImages/480p/soapbox /home/wangzd/datasets/uvc/DAVIS/Annotations/480p/soapbox\n"
},
{
"path": "eval/eval_mot.py",
"content": "import os\r\nimport numpy as np\r\nimport copy\r\nimport motmetrics as mm\r\nmm.lap.default_solver = 'lap'\r\nfrom utils.io import read_mot_results, unzip_objs\r\n\r\n\r\nclass Evaluator(object):\r\n\r\n def __init__(self, data_root, seq_name, data_type='mot'):\r\n self.data_root = data_root\r\n self.seq_name = seq_name\r\n self.data_type = data_type\r\n\r\n self.load_annotations()\r\n self.reset_accumulator()\r\n\r\n def load_annotations(self):\r\n assert self.data_type == 'mot'\r\n\r\n gt_filename = os.path.join(self.data_root, self.seq_name, 'gt', 'gt.txt')\r\n self.gt_frame_dict = read_mot_results(gt_filename, self.data_type, is_gt=True)\r\n self.gt_ignore_frame_dict = read_mot_results(gt_filename, self.data_type, is_ignore=True)\r\n\r\n def reset_accumulator(self):\r\n self.acc = mm.MOTAccumulator(auto_id=True)\r\n\r\n def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False):\r\n # results\r\n trk_tlwhs = np.copy(trk_tlwhs)\r\n trk_ids = np.copy(trk_ids)\r\n\r\n # gts\r\n gt_objs = self.gt_frame_dict.get(frame_id, [])\r\n gt_tlwhs, gt_ids = unzip_objs(gt_objs)[:2]\r\n\r\n # ignore boxes\r\n ignore_objs = self.gt_ignore_frame_dict.get(frame_id, [])\r\n ignore_tlwhs = unzip_objs(ignore_objs)[0]\r\n\r\n\r\n # remove ignored results\r\n keep = np.ones(len(trk_tlwhs), dtype=bool)\r\n iou_distance = mm.distances.iou_matrix(ignore_tlwhs, trk_tlwhs, max_iou=0.5)\r\n if len(iou_distance) > 0:\r\n match_is, match_js = mm.lap.linear_sum_assignment(iou_distance)\r\n match_is, match_js = map(lambda a: np.asarray(a, dtype=int), [match_is, match_js])\r\n match_ious = iou_distance[match_is, match_js]\r\n\r\n match_js = np.asarray(match_js, dtype=int)\r\n match_js = match_js[np.logical_not(np.isnan(match_ious))]\r\n keep[match_js] = False\r\n trk_tlwhs = trk_tlwhs[keep]\r\n trk_ids = trk_ids[keep]\r\n\r\n # get distance matrix\r\n iou_distance = mm.distances.iou_matrix(gt_tlwhs, trk_tlwhs, max_iou=0.5)\r\n\r\n # acc\r\n self.acc.update(gt_ids, trk_ids, iou_distance)\r\n\r\n if rtn_events and iou_distance.size > 0 and hasattr(self.acc, 'last_mot_events'):\r\n events = self.acc.last_mot_events # only supported by https://github.com/longcw/py-motmetrics\r\n else:\r\n events = None\r\n return events\r\n\r\n def eval_file(self, filename):\r\n self.reset_accumulator()\r\n\r\n result_frame_dict = read_mot_results(filename, self.data_type, is_gt=False)\r\n frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys())))\r\n for frame_id in frames:\r\n trk_objs = result_frame_dict.get(frame_id, [])\r\n trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2]\r\n self.eval_frame(frame_id, trk_tlwhs, trk_ids, rtn_events=False)\r\n\r\n return self.acc\r\n\r\n @staticmethod\r\n def get_summary(accs, names, metrics=('mota', 'num_switches', 'idp', 'idr', 'idf1', 'precision', 'recall')):\r\n names = copy.deepcopy(names)\r\n if metrics is None:\r\n metrics = mm.metrics.motchallenge_metrics\r\n metrics = copy.deepcopy(metrics)\r\n\r\n mh = mm.metrics.create()\r\n summary = mh.compute_many(\r\n accs,\r\n metrics=metrics,\r\n names=names,\r\n generate_overall=True\r\n )\r\n\r\n return summary\r\n\r\n @staticmethod\r\n def save_summary(summary, filename):\r\n import pandas as pd\r\n writer = pd.ExcelWriter(filename)\r\n summary.to_excel(writer)\r\n writer.save()\r\n"
},
{
"path": "eval/eval_pck.py",
"content": "from scipy.io import loadmat\nfrom numpy import transpose\nimport skimage.io as sio\nimport numpy as np\nimport os\nimport cv2\n\nimport scipy.io as sio\n\nfilelist = '/home/wangzd/datasets/GOT/JHMDB/split.txt'\nsrc_folder = 'results/poseprop/womotion_resnet18_s3/'\n\nf = open(filelist, 'r')\ngts = []\nheights = []\nwidths = []\npreds = []\njnt_visible_set = []\nhuman_boxes = []\n\nfeat_res = 40\n\nfor cnt, line in enumerate(f):\n rows = line.strip().split()\n lblpath = rows[0] #+ '/joint_positions.mat'\n lbls_mat = sio.loadmat(lblpath)\n lbls_coord = lbls_mat['pos_img']\n lbls_coord = lbls_coord - 1\n\n gts.append(lbls_coord)\n\n imgpath = rows[1] + '/00001.png'\n img = cv2.imread(imgpath)\n heights.append(img.shape[0])\n widths.append(img.shape[1])\n\nf.close()\n\n# gts = gts[0: 200]\nprint('read gt')\n\n# read prediction results\nfor i in range(len(gts)):\n\n # import pdb; pdb.set_trace()\n predfile = src_folder + str(i) + '.dat'\n predres = np.load(predfile, allow_pickle=True)\n\n # import pdb; pdb.set_trace()\n\n jnt_visible = np.ones((predres.shape[1], predres.shape[2]))\n\n for j in range(predres.shape[1]):\n for k in range(predres.shape[2]):\n if predres[0, j, k] < 0:\n jnt_visible[j, k] = 0\n\n jnt_visible_set.append(jnt_visible)\n\n now_height = heights[i]\n now_width = widths[i]\n predres[0, :, :] = predres[0, :, :] / float(feat_res) * now_width\n predres[1, :, :] = predres[1, :, :] / float(feat_res) * now_height\n\n preds.append(predres)\n\nprint('read prediction')\n\n# compute the human box for normalization\nfor i in range(len(gts)):\n\n nowgt = gts[i]\n jnt_visible = jnt_visible_set[i]\n now_boxes = np.zeros(nowgt.shape[2])\n\n for k in range(nowgt.shape[2]):\n minx = 1e6\n maxx = -1\n miny = 1e6\n maxy = -1\n for j in range(nowgt.shape[1]):\n\n if jnt_visible[j, k] == 0:\n continue\n\n minx = np.min([minx, nowgt[0, j, k]])\n miny = np.min([miny, nowgt[1, j, k]])\n maxx = np.max([maxx, nowgt[0, j, k]])\n maxy = np.max([maxy, nowgt[1, j, k]])\n\n now_boxes[k] = 0.6 * np.linalg.norm(np.subtract([maxx,maxy],[minx,miny]))\n # now_boxes[k] = np.max([maxy - miny, maxx - minx])\n\n\n human_boxes.append(now_boxes)\n\nprint('done box')\n\n# compute distances\ndistAll = {}\nfor pidx in range(15):\n distAll[pidx] = np.zeros([0,0])\n\nfor i in range(len(gts)):\n\n predres = preds[i]\n nowgt = gts[i]\n now_boxes = human_boxes[i]\n jnt_visible = jnt_visible_set[i]\n for j in range(nowgt.shape[1]):\n for k in range(nowgt.shape[2]):\n\n if jnt_visible[j, k] == 0:\n continue\n\n if k == 0:\n continue\n\n predx = predres[0, j, k]\n predy = predres[1, j, k]\n gtx = nowgt[0, j, k]\n gty = nowgt[1, j, k]\n d = np.linalg.norm(np.subtract([predx, predy],[gtx, gty]))\n dNorm = d / now_boxes[k]\n\n distAll[j] = np.append(distAll[j],[[dNorm]])\n\nprint('done distances')\n\ndef computePCK(distAll,distThresh):\n\n pckAll = np.zeros([len(distAll)+1,1])\n nCorrect = 0\n nTotal = 0\n for pidx in range(len(distAll)):\n idxs = np.argwhere(distAll[pidx] <= distThresh)\n pck = 100.0*len(idxs)/len(distAll[pidx])\n pckAll[pidx,0] = pck\n nCorrect += len(idxs)\n nTotal += len(distAll[pidx])\n\n pckAll[len(distAll),0] = np.mean(pckAll[0 :len(distAll),0]) # 100.0*nCorrect/nTotal\n\n return pckAll\n\nrng = [0.1, 0.2, 0.3, 0.4, 0.5]\n\nfor i in range(len(rng)):\n pckall = computePCK(distAll, rng[i])\n print(str(rng[i]) + ': ' + str(pckall[-1]) )\n # print(pckall[-1])\n"
},
{
"path": "eval/mots/Evaluator.py",
"content": "import pdb\nimport sys, os\nsys.path.append(os.getcwd())\nimport argparse\n\nimport traceback\nimport time\nimport pickle\nimport pandas as pd\nimport glob\nfrom os import path\nimport numpy as np\n\n\n\nclass Evaluator(object):\n \"\"\" The `Evaluator` class runs evaluation per sequence and computes the overall performance on the benchmark\"\"\"\n def __init__(self):\n pass\n\n def run(self, benchmark_name = None , gt_dir = None, res_dir = None, save_pkl = None, eval_mode = \"train\", seqmaps_dir = \"seqmaps\"):\n \"\"\"\n Params\n -----\n benchmark_name: Name of benchmark, e.g. MOT17\n gt_dir: directory of folders with gt data, including the c-files with sequences\n res_dir: directory with result files\n .txt\n .txt\n ...\n .txt\n eval_mode:\n seqmaps_dir:\n seq_file: File name of file containing sequences, e.g. 'c10-train.txt'\n save_pkl: path to output directory for final results\n \"\"\"\n\n start_time = time.time()\n\n self.benchmark_gt_dir = gt_dir\n self.seq_file = \"{}-{}.txt\".format(benchmark_name, eval_mode)\n\n res_dir = res_dir\n self.benchmark_name = benchmark_name\n self.seqmaps_dir = seqmaps_dir\n\n self.mode = eval_mode\n\n self.datadir = os.path.join(gt_dir, self.mode)\n\n # getting names of sequences to evaluate\n error_traceback = \"\"\n assert self.mode in [\"train\", \"test\", \"all\"], \"mode: %s not valid \" %s\n\n print(\"Evaluating Benchmark: %s\" % self.benchmark_name)\n\n # ======================================================\n # Handle evaluation\n # ======================================================\n\n\n\n # load list of all sequences\n self.sequences = os.listdir(self.datadir) \n\n self.gtfiles = []\n self.tsfiles = []\n for seq in self.sequences:\n gtf = os.path.join(self.benchmark_gt_dir, self.mode ,seq, 'gt/gt.txt')\n if path.exists(gtf): self.gtfiles.append(gtf)\n else: raise Exception(\"Ground Truth %s missing\" % gtf)\n tsf = os.path.join( res_dir, \"%s.txt\" % seq)\n if path.exists(gtf): self.tsfiles.append(tsf)\n else: raise Exception(\"Result file %s missing\" % tsf)\n\n\n print('Found {} ground truth files and {} test files.'.format(len(self.gtfiles), len(self.tsfiles)))\n print( self.tsfiles)\n\n self.MULTIPROCESSING = False \n MAX_NR_CORES = 10\n # set number of core for mutliprocessing\n if self.MULTIPROCESSING: self.NR_CORES = np.minimum( MAX_NR_CORES, len(self.tsfiles))\n try:\n\n \"\"\" run evaluation \"\"\"\n results = self.eval()\n\n # calculate overall results\n results_attributes = self.Overall_Results.metrics.keys()\n\n for attr in results_attributes:\n \"\"\" accumulate evaluation values over all sequences \"\"\"\n try:\n self.Overall_Results.__dict__[attr] = sum(obj.__dict__[attr] for obj in self.results)\n except:\n pass\n cache_attributes = self.Overall_Results.cache_dict.keys()\n for attr in cache_attributes:\n \"\"\" accumulate cache values over all sequences \"\"\"\n try:\n self.Overall_Results.__dict__[attr] = self.Overall_Results.cache_dict[attr]['func']([obj.__dict__[attr] for obj in self.results])\n except:\n pass\n print(\"evaluation successful\")\n\n\n # Compute clearmot metrics for overall and all sequences\n for res in self.results:\n res.compute_clearmot()\n self.Overall_Results.compute_clearmot()\n\n\n self.accumulate_df(type = \"mail\")\n self.failed = False\n error = None\n\n\n except Exception as e:\n print(str(traceback.format_exc()))\n print (\"\")\n error_traceback = [m.split(\"\")[0] for m in startExc[1:]]\n\n error = \"\"\n\n for err in error_traceback:\n error+=\"Error: %s\" % err\n\n\n print( \"Error Message\", error)\n self.error = error\n print(\"ERROR %s\" % error)\n\n\n print (\"Evaluation Finished\")\n print(\"Your Results\")\n print(self.render_summary())\n # save results if path set\n if save_pkl:\n\n\n self.Overall_Results.save_dict(os.path.join( save_pkl, \"%s-%s-overall.pkl\" % (self.benchmark_name, self.mode)))\n for res in self.results:\n res.save_dict(os.path.join( save_pkl, \"%s-%s-%s.pkl\" % (self.benchmark_name, self.mode, res.seqName)))\n print(\"Successfully save results\")\n\n return self.Overall_Results, self.results\n def eval(self):\n raise NotImplementedError\n\n\n def accumulate_df(self, type = None):\n \"\"\" create accumulated dataframe with all sequences \"\"\"\n for k, res in enumerate(self.results):\n res.to_dataframe(display_name = True, type = type )\n if k == 0: summary = res.df\n else: summary = summary.append(res.df)\n summary = summary.sort_index()\n\n\n self.Overall_Results.to_dataframe(display_name = True, type = type )\n\n self.summary = summary.append(self.Overall_Results.df)\n\n\n\n def render_summary( self, buf = None):\n \"\"\"Render metrics summary to console friendly tabular output.\n\n Params\n ------\n summary : pd.DataFrame\n Dataframe containing summaries in rows.\n\n Kwargs\n ------\n buf : StringIO-like, optional\n Buffer to write to\n formatters : dict, optional\n Dicionary defining custom formatters for individual metrics.\n I.e `{'mota': '{:.2%}'.format}`. You can get preset formatters\n from MetricsHost.formatters\n namemap : dict, optional\n Dictionary defining new metric names for display. I.e\n `{'num_false_positives': 'FP'}`.\n\n Returns\n -------\n string\n Formatted string\n \"\"\"\n\n output = self.summary.to_string(\n buf=buf,\n formatters=self.Overall_Results.formatters,\n justify = \"left\"\n )\n\n return output\ndef run_metrics( metricObject, args ):\n \"\"\" Runs metric for individual sequences\n Params:\n -----\n metricObject: metricObject that has computer_compute_metrics_per_sequence function\n args: dictionary with args for evaluation function\n \"\"\"\n metricObject.compute_metrics_per_sequence(**args)\n return metricObject\n\n\nif __name__ == \"__main__\":\n Evaluator()\n"
},
{
"path": "eval/mots/LICENSE",
"content": "MIT License\n\nCopyright (c) 2019 Visual Computing Institute\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "eval/mots/MOTSVisualization.py",
"content": "import sys\nfrom Visualize import Visualizer\nfrom mots_common.io import load_sequences, load_seqmap, load_txt\nimport pycocotools.mask as rletools\nimport glob\nimport os\nimport cv2\nimport colorsys\nimport numpy as np\n\ndef apply_mask(image, mask, color, alpha=0.5):\n\t\"\"\"\n\t Apply the given mask to the image.\n\t\"\"\"\n\tfor c in range(3):\n\t\timage[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - alpha) + alpha * color[c],image[:, :, c])\n\treturn image\n\n\nclass MOTSVisualizer(Visualizer):\n\n\tdef load(self, FilePath):\n\t\treturn load_txt(FilePath)\n\n\n\n\tdef drawResults(self, im = None, t = 0):\n\t\tself.draw_boxes = False\n\n\n\t\tfor obj in self.resFile[t]:\n\n\t\t\tcolor = self.colors[obj.track_id % len(self.colors)]\n\n\t\t\tcolor = tuple([int(c*255) for c in color])\n\n\n\t\t\tif obj.class_id == 1:\n\t\t\t\tcategory_name = \"Car\"\n\t\t\telif obj.class_id == 2:\n\t\t\t\tcategory_name = \"Ped\"\n\t\t\telse:\n\t\t\t\tcategory_name = \"Ignore\"\n\t\t\t\tcolor = (0.7*255, 0.7*255, 0.7*255)\n\n\t\t\tif obj.class_id == 1 or obj.class_id == 2: # Don't show boxes or ids for ignore regions\n\t\t\t\tx, y, w, h = rletools.toBbox(obj.mask)\n\n\t\t\t\tpt1=(int(x),int(y))\n\t\t\t\tpt2=(int(x+w),int(y+h))\n\n\t\t\t\tcategory_name += \":\" + str(obj.track_id)\n\t\t\t\tcv2.putText(im, category_name, (int(x + 0.5 * w), int( y + 0.5 * h)), cv2.FONT_HERSHEY_TRIPLEX,self.imScale,color,thickness =2)\n\t\t\t\tif self.draw_boxes:\n\t\t\t\t\tcv2.rectangle(im,pt1,pt2,color,2)\n\n\n\t\t\tbinary_mask = rletools.decode(obj.mask)\n\n\t\t\tim = apply_mask(im, binary_mask, color)\n\t\treturn im\n\nif __name__ == \"__main__\":\n\tvisualizer = MOTSVisualizer(\n\tseqName = \"MOTS20-11\",\n\tFilePath =\"data/MOTS/train/MOTS20-11/gt/gt.txt\",\n\timage_dir = \"data/MOTS/train/MOTS20-11/img1\",\n\tmode = \"gt\",\n\toutput_dir = \"vid\")\n\n\tvisualizer.generateVideo(\n\t displayTime = True,\n\t displayName = \"seg\",\n\t showOccluder = True,\n\t fps = 25 )\n"
},
{
"path": "eval/mots/MOTS_metrics.py",
"content": "import math\nfrom collections import defaultdict\nimport pycocotools.mask as rletools\nfrom mots_common.io import SegmentedObject\nfrom mots_common.io import load_seqmap, load_sequences, load_txt\nfrom Metrics import Metrics\nimport os, sys\nimport numpy as np\nfrom scipy.optimize import linear_sum_assignment as linear_assignment\n\n# we only consider pedestrians\nIGNORE_CLASS = 10\nCLASS_ID = 2\n\n\n\ndef mask_iou(a, b, criterion=\"union\"):\n is_crowd = criterion != \"union\"\n return rletools.iou([a.mask], [b.mask], [is_crowd])[0][0]\n\n\n\n\nclass MOTSMetrics(Metrics):\n\tdef __init__(self, seqName = None):\n\t\tsuper().__init__()\n\t\tif seqName:\n\t\t\tself.seqName = seqName\n\t\telse: self.seqName = 0\n\t\t# Evaluation metrics\n\n\t\tself.register(name = \"sMOTSA\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"MOTSA\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"MOTSP\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"MOTSAL\", formatter='{:.2f}'.format, write_mail = False)\n\t\tself.register(name = \"MODSA\", formatter='{:.2f}'.format, write_mail = False)\n\t\tself.register(name = \"MODSP\", formatter='{:.2f}'.format, write_mail = False)\n\n\n\t\tself.register(name = \"IDF1\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"IDTP\", formatter='{:.2f}'.format, write_mail = False)\n\n\n\t\tself.register(name = \"MT\", formatter='{:.0f}'.format)\n\t\tself.register(name = \"PT\", formatter='{:.0f}'.format, write_mail = False )\n\t\tself.register(name = \"ML\", formatter='{:.0f}'.format)\n\n\t\tself.register(name = \"MTR\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"PTR\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"MLR\", formatter='{:.2f}'.format)\n\n\n\t\tself.register(name = \"n_gt_trajectories\", display_name = \"GT\",formatter='{:.0f}'.format, write_mail = True)\n\n\t\tself.register(name = \"tp\", display_name=\"TP\", formatter='{:.0f}'.format) # number of true positives\n\t\tself.register(name = \"fp\", display_name=\"FP\", formatter='{:.0f}'.format) # number of false positives\n\t\tself.register(name = \"fn\", display_name=\"FN\", formatter='{:.0f}'.format) # number of false negatives\n\n\t\tself.register(name = \"recall\", display_name=\"Rcll\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"precision\", display_name=\"Prcn\", formatter='{:.2f}'.format)\n\n\t\tself.register(name = \"F1\", display_name=\"F1\", formatter='{:.2f}'.format, write_mail = False)\n\t\tself.register(name = \"FAR\", formatter='{:.2f}'.format, write_mail = False)\n\t\tself.register(name = \"total_cost\", display_name=\"COST\", formatter='{:.0f}'.format, write_mail = False)\n\t\tself.register(name = \"fragments\", display_name=\"FM\", formatter='{:.0f}'.format)\n\t\tself.register(name = \"fragments_rel\", display_name=\"FMR\", formatter='{:.2f}'.format)\n\t\tself.register(name = \"id_switches\", display_name=\"IDSW\", formatter='{:.0f}'.format)\n\n\t\tself.register(name = \"id_switches_rel\", display_name=\"IDSWR\", formatter='{:.1f}'.format)\n\n\t\tself.register(name = \"n_tr_trajectories\", display_name = \"TR\", formatter='{:.0f}'.format, write_mail = False)\n\t\tself.register(name = \"total_num_frames\", display_name=\"TOTAL_NUM\", formatter='{:.0f}'.format, write_mail = False)\n\n\n\t\tself.register(name = \"n_gt\", display_name = \"GT_OBJ\", formatter='{:.0f}'.format, write_mail = False) # number of ground truth detections\n\t\tself.register(name = \"n_tr\", display_name = \"TR_OBJ\", formatter='{:.0f}'.format, write_mail = False) # number of tracker detections minus ignored tracker detections\n\t\tself.register(name = \"n_itr\",display_name=\"IGNORED\", formatter='{:.0f}'.format, write_mail = False) # number of ignored tracker detections\n\t\tself.register(name = \"id_n_tr\", display_name = \"ID_TR_OBJ\", formatter='{:.0f}'.format, write_mail = False)\n\t\tself.register(name = \"nbox_gt\", display_name = \"NBOX_GT\", formatter='{:.0f}'.format, write_mail = False)\n\n\n\n\n\t\n\tdef compute_clearmot(self):\n\t # precision/recall etc.\n\t if (self.fp + self.tp) == 0 or (self.tp + self.fn) == 0:\n\t self.recall = 0.\n\t self.precision = 0.\n\t else:\n\t self.recall = self.tp / float(self.tp + self.fn) * 100.\n\t self.precision = self.tp / float(self.fp + self.tp) * 100.\n\t if (self.recall + self.precision) == 0:\n\t self.F1 = 0.\n\t else:\n\t self.F1 = (2. * (self.precision * self.recall) / (self.precision + self.recall) ) * 100.\n\t if self.total_num_frames == 0:\n\t self.FAR = \"n/a\"\n\t else:\n\t self.FAR = (self.fp / float(self.total_num_frames) ) \n\t # compute CLEARMOT\n\t if self.n_gt == 0:\n\t self.MOTSA = -float(\"inf\")\n\t self.MODSA = -float(\"inf\")\n\t self.sMOTSA = -float(\"inf\")\n\t else:\n\t self.MOTSA = (1 - (self.fn + self.fp + self.id_switches) / float(self.n_gt) ) * 100.\n\t self.MODSA = (1 - (self.fn + self.fp) / float(self.n_gt)) * 100.\n\t self.sMOTSA = ((self.total_cost - self.fp - self.id_switches) / float(self.n_gt)) * 100.\n\t if self.tp == 0:\n\t self.MOTSP = float(\"inf\")\n\t else:\n\t self.MOTSP = self.total_cost / float(self.tp) * 100.\n\t if self.n_gt != 0:\n\t if self.id_switches == 0:\n\t self.MOTSAL = (1 - (self.fn + self.fp + self.id_switches) / float(self.n_gt)) * 100.\n\t else:\n\t self.MOTSAL = (1 - (self.fn + self.fp + math.log10(self.id_switches)) / float(\n\t self.n_gt))*100.\n\t else:\n\t self.MOTSAL = -float(\"inf\")\n\t\n\t if self.total_num_frames == 0:\n\t self.MODSP = \"n/a\"\n\t else:\n\t self.MODSP = self.MODSP / float(self.total_num_frames) * 100.\n\n\n\n\n\t if self.n_gt_trajectories == 0:\n\t self.MTR = 0.\n\t self.PTR = 0.\n\t self.MLR = 0.\n\t else:\n\t self.MTR = self.MT * 100. / float(self.n_gt_trajectories)\n\t self.PTR = self.PT * 100. / float(self.n_gt_trajectories)\n\t self.MLR = self.ML * 100. / float(self.n_gt_trajectories)\n\n\n\t # calculate relative IDSW and FM\n\n\t if self.recall != 0:\n\t self.id_switches_rel = self.id_switches/self.recall*100\n\t self.fragments_rel = self.fragments/self.recall* 100 \n\n\t else:\n\t self.id_switches_rel = float(\"inf\")\n\t self.fragments_rel = float(\"inf\")\n\n\n\n\n\t # IDF1\n\t if self.n_gt_trajectories == 0:\n\t self.IDF1 = 0.\n\t else:\n\t self.IDF1 = (2 * self.IDTP) / (self.nbox_gt + self.id_n_tr) * 100.\n\t return self\n\n\n\n\n\t# go through all frames and associate ground truth and tracker results\n\n\tdef compute_metrics_per_sequence(self, sequence, pred_file, gt_file, gtDataDir, benchmark_name,\n\t\t\t\t\t\tignore_class = IGNORE_CLASS, class_id = CLASS_ID, overlap_function = mask_iou):\n\n\t\tgt_seq = load_txt(gt_file)\n\t\tresults_seq = load_txt(pred_file)\n\n\n\t\t# load information about sequence\n\t\timport configparser\n\t\tconfig = configparser.ConfigParser()\n\n\n\t\tconfig.read(os.path.join(gtDataDir, \"seqinfo.ini\"))\n\n\t\tmax_frames = int(config['Sequence'][\"seqlength\"])\n\n\t\tself.total_num_frames = max_frames + 1\n\n\t\tseq_trajectories = defaultdict(list)\n\n\t\t# To count number of track ids\n\t\tgt_track_ids = set()\n\t\ttr_track_ids = set()\n\n\t\t# Statistics over the current sequence\n\t\tseqtp = 0\n\t\tseqfn = 0\n\t\tseqfp = 0\n\t\tseqitr = 0\n\n\t\tn_gts = 0\n\t\tn_trs = 0\n\n\t\tframe_to_ignore_region = {}\n\n\t\t# Iterate over frames in this sequence\n\t\tfor f in range(max_frames + 1):\n\t\t\tg = []\n\t\t\tdc = []\n\t\t\tt = []\n\n\t\t\tif f in gt_seq:\n\t\t\t\tfor obj in gt_seq[f]:\n\t\t\t\t\tif obj.class_id == ignore_class:\n\t\t\t\t\t\tdc.append(obj)\n\t\t\t\t\telif obj.class_id == class_id:\n\t\t\t\t\t\tg.append(obj)\n\t\t\t\t\t\tgt_track_ids.add(obj.track_id)\n\t\t\tif f in results_seq:\n\t\t\t\tfor obj in results_seq[f]:\n\t\t\t\t\tif obj.class_id == class_id:\n\t\t\t\t\t\tt.append(obj)\n\t\t\t\t\t\ttr_track_ids.add(obj.track_id)\n\n\t\t\t# Handle ignore regions as one large ignore region\n\t\t\tdc = SegmentedObject(mask=rletools.merge([d.mask for d in dc], intersect=False),\n\t\t\t class_id=ignore_class, track_id=ignore_class)\n\t\t\tframe_to_ignore_region[f] = dc\n\n\t\t\ttracks_valid = [False for _ in range(len(t))]\n\n\t\t\t# counting total number of ground truth and tracker objects\n\t\t\tself.n_gt += len(g)\n\t\t\tself.n_tr += len(t)\n\n\t\t\tn_gts += len(g)\n\t\t\tn_trs += len(t)\n\n\t\t\t# tmp variables for sanity checks and MODSP computation\n\t\t\ttmptp = 0\n\t\t\ttmpfp = 0\n\t\t\ttmpfn = 0\n\t\t\ttmpc = 0 # this will sum up the overlaps for all true positives\n\t\t\ttmpcs = [0] * len(g) # this will save the overlaps for all true positives\n\t\t\t# the reason is that some true positives might be ignored\n\t\t\t# later such that the corrsponding overlaps can\n\t\t\t# be subtracted from tmpc for MODSP computation\n\n\t\t\t# To associate, simply take for each ground truth the (unique!) detection with IoU>0.5 if it exists\n\n\t\t\t# all ground truth trajectories are initially not associated\n\t\t\t# extend groundtruth trajectories lists (merge lists)\n\t\t\tfor gg in g:\n\t\t\t\tseq_trajectories[gg.track_id].append(-1)\n\t\t\tnum_associations = 0\n\t\t\tfor row, gg in enumerate(g):\n\t\t\t\tfor col, tt in enumerate(t):\n\t\t\t\t\tc = overlap_function(gg, tt)\n\t\t\t\t\tif c > 0.5:\n\t\t\t\t\t\ttracks_valid[col] = True\n\t\t\t\t\t\tself.total_cost += c\n\t\t\t\t\t\ttmpc += c\n\t\t\t\t\t\ttmpcs[row] = c\n\t\t\t\t\t\tseq_trajectories[g[row].track_id][-1] = t[col].track_id\n\n\t\t\t\t\t\t# true positives are only valid associations\n\t\t\t\t\t\tself.tp += 1\n\t\t\t\t\t\ttmptp += 1\n\n\t\t\t\t\t\tnum_associations += 1\n\n\t\t\t# associate tracker and DontCare areas\n\t\t\t# ignore tracker in neighboring classes\n\t\t\tnignoredtracker = 0 # number of ignored tracker detections\n\n\t\t\tfor i, tt in enumerate(t):\n\t\t\t\toverlap = overlap_function(tt, dc, \"a\")\n\t\t\t\tif overlap > 0.5 and not tracks_valid[i]:\n\t\t\t\t\tnignoredtracker += 1\n\n\t\t\t# count the number of ignored tracker objects\n\t\t\tself.n_itr += nignoredtracker\n\n\t\t\t# false negatives = non-associated gt instances\n\t\t\t#\n\t\t\ttmpfn += len(g) - num_associations\n\t\t\tself.fn += len(g) - num_associations\n\n\t\t\t# false positives = tracker instances - associated tracker instances\n\t\t\t# mismatches (mme_t)\n\t\t\ttmpfp += len(t) - tmptp - nignoredtracker\n\t\t\tself.fp += len(t) - tmptp - nignoredtracker\n\t\t\t# tmpfp = len(t) - tmptp - nignoredtp # == len(t) - (tp - ignoredtp) - ignoredtp\n\t\t\t# self.fp += len(t) - tmptp - nignoredtp\n\n\t\t\t# update sequence data\n\t\t\tseqtp += tmptp\n\t\t\tseqfp += tmpfp\n\t\t\tseqfn += tmpfn\n\t\t\tseqitr += nignoredtracker\n\n\t\t\t# sanity checks\n\t\t\t# - the number of true positives minues ignored true positives\n\t\t\t# should be greater or equal to 0\n\t\t\t# - the number of false negatives should be greater or equal to 0\n\t\t\t# - the number of false positives needs to be greater or equal to 0\n\t\t\t# otherwise ignored detections might be counted double\n\t\t\t# - the number of counted true positives (plus ignored ones)\n\t\t\t# and the number of counted false negatives (plus ignored ones)\n\t\t\t# should match the total number of ground truth objects\n\t\t\t# - the number of counted true positives (plus ignored ones)\n\t\t\t# and the number of counted false positives\n\t\t\t# plus the number of ignored tracker detections should\n\t\t\t# match the total number of tracker detections\n\t\t\tif tmptp < 0:\n\t\t\t\tprint(tmptp)\n\t\t\t\traise NameError(\"Something went wrong! TP is negative\")\n\t\t\tif tmpfn < 0:\n\t\t\t\tprint(tmpfn, len(g), num_associations)\n\t\t\t\traise NameError(\"Something went wrong! FN is negative\")\n\t\t\tif tmpfp < 0:\n\t\t\t\tprint(tmpfp, len(t), tmptp, nignoredtracker)\n\t\t\t\traise NameError(\"Something went wrong! FP is negative\")\n\t\t\tif tmptp + tmpfn != len(g):\n\t\t\t\tprint(\"seqname\", seq_name)\n\t\t\t\tprint(\"frame \", f)\n\t\t\t\tprint(\"TP \", tmptp)\n\t\t\t\tprint(\"FN \", tmpfn)\n\t\t\t\tprint(\"FP \", tmpfp)\n\t\t\t\tprint(\"nGT \", len(g))\n\t\t\t\tprint(\"nAss \", num_associations)\n\t\t\t\traise NameError(\"Something went wrong! nGroundtruth is not TP+FN\")\n\t\t\tif tmptp + tmpfp + nignoredtracker != len(t):\n\t\t\t\tprint(seq_name, f, len(t), tmptp, tmpfp)\n\t\t\t\tprint(num_associations)\n\t\t\t\traise NameError(\"Something went wrong! nTracker is not TP+FP\")\n\n\t\t\t# compute MODSP\n\t\t\tMODSP_f = 1\n\t\t\tif tmptp != 0:\n\t\t\t\tMODSP_f = tmpc / float(tmptp)\n\t\t\tself.MODSP += MODSP_f\n\n\t\tassert len(seq_trajectories) == len(gt_track_ids)\n\t\tself.n_gt_trajectories = len(gt_track_ids)\n\t\tself.n_tr_trajectories = len(tr_track_ids)\n\n\t\t# compute MT/PT/ML, fragments, idswitches for all groundtruth trajectories\n\t\tif len(seq_trajectories) != 0:\n\t\t\tfor g in seq_trajectories.values():\n\t\t\t\t# all frames of this gt trajectory are not assigned to any detections\n\t\t\t\tif all([this == -1 for this in g]):\n\t\t\t\t\tself.ML += 1\n\t\t\t\t\tcontinue\n\t\t\t\t# compute tracked frames in trajectory\n\t\t\t\tlast_id = g[0]\n\t\t\t\t# first detection (necessary to be in gt_trajectories) is always tracked\n\t\t\t\ttracked = 1 if g[0] >= 0 else 0\n\t\t\t\tfor f in range(1, len(g)):\n\t\t\t\t\tif last_id != g[f] and last_id != -1 and g[f] != -1:\n\t\t\t\t\t\tself.id_switches += 1\n\t\t\t\t\tif f < len(g) - 1 and g[f - 1] != g[f] and last_id != -1 and g[f] != -1 and g[f + 1] != -1:\n\t\t\t\t\t\tself.fragments += 1\n\t\t\t\t\tif g[f] != -1:\n\t\t\t\t\t\ttracked += 1\n\t\t\t\t\t\tlast_id = g[f]\n\t\t\t\t# handle last frame; tracked state is handled in for loop (g[f]!=-1)\n\t\t\t\tif len(g) > 1 and g[f - 1] != g[f] and last_id != -1 and g[f] != -1:\n\t\t\t\t\tself.fragments += 1\n\n\t\t\t\t# compute MT/PT/ML\n\t\t\t\ttracking_ratio = tracked / float(len(g))\n\t\t\t\tif tracking_ratio > 0.8:\n\t\t\t\t\tself.MT += 1\n\t\t\t\telif tracking_ratio < 0.2:\n\t\t\t\t\tself.ML += 1\n\t\t\t\telse: # 0.2 <= tracking_ratio <= 0.8\n\t\t\t\t\tself.PT += 1\n\n\t\t\t# compute IDF1\n\t\t\tidf1, idtp, nbox_gt, id_n_tr = compute_idf1_and_idtp_for_sequence(gt_seq, results_seq, gt_track_ids, tr_track_ids, frame_to_ignore_region)\n\t\t\tself.IDTP = idtp\n\t\t\t#self.id_ign = id_ign\n\t\t\tself.id_n_tr = id_n_tr\n\t\t\tself.nbox_gt = nbox_gt\n\t\t\t\n\t\treturn self\n\n\n\n### IDF1 stuff\n### code below adapted from https://github.com/shenh10/mot_evaluation/blob/5dd51e5cb7b45992774ea150e4386aa0b02b586f/utils/measurements.py\ndef compute_idf1_and_idtp_for_sequence(frame_to_gt, frame_to_pred, gt_ids, st_ids, frame_to_ignore_region):\n\tframe_to_can_be_ignored = {}\n\tfor t in frame_to_pred.keys():\n\t\tpreds_t = frame_to_pred[t]\n\t\tpred_masks_t = [p.mask for p in preds_t]\n\t\tignore_region_t = frame_to_ignore_region[t].mask\n\t\toverlap = np.squeeze(rletools.iou(pred_masks_t, [ignore_region_t], [1]), axis=1)\n\t\tframe_to_can_be_ignored[t] = overlap > 0.5\n\n\tgt_ids = sorted(gt_ids)\n\tst_ids = sorted(st_ids)\n\tgroundtruth = [[] for _ in gt_ids]\n\tprediction = [[] for _ in st_ids]\n\tfor t, gts_t in frame_to_gt.items():\n\t\tfor gt_t in gts_t:\n\t\t\tif gt_t.track_id in gt_ids:\n\t\t\t\tgroundtruth[gt_ids.index(gt_t.track_id)].append((t, gt_t))\n\tfor t in frame_to_pred.keys():\n\t\tpreds_t = frame_to_pred[t]\n\t\tcan_be_ignored_t = frame_to_can_be_ignored[t]\n\t\tassert len(preds_t) == len(can_be_ignored_t)\n\t\tfor pred_t, ign_t in zip(preds_t, can_be_ignored_t):\n\t\t\tif pred_t.track_id in st_ids:\n\t\t\t\tprediction[st_ids.index(pred_t.track_id)].append((t, pred_t, ign_t))\n\tfor gt in groundtruth:\n\t\tgt.sort(key=lambda x: x[0])\n\tfor pred in prediction:\n\t\tpred.sort(key=lambda x: x[0])\n\n\tn_gt = len(gt_ids)\n\tn_st = len(st_ids)\n\tcost = np.zeros((n_gt + n_st, n_st + n_gt), dtype=float)\n\tcost[n_gt:, :n_st] = sys.maxsize # float('inf')\n\tcost[:n_gt, n_st:] = sys.maxsize # float('inf')\n\n\tfp = np.zeros(cost.shape)\n\tfn = np.zeros(cost.shape)\n\tign = np.zeros(cost.shape)\n\t# cost matrix of all trajectory pairs\n\tcost_block, fp_block, fn_block, ign_block = cost_between_gt_pred(groundtruth, prediction)\n\tcost[:n_gt, :n_st] = cost_block\n\tfp[:n_gt, :n_st] = fp_block\n\tfn[:n_gt, :n_st] = fn_block\n\tign[:n_gt, :n_st] = ign_block\n\n\t# computed trajectory match no groundtruth trajectory, FP\n\tfor i in range(n_st):\n\t\t#cost[i + n_gt, i] = prediction[i].shape[0]\n\t\t#fp[i + n_gt, i] = prediction[i].shape[0]\n\t\t# don't count fp in case of ignore region\n\t\tfps = sum([~x[2] for x in prediction[i]])\n\t\tig = sum([x[2] for x in prediction[i]])\n\t\tcost[i + n_gt, i] = fps\n\t\tfp[i + n_gt, i] = fps\n\t\tign[i + n_gt, i] = ig\n\n\t# groundtruth trajectory match no computed trajectory, FN\n\tfor i in range(n_gt):\n\t\t#cost[i, i + n_st] = groundtruth[i].shape[0]\n\t\t#fn[i, i + n_st] = groundtruth[i].shape[0]\n\t\tcost[i, i + n_st] = len(groundtruth[i])\n\t\tfn[i, i + n_st] = len(groundtruth[i])\n\t# TODO: add error handling here?\n\tmatched_indices = linear_assignment(cost)\n\t#nbox_gt = sum([groundtruth[i].shape[0] for i in range(n_gt)])\n\t#nbox_st = sum([prediction[i].shape[0] for i in range(n_st)])\n\tnbox_gt = sum([len(groundtruth[i]) for i in range(n_gt)])\n\n\tnbox_st = sum([len(prediction[i]) for i in range(n_st)])\n\n\t#IDFP = 0\n\tIDFN = 0\n\tid_ign = 0\n\tfor matched in zip(*matched_indices):\n\t\t#IDFP += fp[matched[0], matched[1]]\n\t\tIDFN += fn[matched[0], matched[1]]\n\t\t# exclude detections which are not matched and ignored from total count\n\t\tid_ign += ign[matched[0], matched[1]]\n\tid_n_tr = nbox_st - id_ign\n\n\tIDTP = nbox_gt - IDFN\n\n\tIDF1 = 2 * IDTP / (nbox_gt + id_n_tr)\n\n\treturn IDF1, IDTP, nbox_gt, id_n_tr\n\n\n\n\ndef cost_between_gt_pred(groundtruth, prediction):\n\tn_gt = len(groundtruth)\n\tn_st = len(prediction)\n\tcost = np.zeros((n_gt, n_st), dtype=float)\n\tfp = np.zeros((n_gt, n_st), dtype=float)\n\tfn = np.zeros((n_gt, n_st), dtype=float)\n\tign = np.zeros((n_gt, n_st), dtype=float)\n\tfor i in range(n_gt):\n\t\tfor j in range(n_st):\n\t\t\tfp[i, j], fn[i, j], ign[i, j] = cost_between_trajectories(groundtruth[i], prediction[j])\n\t\t\tcost[i, j] = fp[i, j] + fn[i, j]\n\treturn cost, fp, fn, ign\n\n\ndef cost_between_trajectories(traj1, traj2):\n\t#[npoints1, dim1] = traj1.shape\n\t#[npoints2, dim2] = traj2.shape\n\tnpoints1 = len(traj1)\n\tnpoints2 = len(traj2)\n\t# find start and end frame of each trajectories\n\t#start1 = traj1[0, 0]\n\t#end1 = traj1[-1, 0]\n\t#start2 = traj2[0, 0]\n\t#end2 = traj2[-1, 0]\n\ttimes1 = [x[0] for x in traj1]\n\ttimes2 = [x[0] for x in traj2]\n\tstart1 = min(times1)\n\tstart2 = min(times2)\n\tend1 = max(times1)\n\tend2 = max(times2)\n\n\tign = [traj2[i][2] for i in range(npoints2)]\n\n\t# check frame overlap\n\t#has_overlap = max(start1, start2) < min(end1, end2)\n\t# careful, changed this to <=, but I think now it's right\n\thas_overlap = max(start1, start2) <= min(end1, end2)\n\tif not has_overlap:\n\t\tfn = npoints1\n\t\t#fp = npoints2\n\t\t# disregard detections which can be ignored\n\t\tfp = sum([~x for x in ign])\n\t\tig = sum(ign)\n\t\treturn fp, fn, ig\n\n\t# gt trajectory mapping to st, check gt missed\n\tmatched_pos1 = corresponding_frame(times1, npoints1, times2, npoints2)\n\t# st trajectory mapping to gt, check computed one false alarms\n\tmatched_pos2 = corresponding_frame(times2, npoints2, times1, npoints1)\n\toverlap1 = compute_overlap(traj1, traj2, matched_pos1)\n\toverlap2 = compute_overlap(traj2, traj1, matched_pos2)\n\t# FN\n\tfn = sum([1 for i in range(npoints1) if overlap1[i] < 0.5])\n\t# FP\n\t# don't count false positive in case of ignore region\n\tunmatched = [overlap2[i] < 0.5 for i in range(npoints2)]\n\t#fp = sum([1 for i in range(npoints2) if overlap2[i] < 0.5 and not traj2[i][2]])\n\tfp = sum([1 for i in range(npoints2) if unmatched[i] and not ign[i]])\n\tig = sum([1 for i in range(npoints2) if unmatched[i] and ign[i]])\n\treturn fp, fn, ig\n\n\ndef corresponding_frame(traj1, len1, traj2, len2):\n\t\"\"\"\n\tFind the matching position in traj2 regarding to traj1\n\tAssume both trajectories in ascending frame ID\n\t\"\"\"\n\tp1, p2 = 0, 0\n\tloc = -1 * np.ones((len1,), dtype=int)\n\twhile p1 < len1 and p2 < len2:\n\t\tif traj1[p1] < traj2[p2]:\n\t\t\tloc[p1] = -1\n\t\t\tp1 += 1\n\t\telif traj1[p1] == traj2[p2]:\n\t\t\tloc[p1] = p2\n\t\t\tp1 += 1\n\t\t\tp2 += 1\n\t\telse:\n\t\t\tp2 += 1\n\treturn loc\n\n\ndef compute_overlap(traj1, traj2, matched_pos):\n\t\"\"\"\n\tCompute the loss hit in traj2 regarding to traj1\n\t\"\"\"\n\toverlap = np.zeros((len(matched_pos),), dtype=float)\n\tfor i in range(len(matched_pos)):\n\t\tif matched_pos[i] == -1:\n\t\t\tcontinue\n\t\telse:\n\t\t\tmask1 = traj1[i][1].mask\n\t\t\tmask2 = traj2[matched_pos[i]][1].mask\n\t\t\tiou = rletools.iou([mask1], [mask2], [False])[0][0]\n\t\t\toverlap[i] = iou\n\treturn overlap\n\n"
},
{
"path": "eval/mots/Metrics.py",
"content": "from __future__ import division\nfrom collections import OrderedDict, Iterable\nimport pandas as pd\nimport numpy as np\nimport pickle\n\nclass Metrics(object):\n\tdef __init__(self):\n\t\tself.metrics = OrderedDict()\n\t\tself.cache_dict = OrderedDict()\n\n\n\tdef register(self, name=None, value=None, formatter=None,\n\t\t\t\tdisplay_name=None, write_db = True, write_mail = True):\n\t\t\"\"\"Register a new metric.\n\t\tParams\n\t\t------\n\t\tname: str\n\t\t\tName of the metric. Name is used for computation and set as attribute.\n display_name: str or None\n Disoplay name of variable written in db and mail\n\t\tvalue:\n\t\tformatter:\n\t\t\tFormatter to present value of metric. E.g. `'{:.2f}'.format`\n\t\twrite_db: boolean, default = True\n\t\t\tWrite value into db\n\t\twrite_mail: boolean, default = True\n\t\t\tWrite metric in result mail to user\n\t\t\"\"\"\n\t\tassert not name is None, 'No name specified'.format(name)\n\n\t\tif not value:\n\t\t\tvalue = 0\n\n\t\tself.__setattr__( name, value)\n\n\t\tif not display_name: display_name = name\n\t\tself.metrics[name] = {\n\t\t 'name' : name,\n\t\t 'write_db' : write_db,\n\t\t 'formatter' : formatter,\n\t\t 'write_mail' : write_mail,\n\t\t 'display_name' : display_name\n\t\t}\n\n\tdef cache(self, name=None, value=None, func=None):\n\t\tassert not name is None, 'No name specified'.format(name)\n\n\n\t\tself.__setattr__( name, value)\n\n\n\t\tself.cache_dict[name] = {\n\t\t 'name' : name,\n\t\t 'func' : func\n\t\t}\n\n\n\n\tdef __call__(self, name):\n\t\treturn self.metrics[name]\n\n\t@property\n\tdef names(self):\n\t\t\"\"\"Returns the name identifiers of all registered metrics.\"\"\"\n\t\treturn [v['name'] for v in self.metrics.values()]\n\n\t@property\n\tdef display_names(self):\n\t\t\"\"\"Returns the display name identifiers of all registered metrics.\"\"\"\n\t\treturn [v['display_name'] for v in self.metrics.values()]\n\n\n\t@property\n\tdef formatters(self):\n\t\t\"\"\"Returns the formatters for all metrics that have associated formatters.\"\"\"\n\t\treturn dict([(v['display_name'], v['formatter']) for k, v in self.metrics.items() if not v['formatter'] is None])\n\n\t#@property\n\tdef val_dict(self, display_name = False, object = \"metrics\"):\n\t\t\"\"\"Returns dictionary of all registered values of object name or display_name as key.\n\t\tParams\n ------\n\n display_name: boolean, default = False\n If True, display_name of keys in dict. (default names)\n object: \"cache\" or \"metrics\", default = \"metrics\"\n\t\t\"\"\"\n\t\tif display_name: key_string = \"display_name\"\n\t\telse: key_string = \"name\"\n\t\tprint(\"object dict: \", object)\n\t\tval_dict = dict([(self.__getattribute__(object)[key][key_string], self.__getattribute__(key)) for key in self.__getattribute__(object).keys() ])\n\t\treturn val_dict\n\n\tdef val_db(self, display_name = True):\n\t\t\"\"\"Returns dictionary of all registered values metrics to write in db.\"\"\"\n\t\tif display_name: key_string = \"display_name\"\n\t\telse: key_string = \"name\"\n\t\tval_dict = dict([(self.metrics[key][key_string], self.__getattribute__(key)) for key in self.metrics.keys() if self.metrics[key][\"write_db\"] ])\n\t\treturn val_dict\n\n\n\tdef val_mail(self, display_name = True):\n\t\t\"\"\"Returns dictionary of all registered values metrics to write in mail.\"\"\"\n\t\tif display_name: key_string = \"display_name\"\n\t\telse: key_string = \"name\"\n\t\tval_dict = dict([(self.metrics[key][key_string], self.__getattribute__(key)) for key in self.metrics.keys() if self.metrics[key][\"write_mail\"] ])\n\t\treturn val_dict\n\n\n\tdef to_dataframe(self, display_name = False, type = None):\n\t\t\"\"\"Returns pandas dataframe of all registered values metrics. \"\"\"\n\t\tif type==\"mail\":\n\t\t\tself.df = pd.DataFrame(self.val_mail(display_name = display_name), index=[self.seqName])\n\t\telse:\n\t\t\tself.df = pd.DataFrame(self.val_dict(display_name = display_name), index=[self.seqName])\n\tdef update_values(self, value_dict = None):\n\t\t\"\"\"Updates registered metrics with new values in value_dict. \"\"\"\n\t\tif value_dict:\n\t\t\tfor key, value in value_dict.items() :\n\t\t\t\tif hasattr(self, key):\n\t\t\t\t\tself.__setattr__(key, value)\n\t\t\t\t\n\n\tdef print_type(self, object = \"metrics\"):\n\t\t\"\"\"Prints variable type of registered metrics or caches. \"\"\"\n\t\tprint( \"OBJECT \" , object)\n\t\tval_dict = self.val_dict(object = object)\n\t\tfor key, item in val_dict.items() :\n\t\t\tprint(\"%s: %s; Shape: %s\" %(key, type(item), np.shape(item)))\n\n\tdef print_results(self):\n\t\t\"\"\"Prints metrics. \"\"\"\n\t\tresult_dict = self.val_dict()\n\t\tfor key, item in result_dict.items():\n\t\t\tprint(key)\n\t\t\tprint(\"%s: %s\" %(key, self.metrics[key][\"formatter\"](item)))\n\n\n\tdef save_dict(self, path):\n\t\t\"\"\"Save value dict to path as pickle file.\"\"\"\n\t\twith open(path, 'wb') as handle:\n\t\t\tpickle.dump(self.__dict__, handle, protocol=pickle.HIGHEST_PROTOCOL)\n\n\n\tdef compute_metrics_per_sequence(self):\n\t\traise NotImplementedError\n\n\n"
},
{
"path": "eval/mots/README.md",
"content": "# MOTS\r\n\r\n\r\n## Requirements\r\n* Python 3.6.9\r\n* install [requirements.txt](requirements.txt)\r\n\r\n## Usage\r\n1) Run \r\n```\r\npython MOTS/evalMOTS.py\r\n```\r\n\r\n\r\n\r\n## Evaluation\r\nTo run the evaluation for your method please adjust the file ```MOTS/evalMOTS.py``` using the following arguments:\r\n\r\n```benchmark_name```: Name of the benchmark, e.g. MOTS \r\n```gt_dir```: Directory containing ground truth files in ``` ![](\"docs/logo.png\"){kind=logo}
\n\n--------------------------------------------------------------------------------\n\n**[NeurIPS 2021] Do different tracking tasks require different appearance model?**\n\n**[[ArXiv](https://arxiv.org/abs/2107.02156)]** **[[Project Page](https://zhongdao.github.io/UniTrack)]**\n\nUniTrack is a simple and Unified framework for addressing multiple tracking tasks. \n\nBeing a fundamental problem in computer vision, tracking has been fragmented into a multitude of different experimental setups. As a consequence, the literature has fragmented too, and now the novel approaches proposed by the community are usually specialized to fit only one specific setup. To understand to what extent this specialization is actually necessary, we present UniTrack, a solution to address multiple different tracking tasks within the same framework. All tasks share the same [appearance model](#appearance-model). UniTrack\n\n- Does **NOT** need training on a specific tracking task.\n\n- Shows [competitive performance](docs/RESULTS.md) on six out of seven tracking tasks considered.\n\n- Can be easily adapted to even [more tasks](##Demo).\n\n- Can be used as an evaluation platform to [test pre-trained self-supervised models](docs/MODELZOO.md).\n \n \n## Demo\n**Multi-Object Tracking demo for 80 COCO classes ([YOLOX](https://github.com/Megvii-BaseDetection/YOLOX) + UniTrack)**\n![](\"docs/unitrack_yolox.gif\")
\n\nIn this demo we run the YOLOX detector and perform MOT for the 80 COCO classes. Try the demo by:\n```python\npython demo/mot_demo.py --classes cls1 cls2 ... clsN\n```\nwhere cls1 to clsN represent the indices of classes you would like to detect and track. See [here](https://gist.github.com/AruniRC/7b3dadd004da04c80198557db5da4bda) for the index list. By default all 80 classes are detected and tracked.\n \n**Single-Object Tracking demo for custom videos**\n```python\npython demo/sot_demo.py --config ./config/imagenet_resnet18_s3.yaml --input /path/to/your/video\n```\nIn this demo, you are asked to annotate the target to be tracked, by drawing a rectangle in the first frame of the video. Then the algorithm tracks the target in following timesteps without object detection.\n \n## Tasks & Framework\n\n\n### Tasks\nWe classify existing tracking tasks along four axes: (1) Single or multiple targets; (2) Users specify targets or automatic detectors specify targets; (3) Observation formats (bounding box/mask/pose); (2) Class-agnostic or class-specific (i.e. human/vehicles). We mainly experiment on 5 tasks: **SOT, VOS, MOT, MOTS, and PoseTrack**. Task setups are summarized in the above figure.\n\n### Appearance model\nAn appearance model is the only learnable component in UniTrack. It should provide universal visual representation, and is usually pre-trained on large-scale dataset in supervised or unsupervised manners. Typical examples include ImageNet pre-trained ResNets (supervised), and recent self-supervised models such as MoCo and SimCLR (unsupervised).\n\n### Propagation and Association\n*Propagation* and *Association* are the two core primitives used in UniTrack to address a wide variety of tracking tasks (currently 7, but more can be added), Both use the features extracted by the pre-trained appearance model. For propagation, we adopt exiting methods such as [cross correlation](https://www.robots.ox.ac.uk/~luca/siamese-fc.html), [DCF](https://openaccess.thecvf.com/content_cvpr_2017/html/Valmadre_End-To-End_Representation_Learning_CVPR_2017_paper.html), and [mask propation](https://github.com/ajabri/videowalk). For association we employ a simple algorithm as in [JDE](https://github.com/Zhongdao/Towards-Realtime-MOT) and develop a novel reconstruction-based similairty metric that allows to compare objects across shapes and sizes.\n \n \n## Getting started\n\n1. Installation: Please check out [docs/INSTALL.md](docs/INSTALL.md)\n2. Data preparation: Please check out [docs/DATA.md](docs/DATA.md)\n3. Appearance model preparation: Please check out [docs/MODELZOO.md](docs/MODELZOO.md)\n4. Run evaluation on all datasets: Please check out [docs/RUN.md](docs/RUN.md)\n\n## Results\nBelow we show results of UniTrack with a simple **ImageNet Pre-trained ResNet-18** as the appearance model. More results can be found in [RESULTS.md](docs/RESULTS.md).\n\n**Single Object Tracking (SOT) on OTB-2015**\n\n![](\"docs/sot1.gif\")
![](\"docs/sot2.gif\")
\n\n**Video Object Segmentation (VOS) on DAVIS-2017 *val* split**\n\n![](\"docs/vos1.gif\")
![](\"docs/vos2.gif\")
\n\n**Multiple Object Tracking (MOT) on MOT-16 [*test* set *private detector* track](https://motchallenge.net/method/MOT=3856&chl=5)** (Detections from FairMOT)\n\n![](\"docs/MOT1.gif\")
![](\"docs/MOT2.gif\")
\n\n**Multiple Object Tracking and Segmentation (MOTS) on MOTS challenge [*test* set](https://motchallenge.net/method/MOTS=109&chl=17)** (Detections from COSTA_st)\n\n![](\"docs/MOTS1.gif\")
![](\"docs/MOTS2.gif\")
\n\n**Pose Tracking on PoseTrack-2018 *val* split** (Detections from LightTrack)\n\n![](\"docs/posetrack1.gif\")
![](\"docs/posetrack2.gif\")
\n\n## Acknowledgement\nA part of code is borrowed from \n \n[VideoWalk](https://github.com/ajabri/videowalk) by Allan A. Jabri\n\n[SOT code](https://github.com/JudasDie/SOTS) by Zhipeng Zhang\n \n## Citation\n```bibtex\n@article{wang2021different,\n author = {Wang, Zhongdao and Zhao, Hengshuang and Li, Ya-Li and Wang, Shengjin and Torr, Philip and Bertinetto, Luca},\n title = {Do different tracking tasks require different appearance models?},\n journal = {Thirty-Fifth Conference on Neural Infromation Processing Systems},\n year = {2021},\n}\n```\n"
},
{
"path": "config/crw_resnet18_s3.yaml",
"content": "common:\n exp_name: crw_resnet18_s3\n \n # Model related\n model_type: crw\n remove_layers: ['layer4']\n im_mean: [0.4914, 0.4822, 0.4465]\n im_std: [0.2023, 0.1994, 0.2010]\n nopadding: False\n head_depth: -1\n resume: 'weights/crw.pth'\n \n # Misc\n down_factor: 8\n infer2D: True \n workers: 4\n gpu_id: 0\n device: cuda\n\nsot:\n dataset: 'OTB2015'\n dataroot: '/home/wangzd/datasets/GOT/OTB100/'\n epoch_test: False\n\nvos:\n davisroot: '/home/wangzd/datasets/uvc/DAVIS/'\n split: 'val'\n temperature: 0.05\n topk: 10\n radius: 12\n videoLen: 5 \n cropSize: -1\n head_depth: -1\n no_l2: False\n long_mem: [0]\n infer2D: False\n norm_mask: False\n\nmot:\n obid: 'FairMOT'\n mot_root: '/home/wangzd/datasets/MOT/MOT16'\n feat_size: [4,10]\n save_videos: True\n save_images: False\n test_mot16: False\n track_buffer: 30\n min_box_area: 200\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7\n img_size: [1088, 608]\n prop_flag: False\n use_kalman: True \n asso_with_motion: True \n motion_lambda: 0.98\n motion_gated: True \n\nmots:\n obid: 'COSTA'\n mots_root: '/home/wangzd/datasets/GOT/MOTS'\n save_videos: False\n save_images: True\n test: False\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n prop_flag: False\n max_mask_area: 200\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7 \n first_stage_thres: 0.7\n feat_size: [4,10]\n use_kalman: True \n asso_with_motion: True \n motion_lambda: 0.98 \n motion_gated: False\n\nposetrack:\n obid: 'lighttrack_MSRA152'\n data_root: '/home/wangzd/datasets/GOT/Posetrack2018'\n split: 'val'\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n frame_rate: 6\n save_videos: False\n save_images: True\n prop_flag: False\n feat_size: [4,10]\n max_mask_area: 400\n dup_iou_thres: 0.2\n confirm_iou_thres: 0.6\n first_stage_thres: 0.7\n use_kalman: True \n asso_with_motion: True \n motion_lambda: 0.9999\n motion_gated: False\n only_position: True\n\n"
},
{
"path": "config/crw_resnet18_s3_womotion.yaml",
"content": "common:\n exp_name: crw_resnet18_s3_womotion\n \n # Model related\n model_type: crw\n remove_layers: ['layer4']\n im_mean: [0.4914, 0.4822, 0.4465]\n im_std: [0.2023, 0.1994, 0.2010]\n nopadding: False\n head_depth: -1\n resume: 'weights/crw.pth'\n \n # Misc\n down_factor: 8\n infer2D: True \n workers: 4\n gpu_id: 0\n device: cuda\n\nsot:\n dataset: 'OTB2015'\n dataroot: '/home/wangzd/datasets/GOT/OTB100/'\n epoch_test: False\n\nvos:\n davisroot: '/home/wangzd/datasets/uvc/DAVIS/'\n split: 'val'\n temperature: 0.05\n topk: 10\n radius: 12\n videoLen: 5 \n cropSize: -1\n head_depth: -1\n no_l2: False\n long_mem: [0]\n infer2D: False\n norm_mask: False\n\nmot:\n obid: 'FairMOT'\n mot_root: '/home/wangzd/datasets/MOT/MOT16'\n feat_size: [4,10]\n save_videos: True\n save_images: False\n test_mot16: False\n track_buffer: 30\n min_box_area: 200\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7\n img_size: [1088, 608]\n prop_flag: False\n use_kalman: True \n asso_with_motion: False \n motion_lambda: 1\n motion_gated: False\n\nmots:\n obid: 'COSTA'\n mots_root: '/home/wangzd/datasets/GOT/MOTS'\n save_videos: False\n save_images: True\n test: False\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n prop_flag: False\n max_mask_area: 200\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7 \n first_stage_thres: 0.7\n feat_size: [4,10]\n use_kalman: True \n asso_with_motion: False \n motion_lambda: 1 \n motion_gated: False\n\nposetrack:\n obid: 'lighttrack_MSRA152'\n data_root: '/home/wangzd/datasets/GOT/Posetrack2018'\n split: 'val'\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n frame_rate: 6\n save_videos: False\n save_images: True\n prop_flag: False\n feat_size: [4,10]\n max_mask_area: 400\n dup_iou_thres: 0.2\n confirm_iou_thres: 0.6\n first_stage_thres: 0.7\n use_kalman: True \n asso_with_motion: False \n motion_lambda: 1 \n motion_gated: False\n only_position: True\n\n"
},
{
"path": "config/imagenet_resnet18_s3.yaml",
"content": "common:\n exp_name: imagenet_resnet18_s3 \n \n # Model related\n model_type: imagenet18\n remove_layers: ['layer4']\n im_mean: [0.485, 0.456, 0.406]\n im_std: [0.229, 0.224, 0.225]\n nopadding: False\n resume: None\n \n # Misc\n down_factor: 8\n infer2D: True \n workers: 4\n gpu_id: 0\n device: cuda\n\nsot:\n dataset: 'OTB2015'\n dataroot: '/home/wangzd/datasets/GOT/OTB100/'\n epoch_test: False\n\nvos:\n davisroot: '/home/wangzd/datasets/uvc/DAVIS/'\n split: 'val'\n temperature: 0.05\n topk: 10\n radius: 12\n videoLen: 5 \n cropSize: -1\n head_depth: -1\n no_l2: False\n long_mem: [0]\n infer2D: False\n norm_mask: False\n\nmot:\n obid: 'FairMOT'\n mot_root: '/home/wangzd/datasets/MOT/MOT16'\n feat_size: [4,10]\n save_videos: True\n save_images: False\n test_mot16: False\n track_buffer: 30\n min_box_area: 200\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7\n img_size: [1088, 608]\n prop_flag: False\n use_kalman: True \n asso_with_motion: True \n motion_lambda: 0.98\n motion_gated: True \n\nmots:\n obid: 'COSTA'\n mots_root: '/home/wangzd/datasets/GOT/MOTS'\n save_videos: False\n save_images: True\n test: False\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n prop_flag: False\n max_mask_area: 200\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7 \n first_stage_thres: 0.7\n feat_size: [4,10]\n use_kalman: True \n asso_with_motion: True \n motion_lambda: 0.98 \n motion_gated: False\n\nposetrack:\n obid: 'lighttrack_MSRA152'\n data_root: '/home/wangzd/datasets/GOT/Posetrack2018'\n split: 'val'\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n frame_rate: 6\n save_videos: False\n save_images: True\n prop_flag: False\n feat_size: [4,10]\n max_mask_area: 400\n dup_iou_thres: 0.2\n confirm_iou_thres: 0.6\n first_stage_thres: 0.7\n use_kalman: True \n asso_with_motion: True \n motion_lambda: 0.9999\n motion_gated: False\n only_position: True\n\n"
},
{
"path": "config/imagenet_resnet18_s3_womotion.yaml",
"content": "common:\n exp_name: imagenet_resnet18_s3_womotion\n \n # Model related\n model_type: imagenet18\n remove_layers: ['layer4']\n im_mean: [0.485, 0.456, 0.406]\n im_std: [0.229, 0.224, 0.225]\n nopadding: False\n resume: None\n \n # Misc\n down_factor: 8\n infer2D: True \n workers: 4\n gpu_id: 0\n device: cuda\n\nsot:\n dataset: 'OTB2015'\n dataroot: '/home/wangzd/datasets/GOT/OTB100/'\n epoch_test: False\n\nvos:\n davisroot: '/home/wangzd/datasets/uvc/DAVIS/'\n split: 'val'\n temperature: 0.05\n topk: 10\n radius: 12\n videoLen: 5 \n cropSize: -1\n head_depth: -1\n no_l2: False\n long_mem: [0]\n infer2D: False\n norm_mask: False\n\nmot:\n obid: 'FairMOT'\n mot_root: '/home/wangzd/datasets/MOT/MOT16'\n feat_size: [4,10]\n save_videos: True\n save_images: False\n test_mot16: False\n track_buffer: 30\n min_box_area: 200\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7\n img_size: [1088, 608]\n prop_flag: False\n use_kalman: True \n asso_with_motion: False \n motion_lambda: 1\n motion_gated: False\n\nmots:\n obid: 'COSTA'\n mots_root: '/home/wangzd/datasets/GOT/MOTS'\n save_videos: False\n save_images: True\n test: False\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n prop_flag: False\n max_mask_area: 200\n dup_iou_thres: 0.15\n confirm_iou_thres: 0.7 \n first_stage_thres: 0.7\n feat_size: [4,10]\n use_kalman: True \n asso_with_motion: False\n motion_lambda: 1 \n motion_gated: False\n\nposetrack:\n obid: 'lighttrack_MSRA152'\n data_root: '/home/wangzd/datasets/GOT/Posetrack2018'\n split: 'val'\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n frame_rate: 6\n save_videos: False\n save_images: True\n prop_flag: False\n feat_size: [4,10]\n max_mask_area: 400\n dup_iou_thres: 0.2\n confirm_iou_thres: 0.6\n first_stage_thres: 0.7\n use_kalman: True\n asso_with_motion: False\n motion_lambda: 1 \n motion_gated: False\n only_position: True\n\nvis:\n obid: 'MaskTrackRCNN'\n data_root: '/home/wangzd/datasets/GOT/YoutubeVIS/'\n split: 'val'\n track_buffer: 30\n nms_thres: 0.4\n conf_thres: 0.5\n iou_thres: 0.5\n frame_rate: 6\n save_videos: False\n save_images: True\n prop_flag: False\n feat_size: [12,12]\n max_mask_area: 1000\n dup_iou_thres: 0.2\n confirm_iou_thres: 0.6\n first_stage_thres: 0.9\n use_kalman: True\n asso_with_motion: False\n motion_lambda: 1 \n motion_gated: False\n\n\n\n \n\n\n\n\n\n"
},
{
"path": "core/association/__init__.py",
"content": ""
},
{
"path": "core/association/matching.py",
"content": "import pdb\r\nimport cv2\r\nimport torch\r\nimport torch.nn.functional as F\r\nimport numpy as np\r\nimport scipy\r\nfrom scipy.spatial.distance import cdist\r\nimport lap\r\n\r\nfrom cython_bbox import bbox_overlaps as bbox_ious\r\nfrom core.motion import kalman_filter\r\nimport time\r\n\r\ndef merge_matches(m1, m2, shape):\r\n O,P,Q = shape\r\n m1 = np.asarray(m1)\r\n m2 = np.asarray(m2)\r\n\r\n M1 = scipy.sparse.coo_matrix((np.ones(len(m1)), (m1[:, 0], m1[:, 1])), shape=(O, P))\r\n M2 = scipy.sparse.coo_matrix((np.ones(len(m2)), (m2[:, 0], m2[:, 1])), shape=(P, Q))\r\n\r\n mask = M1*M2\r\n match = mask.nonzero()\r\n match = list(zip(match[0], match[1]))\r\n unmatched_O = tuple(set(range(O)) - set([i for i, j in match]))\r\n unmatched_Q = tuple(set(range(Q)) - set([j for i, j in match]))\r\n\r\n return match, unmatched_O, unmatched_Q\r\n\r\n\r\ndef linear_assignment(cost_matrix, thresh):\r\n if cost_matrix.size == 0:\r\n return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))\r\n matches, unmatched_a, unmatched_b = [], [], []\r\n cost, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)\r\n for ix, mx in enumerate(x):\r\n if mx >= 0:\r\n matches.append([ix, mx])\r\n unmatched_a = np.where(x < 0)[0]\r\n unmatched_b = np.where(y < 0)[0]\r\n matches = np.asarray(matches)\r\n return matches, unmatched_a, unmatched_b\r\n \r\n\r\ndef ious(atlbrs, btlbrs):\r\n \"\"\"\r\n Compute cost based on IoU\r\n :type atlbrs: list[tlbr] | np.ndarray\r\n :type atlbrs: list[tlbr] | np.ndarray\r\n\r\n :rtype ious np.ndarray\r\n \"\"\"\r\n ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float)\r\n if ious.size == 0:\r\n return ious\r\n\r\n ious = bbox_ious(\r\n np.ascontiguousarray(atlbrs, dtype=np.float),\r\n np.ascontiguousarray(btlbrs, dtype=np.float)\r\n )\r\n\r\n return ious\r\n\r\n\r\ndef iou_distance(atracks, btracks):\r\n \"\"\"\r\n Compute cost based on IoU\r\n :type atracks: list[STrack]\r\n :type btracks: list[STrack]\r\n\r\n :rtype cost_matrix np.ndarray\r\n \"\"\"\r\n\r\n if (len(atracks)>0 and isinstance(atracks[0], np.ndarray)) or (len(btracks) > 0 and isinstance(btracks[0], np.ndarray)):\r\n atlbrs = atracks\r\n btlbrs = btracks\r\n else:\r\n atlbrs = [track.tlbr for track in atracks]\r\n btlbrs = [track.tlbr for track in btracks]\r\n _ious = ious(atlbrs, btlbrs)\r\n cost_matrix = 1 - _ious\r\n\r\n return cost_matrix\r\n\r\ndef embedding_distance(tracks, detections, metric='cosine'):\r\n \"\"\"\r\n :param tracks: list[STrack]\r\n :param detections: list[BaseTrack]\r\n :param metric:\r\n :return: cost_matrix np.ndarray\r\n \"\"\"\r\n\r\n cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float)\r\n if cost_matrix.size == 0:\r\n return cost_matrix\r\n det_features = np.asarray([track.curr_feat for track in detections], dtype=np.float)\r\n track_features = np.asarray([track.smooth_feat for track in tracks], dtype=np.float)\r\n cost_matrix = np.maximum(0.0, cdist(track_features, det_features)) # Nomalized features\r\n return cost_matrix\r\n\r\n\r\ndef fuse_motion(kf, cost_matrix, tracks, detections, only_position=False, lambda_=0.98, gate=True):\r\n if cost_matrix.size == 0:\r\n return cost_matrix\r\n gating_dim = 2 if only_position else 4\r\n gating_threshold = kalman_filter.chi2inv95[gating_dim]\r\n measurements = np.asarray([det.to_xyah() for det in detections])\r\n for row, track in enumerate(tracks):\r\n gating_distance = kf.gating_distance(\r\n track.mean, track.covariance, measurements, only_position, metric='maha')\r\n if gate:\r\n cost_matrix[row, gating_distance > gating_threshold] = np.inf\r\n cost_matrix[row] = lambda_ * cost_matrix[row] + (1-lambda_)* gating_distance\r\n return cost_matrix\r\n\r\n\r\ndef center_emb_distance(tracks, detections, metric='cosine'):\r\n \"\"\"\r\n :param tracks: list[STrack]\r\n :param detections: list[BaseTrack]\r\n :param metric:\r\n :return: cost_matrix np.ndarray\r\n \"\"\"\r\n\r\n cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float)\r\n if cost_matrix.size == 0:\r\n return cost_matrix\r\n det_features = torch.stack([track.curr_feat.squeeze() for track in detections])\r\n track_features = torch.stack([track.smooth_feat.squeeze() for track in tracks])\r\n normed_det = F.normalize(det_features)\r\n normed_track = F.normalize(track_features)\r\n cost_matrix = torch.mm(normed_track, normed_det.T)\r\n cost_matrix = 1 - cost_matrix.detach().cpu().numpy()\r\n return cost_matrix\r\n\r\ndef recons_distance(tracks, detections, tmp=100):\r\n \"\"\"\r\n :param tracks: list[STrack]\r\n :param detections: list[BaseTrack]\r\n :param metric:\r\n :return: cost_matrix np.ndarray\r\n \"\"\"\r\n\r\n cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float)\r\n if cost_matrix.size == 0:\r\n return cost_matrix\r\n det_features_ = torch.stack([track.curr_feat.squeeze() for track in detections])\r\n track_features_ = torch.stack([track.smooth_feat for track in tracks])\r\n det_features = F.normalize(det_features_, dim=1)\r\n track_features = F.normalize(track_features_, dim=1)\r\n\r\n ndet, ndim, nw, nh = det_features.shape\r\n ntrk, _, _, _ = track_features.shape\r\n fdet = det_features.permute(0,2,3,1).reshape(-1, ndim).cuda() # ndet*nw*nh, ndim\r\n ftrk = track_features.permute(0,2,3,1).reshape(-1, ndim).cuda() # ntrk*nw*nh, ndim\r\n\r\n aff = torch.mm(ftrk, fdet.transpose(0,1)) # ntrk*nw*nh, ndet*nw*nh\r\n aff_td = F.softmax(tmp*aff, dim=1)\r\n aff_dt = F.softmax(tmp*aff, dim=0).transpose(0,1)\r\n\r\n recons_ftrk = torch.einsum('tds,dsm->tdm', aff_td.view(ntrk*nw*nh, ndet, nw*nh), \r\n fdet.view(ndet, nw*nh, ndim)) # ntrk*nw*nh, ndet, ndim\r\n recons_fdet = torch.einsum('dts,tsm->dtm', aff_dt.view(ndet*nw*nh, ntrk, nw*nh),\r\n ftrk.view(ntrk, nw*nh, ndim)) # ndet*nw*nh, ntrk, ndim\r\n \r\n res_ftrk = (recons_ftrk.permute(0,2,1) - ftrk.unsqueeze(-1)).view(ntrk, nw*nh*ndim, ndet)\r\n res_fdet = (recons_fdet.permute(0,2,1) - fdet.unsqueeze(-1)).view(ndet, nw*nh*ndim, ntrk)\r\n\r\n cost_matrix = (torch.abs(res_ftrk).mean(1) + torch.abs(res_fdet).mean(1).transpose(0,1)) * 0.5\r\n cost_matrix = cost_matrix / cost_matrix.max(1)[0].unsqueeze(-1) \r\n #pdb.set_trace()\r\n cost_matrix = cost_matrix.cpu().numpy()\r\n return cost_matrix\r\n\r\n\r\ndef get_track_feat(tracks, feat_flag='curr'):\r\n if feat_flag == 'curr':\r\n feat_list = [track.curr_feat.squeeze(0) for track in tracks]\r\n elif feat_flag == 'smooth':\r\n feat_list = [track.smooth_feat.squeeze(0) for track in tracks]\r\n else:\r\n raise NotImplementedError\r\n \r\n n = len(tracks)\r\n fdim = feat_list[0].shape[0]\r\n fdim_num = len(feat_list[0].shape)\r\n if fdim_num > 2:\r\n feat_list = [f.view(fdim,-1) for f in feat_list]\r\n numels = [f.shape[1] for f in feat_list]\r\n \r\n ret = torch.zeros(n, fdim, np.max(numels)).to(feat_list[0].device)\r\n for i, f in enumerate(feat_list):\r\n ret[i, :, :numels[i]] = f\r\n return ret \r\n\r\ndef reconsdot_distance(tracks, detections, tmp=100):\r\n \"\"\"\r\n :param tracks: list[STrack]\r\n :param detections: list[BaseTrack]\r\n :param metric:\r\n :return: cost_matrix np.ndarray\r\n \"\"\"\r\n cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float)\r\n if cost_matrix.size == 0:\r\n return cost_matrix, None\r\n det_features_ = get_track_feat(detections)\r\n track_features_ = get_track_feat(tracks, feat_flag='curr')\r\n\r\n det_features = F.normalize(det_features_, dim=1)\r\n track_features = F.normalize(track_features_, dim=1)\r\n\r\n ndet, ndim, nsd = det_features.shape\r\n ntrk, _, nst = track_features.shape\r\n\r\n fdet = det_features.permute(0, 2, 1).reshape(-1, ndim).cuda()\r\n ftrk = track_features.permute(0, 2, 1).reshape(-1, ndim).cuda()\r\n\r\n aff = torch.mm(ftrk, fdet.transpose(0, 1))\r\n aff_td = F.softmax(tmp*aff, dim=1)\r\n aff_dt = F.softmax(tmp*aff, dim=0).transpose(0, 1)\r\n\r\n recons_ftrk = torch.einsum('tds,dsm->tdm', aff_td.view(ntrk*nst, ndet, nsd),\r\n fdet.view(ndet, nsd, ndim))\r\n recons_fdet = torch.einsum('dts,tsm->dtm', aff_dt.view(ndet*nsd, ntrk, nst),\r\n ftrk.view(ntrk, nst, ndim))\r\n\r\n recons_ftrk = recons_ftrk.permute(0, 2, 1).view(ntrk, nst*ndim, ndet)\r\n recons_ftrk_norm = F.normalize(recons_ftrk, dim=1)\r\n recons_fdet = recons_fdet.permute(0, 2, 1).view(ndet, nsd*ndim, ntrk)\r\n recons_fdet_norm = F.normalize(recons_fdet, dim=1)\r\n\r\n dot_td = torch.einsum('tad,ta->td', recons_ftrk_norm,\r\n F.normalize(ftrk.reshape(ntrk, nst*ndim), dim=1))\r\n dot_dt = torch.einsum('dat,da->dt', recons_fdet_norm,\r\n F.normalize(fdet.reshape(ndet, nsd*ndim), dim=1))\r\n\r\n cost_matrix = 1 - 0.5 * (dot_td + dot_dt.transpose(0, 1))\r\n cost_matrix = cost_matrix.detach().cpu().numpy()\r\n\r\n return cost_matrix, None\r\n\r\n\r\ndef category_gate(cost_matrix, tracks, detections):\r\n \"\"\"\r\n :param tracks: list[STrack]\r\n :param detections: list[BaseTrack]\r\n :param metric:\r\n :return: cost_matrix np.ndarray\r\n \"\"\"\r\n if cost_matrix.size == 0:\r\n return cost_matrix\r\n\r\n det_categories = np.array([d.category for d in detections])\r\n trk_categories = np.array([t.category for t in tracks])\r\n\r\n cost_matrix = cost_matrix + np.abs(\r\n det_categories[None, :] - trk_categories[:, None])\r\n return cost_matrix\r\n\r\n\r\n"
},
{
"path": "core/motion/kalman_filter.py",
"content": "# vim: expandtab:ts=4:sw=4\r\nimport numpy as np\r\nimport scipy.linalg\r\n\r\n\r\n\"\"\"\r\nTable for the 0.95 quantile of the chi-square distribution with N degrees of\r\nfreedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv\r\nfunction and used as Mahalanobis gating threshold.\r\n\"\"\"\r\nchi2inv95 = {\r\n 1: 3.8415,\r\n 2: 5.9915,\r\n 3: 7.8147,\r\n 4: 9.4877,\r\n 5: 11.070,\r\n 6: 12.592,\r\n 7: 14.067,\r\n 8: 15.507,\r\n 9: 16.919}\r\n\r\n\r\nclass KalmanFilter(object):\r\n \"\"\"\r\n A simple Kalman filter for tracking bounding boxes in image space.\r\n\r\n The 8-dimensional state space\r\n\r\n x, y, a, h, vx, vy, va, vh\r\n\r\n contains the bounding box center position (x, y), aspect ratio a, height h,\r\n and their respective velocities.\r\n\r\n Object motion follows a constant velocity model. The bounding box location\r\n (x, y, a, h) is taken as direct observation of the state space (linear\r\n observation model).\r\n\r\n \"\"\"\r\n\r\n def __init__(self):\r\n ndim, dt = 4, 1.\r\n\r\n # Create Kalman filter model matrices.\r\n self._motion_mat = np.eye(2 * ndim, 2 * ndim)\r\n for i in range(ndim):\r\n self._motion_mat[i, ndim + i] = dt\r\n self._update_mat = np.eye(ndim, 2 * ndim)\r\n\r\n # Motion and observation uncertainty are chosen relative to the current\r\n # state estimate. These weights control the amount of uncertainty in\r\n # the model. This is a bit hacky.\r\n self._std_weight_position = 1. / 20\r\n self._std_weight_velocity = 1. / 160\r\n\r\n def initiate(self, measurement):\r\n \"\"\"Create track from unassociated measurement.\r\n\r\n Parameters\r\n ----------\r\n measurement : ndarray\r\n Bounding box coordinates (x, y, a, h) with center position (x, y),\r\n aspect ratio a, and height h.\r\n\r\n Returns\r\n -------\r\n (ndarray, ndarray)\r\n Returns the mean vector (8 dimensional) and covariance matrix (8x8\r\n dimensional) of the new track. Unobserved velocities are initialized\r\n to 0 mean.\r\n\r\n \"\"\"\r\n mean_pos = measurement\r\n mean_vel = np.zeros_like(mean_pos)\r\n mean = np.r_[mean_pos, mean_vel]\r\n\r\n std = [\r\n 2 * self._std_weight_position * measurement[3],\r\n 2 * self._std_weight_position * measurement[3],\r\n 1e-2,\r\n 2 * self._std_weight_position * measurement[3],\r\n 10 * self._std_weight_velocity * measurement[3],\r\n 10 * self._std_weight_velocity * measurement[3],\r\n 1e-5,\r\n 10 * self._std_weight_velocity * measurement[3]]\r\n covariance = np.diag(np.square(std))\r\n return mean, covariance\r\n\r\n def predict(self, mean, covariance):\r\n \"\"\"Run Kalman filter prediction step.\r\n\r\n Parameters\r\n ----------\r\n mean : ndarray\r\n The 8 dimensional mean vector of the object state at the previous\r\n time step.\r\n covariance : ndarray\r\n The 8x8 dimensional covariance matrix of the object state at the\r\n previous time step.\r\n\r\n Returns\r\n -------\r\n (ndarray, ndarray)\r\n Returns the mean vector and covariance matrix of the predicted\r\n state. Unobserved velocities are initialized to 0 mean.\r\n\r\n \"\"\"\r\n std_pos = [\r\n self._std_weight_position * mean[3],\r\n self._std_weight_position * mean[3],\r\n 1e-2,\r\n self._std_weight_position * mean[3]]\r\n std_vel = [\r\n self._std_weight_velocity * mean[3],\r\n self._std_weight_velocity * mean[3],\r\n 1e-5,\r\n self._std_weight_velocity * mean[3]]\r\n motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))\r\n\r\n mean = np.dot(mean, self._motion_mat.T)\r\n covariance = np.linalg.multi_dot((\r\n self._motion_mat, covariance, self._motion_mat.T)) + motion_cov\r\n\r\n return mean, covariance\r\n\r\n def project(self, mean, covariance):\r\n \"\"\"Project state distribution to measurement space.\r\n\r\n Parameters\r\n ----------\r\n mean : ndarray\r\n The state's mean vector (8 dimensional array).\r\n covariance : ndarray\r\n The state's covariance matrix (8x8 dimensional).\r\n\r\n Returns\r\n -------\r\n (ndarray, ndarray)\r\n Returns the projected mean and covariance matrix of the given state\r\n estimate.\r\n\r\n \"\"\"\r\n std = [\r\n self._std_weight_position * mean[3],\r\n self._std_weight_position * mean[3],\r\n 1e-1,\r\n self._std_weight_position * mean[3]]\r\n innovation_cov = np.diag(np.square(std))\r\n\r\n mean = np.dot(self._update_mat, mean)\r\n covariance = np.linalg.multi_dot((\r\n self._update_mat, covariance, self._update_mat.T))\r\n return mean, covariance + innovation_cov\r\n \r\n def multi_predict(self, mean, covariance):\r\n \"\"\"Run Kalman filter prediction step (Vectorized version).\r\n\r\n Parameters\r\n ----------\r\n mean : ndarray\r\n The Nx8 dimensional mean matrix of the object states at the previous\r\n time step.\r\n covariance : ndarray\r\n The Nx8x8 dimensional covariance matrics of the object states at the\r\n previous time step.\r\n\r\n Returns\r\n -------\r\n (ndarray, ndarray)\r\n Returns the mean vector and covariance matrix of the predicted\r\n state. Unobserved velocities are initialized to 0 mean.\r\n\r\n \"\"\"\r\n std_pos = [\r\n self._std_weight_position * mean[:, 3],\r\n self._std_weight_position * mean[:, 3],\r\n 1e-2 * np.ones_like(mean[:, 3]),\r\n self._std_weight_position * mean[:, 3]]\r\n std_vel = [\r\n self._std_weight_velocity * mean[:, 3],\r\n self._std_weight_velocity * mean[:, 3],\r\n 1e-5 * np.ones_like(mean[:, 3]),\r\n self._std_weight_velocity * mean[:, 3]]\r\n sqr = np.square(np.r_[std_pos, std_vel]).T\r\n \r\n motion_cov = []\r\n for i in range(len(mean)):\r\n motion_cov.append(np.diag(sqr[i]))\r\n motion_cov = np.asarray(motion_cov)\r\n \r\n mean = np.dot(mean, self._motion_mat.T)\r\n left = np.dot(self._motion_mat, covariance).transpose((1,0,2))\r\n covariance = np.dot(left, self._motion_mat.T) + motion_cov\r\n\r\n return mean, covariance\r\n\r\n def update(self, mean, covariance, measurement):\r\n \"\"\"Run Kalman filter correction step.\r\n\r\n Parameters\r\n ----------\r\n mean : ndarray\r\n The predicted state's mean vector (8 dimensional).\r\n covariance : ndarray\r\n The state's covariance matrix (8x8 dimensional).\r\n measurement : ndarray\r\n The 4 dimensional measurement vector (x, y, a, h), where (x, y)\r\n is the center position, a the aspect ratio, and h the height of the\r\n bounding box.\r\n\r\n Returns\r\n -------\r\n (ndarray, ndarray)\r\n Returns the measurement-corrected state distribution.\r\n\r\n \"\"\"\r\n projected_mean, projected_cov = self.project(mean, covariance)\r\n\r\n chol_factor, lower = scipy.linalg.cho_factor(\r\n projected_cov, lower=True, check_finite=False)\r\n kalman_gain = scipy.linalg.cho_solve(\r\n (chol_factor, lower), np.dot(covariance, self._update_mat.T).T,\r\n check_finite=False).T\r\n innovation = measurement - projected_mean\r\n\r\n new_mean = mean + np.dot(innovation, kalman_gain.T)\r\n new_covariance = covariance - np.linalg.multi_dot((\r\n kalman_gain, projected_cov, kalman_gain.T))\r\n return new_mean, new_covariance\r\n\r\n def gating_distance(self, mean, covariance, measurements,\r\n only_position=False, metric='maha'):\r\n \"\"\"Compute gating distance between state distribution and measurements.\r\n\r\n A suitable distance threshold can be obtained from `chi2inv95`. If\r\n `only_position` is False, the chi-square distribution has 4 degrees of\r\n freedom, otherwise 2.\r\n\r\n Parameters\r\n ----------\r\n mean : ndarray\r\n Mean vector over the state distribution (8 dimensional).\r\n covariance : ndarray\r\n Covariance of the state distribution (8x8 dimensional).\r\n measurements : ndarray\r\n An Nx4 dimensional matrix of N measurements, each in\r\n format (x, y, a, h) where (x, y) is the bounding box center\r\n position, a the aspect ratio, and h the height.\r\n only_position : Optional[bool]\r\n If True, distance computation is done with respect to the bounding\r\n box center position only.\r\n\r\n Returns\r\n -------\r\n ndarray\r\n Returns an array of length N, where the i-th element contains the\r\n squared Mahalanobis distance between (mean, covariance) and\r\n `measurements[i]`.\r\n\r\n \"\"\"\r\n mean, covariance = self.project(mean, covariance)\r\n if only_position:\r\n mean, covariance = mean[:2], covariance[:2, :2]\r\n measurements = measurements[:, :2]\r\n \r\n d = measurements - mean\r\n if metric == 'gaussian':\r\n return np.sum(d * d, axis=1)\r\n elif metric == 'maha':\r\n cholesky_factor = np.linalg.cholesky(covariance)\r\n z = scipy.linalg.solve_triangular(\r\n cholesky_factor, d.T, lower=True, check_finite=False,\r\n overwrite_b=True)\r\n squared_maha = np.sum(z * z, axis=0)\r\n return squared_maha\r\n else:\r\n raise ValueError('invalid distance metric')\r\n\r\n"
},
{
"path": "core/propagation/__init__.py",
"content": "###################################################################\n# File Name: __init__.py\n# Author: Zhongdao Wang\n# mail: wcd17@mails.tsinghua.edu.cn\n# Created Time: Mon Jan 18 15:57:34 2021\n###################################################################\n\nfrom __future__ import print_function\nfrom __future__ import division\nfrom __future__ import absolute_import\n\nfrom .propagate_box import propagate_box\nfrom .propagate_mask import propagate_mask\nfrom .propagate_pose import propagate_pose\n\ndef propagate(temp_feats, obs, img, model, format='box'):\n if format == 'box':\n return propagate_box(temp_feats, obs, img, model)\n elif format == 'mask':\n return propagate_box(temp_feats, obs, img, model)\n elif format == 'pose':\n return propagate_pose(temp_feats, obs, img, model)\n else:\n raise ValueError('Observation format not supported.')\n"
},
{
"path": "core/propagation/propagate_box.py",
"content": "###################################################################\n# File Name: propagate_box.py\n# Author: Zhongdao Wang\n# mail: wcd17@mails.tsinghua.edu.cn\n# Created Time: Mon Jan 18 16:01:46 2021\n###################################################################\n\nfrom __future__ import print_function\nfrom __future__ import division\nfrom __future__ import absolute_import\n\ndef propagate_box(temp_feats, box, img, model):\n pass\n"
},
{
"path": "core/propagation/propagate_mask.py",
"content": "###################################################################\n# File Name: propagate_box.py\n# Author: Zhongdao Wang\n# mail: wcd17@mails.tsinghua.edu.cn\n# Created Time: Mon Jan 18 16:01:46 2021\n###################################################################\n\nfrom __future__ import print_function\nfrom __future__ import division\nfrom __future__ import absolute_import\n\ndef propagate_mask(temp_feats, mask, img, model):\n pass\n"
},
{
"path": "core/propagation/propagate_pose.py",
"content": "###################################################################\n# File Name: propagate_box.py\n# Author: Zhongdao Wang\n# mail: wcd17@mails.tsinghua.edu.cn\n# Created Time: Mon Jan 18 16:01:46 2021\n###################################################################\n\nfrom __future__ import print_function\nfrom __future__ import division\nfrom __future__ import absolute_import\n\ndef propagate_pose(temp_feats, pose, img, model):\n pass\n"
},
{
"path": "data/jhmdb.py",
"content": "from __future__ import print_function, absolute_import\n\nimport os\nimport numpy as np\nimport math\nimport scipy.io as sio\n\nimport cv2\nimport torch\nfrom matplotlib import cm\n\nfrom utils import im_to_numpy, im_to_torch\n\ndef resize(img, owidth, oheight):\n img = im_to_numpy(img)\n img = cv2.resize( img, (owidth, oheight) )\n img = im_to_torch(img)\n return img\n\ndef load_image(img_path):\n # H x W x C => C x H x W\n img = cv2.imread(img_path)\n# print(img_path)\n img = img.astype(np.float32)\n img = img / 255.0\n img = img[:,:,::-1]\n img = img.copy()\n return im_to_torch(img)\n\ndef color_normalize(x, mean, std):\n if x.size(0) == 1:\n x = x.repeat(3, 1, 1)\n\n for t, m, s in zip(x, mean, std):\n t.sub_(m)\n t.div_(s)\n return x\n\nimport time\n\n\n\n\n######################################################################\ndef try_np_load(p):\n try:\n return np.load(p)\n except:\n return None\n\ndef make_lbl_set(lbls):\n print(lbls.shape)\n t00 = time.time()\n\n lbl_set = [np.zeros(3).astype(np.uint8)]\n count_lbls = [0] \n \n flat_lbls_0 = lbls[0].copy().reshape(-1, lbls.shape[-1]).astype(np.uint8)\n lbl_set = np.unique(flat_lbls_0, axis=0)\n\n # print(lbl_set)\n # if (lbl_set > 20).sum() > 0:\n # import pdb; pdb.set_trace()\n # count_lbls = [np.all(flat_lbls_0 == ll, axis=-1).sum() for ll in lbl_set]\n \n print('lbls', time.time() - t00)\n\n return lbl_set\n\n\ndef texturize(onehot):\n flat_onehot = onehot.reshape(-1, onehot.shape[-1])\n lbl_set = np.unique(flat_onehot, axis=0)\n\n count_lbls = [np.all(flat_onehot == ll, axis=-1).sum() for ll in lbl_set]\n object_id = np.argsort(count_lbls)[::-1][1]\n\n hidxs = []\n for h in range(onehot.shape[0]):\n # appears = any(np.all(onehot[h] == lbl_set[object_id], axis=-1))\n appears = np.any(onehot[h, :, 1:] == 1)\n if appears: \n hidxs.append(h)\n\n nstripes = min(10, len(hidxs))\n\n out = np.zeros((*onehot.shape[:2], nstripes+1))\n out[:, :, 0] = 1\n\n for i, h in enumerate(hidxs):\n cidx = int(i // (len(hidxs) / nstripes))\n w = np.any(onehot[h, :, 1:] == 1, axis=-1)\n out[h][w] = 0\n out[h][w, cidx+1] = 1\n # print(i, h, cidx)\n\n return out\n\n\n\nclass JhmdbSet(torch.utils.data.Dataset):\n def __init__(self, args, sigma=0.5):\n\n self.filelist = args.filelist\n self.imgSize = args.imgSize\n self.videoLen = args.videoLen\n self.mapScale = args.mapScale\n\n self.sigma = sigma\n\n f = open(self.filelist, 'r')\n self.jpgfiles = []\n self.lblfiles = []\n\n for line in f:\n rows = line.split()\n jpgfile = rows[1]\n lblfile = rows[0]\n\n self.jpgfiles.append(jpgfile)\n self.lblfiles.append(lblfile)\n\n f.close()\n \n def get_onehot_lbl(self, lbl_path):\n name = '/' + '/'.join(lbl_path.split('.')[:-1]) + '_onehot.npy'\n if os.path.exists(name):\n return np.load(name)\n else:\n return None\n \n\n def make_paths(self, folder_path, label_path):\n I = [ ll for ll in os.listdir(folder_path) if '.png' in ll ]\n\n frame_num = len(I) + self.videoLen\n I.sort(key=lambda x:int(x.split('.')[0]))\n\n I_out, L_out = [], []\n\n for i in range(frame_num):\n i = max(0, i - self.videoLen)\n img_path = \"%s/%s\" % (folder_path, I[i])\n\n I_out.append(img_path)\n\n return I_out\n\n\n def __getitem__(self, index):\n\n folder_path = self.jpgfiles[index]\n label_path = self.lblfiles[index]\n\n imgs = []\n imgs_orig = []\n lbls = []\n lbls_onehot = []\n patches = []\n target_imgs = []\n \n img_paths = self.make_paths(folder_path, label_path)\n frame_num = len(img_paths)\n\n mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]\n\n t000 = time.time()\n\n # frame_num = 30\n for i in range(frame_num):\n t00 = time.time()\n\n img_path = img_paths[i]\n img = load_image(img_path) # CxHxW\n\n # print('loaded', i, time.time() - t00)\n\n ht, wd = img.size(1), img.size(2)\n if self.imgSize > 0:\n newh, neww = ht, wd\n\n if ht <= wd:\n ratio = 1.0 #float(wd) / float(ht)\n # width, height\n img = resize(img, int(self.imgSize * ratio), self.imgSize)\n newh = self.imgSize\n neww = int(self.imgSize * ratio)\n else:\n ratio = 1.0 #float(ht) / float(wd)\n # width, height\n img = resize(img, self.imgSize, int(self.imgSize * ratio))\n newh = int(self.imgSize * ratio)\n neww = self.imgSize\n\n\n img_orig = img.clone()\n img = color_normalize(img, mean, std)\n\n imgs_orig.append(img_orig)\n imgs.append(img)\n\n rsz_h, rsz_w = math.ceil(img.size(1) / self.mapScale[0]), math.ceil(img.size(2) /self.mapScale[1])\n\n lbls_mat = sio.loadmat(label_path)\n\n lbls_coord = lbls_mat['pos_img']\n lbls_coord = lbls_coord - 1\n\n\n lbls_coord[0, :, :] = lbls_coord[0, :, :] * float(neww) / float(wd) / self.mapScale[0]\n lbls_coord[1, :, :] = lbls_coord[1, :, :] * float(newh) / float(ht) / self.mapScale[1]\n lblsize = (rsz_h, rsz_w)\n\n lbls = np.zeros((lbls_coord.shape[2], lblsize[0], lblsize[1], lbls_coord.shape[1]))\n\n for i in range(lbls_coord.shape[2]):\n lbls_coord_now = lbls_coord[:, :, i]\n scales = lbls_coord_now.max(1) - lbls_coord_now.min(1)\n scale = scales.max()\n scale = max(0.5, scale*0.015)\n for j in range(lbls_coord.shape[1]):\n if self.sigma > 0:\n draw_labelmap_np(lbls[i, :, :, j], lbls_coord_now[:, j], scale)\n else:\n tx = int(lbls_coord_now[0, j])\n ty = int(lbls_coord_now[1, j])\n if tx < lblsize[1] and ty < lblsize[0] and tx >=0 and ty >=0:\n lbls[i, ty, tx, j] = 1.0\n\n lbls_tensor = torch.zeros(frame_num, lblsize[0], lblsize[1], lbls_coord.shape[1])\n\n for i in range(frame_num):\n if i < self.videoLen:\n nowlbl = lbls[0]\n else:\n if(i - self.videoLen < len(lbls)):\n nowlbl = lbls[i - self.videoLen]\n lbls_tensor[i] = torch.from_numpy(nowlbl)\n \n lbls_tensor = torch.cat([(lbls_tensor.sum(-1) == 0)[..., None] *1.0, lbls_tensor], dim=-1)\n\n lblset = np.arange(lbls_tensor.shape[-1]-1)\n lblset = np.array([[0, 0, 0]] + [cm.Paired(i)[:3] for i in lblset]) * 255.0\n\n # Meta info\n meta = dict(folder_path=folder_path, img_paths=img_paths, lbl_paths=[])\n \n imgs = torch.stack(imgs)\n imgs_orig = torch.stack(imgs_orig)\n lbls_resize = lbls_tensor #np.stack(resizes)\n\n assert lbls_resize.shape[0] == len(meta['img_paths'])\n #print('vid', i, 'took', time.time() - t000)\n\n return imgs, imgs_orig, lbls_resize, lbls_tensor, lblset, meta\n\n def __len__(self):\n return len(self.jpgfiles)\n\n\ndef draw_labelmap_np(img, pt, sigma, type='Gaussian'):\n # Draw a 2D gaussian\n # Adopted from https://github.com/anewell/pose-hg-train/blob/master/src/pypose/draw.py\n\n # Check that any part of the gaussian is in-bounds\n ul = [int(pt[0] - 3 * sigma), int(pt[1] - 3 * sigma)]\n br = [int(pt[0] + 3 * sigma + 1), int(pt[1] + 3 * sigma + 1)]\n if (ul[0] >= img.shape[1] or ul[1] >= img.shape[0] or\n br[0] < 0 or br[1] < 0):\n # If not, just return the image as is\n return img\n\n # Generate gaussian\n size = 6 * sigma + 1\n x = np.arange(0, size, 1, float)\n y = x[:, np.newaxis]\n x0 = y0 = size // 2\n # The gaussian is not normalized, we want the center value to equal 1\n if type == 'Gaussian':\n g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))\n elif type == 'Cauchy':\n g = sigma / (((x - x0) ** 2 + (y - y0) ** 2 + sigma ** 2) ** 1.5)\n\n\n # Usable gaussian range\n g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]\n g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]\n # Image range\n img_x = max(0, ul[0]), min(br[0], img.shape[1])\n img_y = max(0, ul[1]), min(br[1], img.shape[0])\n\n img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]\n return img\n"
},
{
"path": "data/kinetics.py",
"content": "import torchvision.datasets.video_utils\n\nfrom torchvision.datasets.video_utils import VideoClips\nfrom torchvision.datasets.utils import list_dir\nfrom torchvision.datasets.folder import make_dataset\nfrom torchvision.datasets.vision import VisionDataset\n\nimport numpy as np\n\nclass Kinetics400(VisionDataset):\n \"\"\"\n `Kinetics-400 ![](\"assets/logo.png\"){kind=logo}
![](\"assets/demo.png\")
\n\n## Introduction\nYOLOX is an anchor-free version of YOLO, with a simpler design but better performance! It aims to bridge the gap between research and industrial communities.\nFor more details, please refer to our [report on Arxiv](https://arxiv.org/abs/2107.08430).\n\n![](\"assets/git_fig.png\")
\n\n## Updates!!\n* 【2021/07/20】 We have released our technical report on [Arxiv](https://arxiv.org/abs/2107.08430).\n\n## Comming soon\n- [ ] YOLOX-P6 and larger model.\n- [ ] Objects365 pretrain.\n- [ ] Transformer modules.\n- [ ] More features in need.\n\n## Benchmark\n\n#### Standard Models.\n|Model |size |mAPtest