├── LICENSE ├── README.md ├── assets ├── 100.jpg ├── 89.jpg ├── 99.jpg ├── demo.png ├── dog.jpg ├── git_fig.png ├── logo.png └── zidane.jpg ├── datasets └── README.md ├── demo ├── MegEngine │ ├── cpp │ │ ├── README.md │ │ ├── build.sh │ │ └── yolox.cpp │ └── python │ │ ├── README.md │ │ ├── build.py │ │ ├── coco_classes.py │ │ ├── convert_weights.py │ │ ├── demo.py │ │ ├── dump.py │ │ ├── models │ │ ├── __init__.py │ │ ├── darknet.py │ │ ├── network_blocks.py │ │ ├── yolo_fpn.py │ │ ├── yolo_head.py │ │ ├── yolo_pafpn.py │ │ └── yolox.py │ │ ├── process.py │ │ └── visualize.py ├── ONNXRuntime │ ├── README.md │ └── onnx_inference.py ├── OpenVINO │ ├── README.md │ ├── cpp │ │ ├── CMakeLists.txt │ │ ├── README.md │ │ └── yolox_openvino.cpp │ └── python │ │ ├── README.md │ │ └── openvino_inference.py ├── TensorRT │ ├── cpp │ │ ├── CMakeLists.txt │ │ ├── README.md │ │ ├── logging.h │ │ └── yolox.cpp │ └── python │ │ └── README.md └── ncnn │ ├── android │ ├── README.md │ ├── app │ │ ├── build.gradle │ │ └── src │ │ │ └── main │ │ │ ├── AndroidManifest.xml │ │ │ ├── assets │ │ │ └── yolox.param │ │ │ ├── java │ │ │ └── com │ │ │ │ └── megvii │ │ │ │ └── yoloXncnn │ │ │ │ ├── MainActivity.java │ │ │ │ ├── YOLOXncnn.java │ │ │ │ └── yoloXncnn.java │ │ │ ├── jni │ │ │ ├── CMakeLists.txt │ │ │ └── yoloXncnn_jni.cpp │ │ │ └── res │ │ │ ├── layout │ │ │ └── main.xml │ │ │ └── values │ │ │ └── strings.xml │ ├── build.gradle │ ├── gradle │ │ └── wrapper │ │ │ ├── gradle-wrapper.jar │ │ │ └── gradle-wrapper.properties │ ├── gradlew │ ├── gradlew.bat │ └── settings.gradle │ └── cpp │ ├── README.md │ └── yolox.cpp ├── docs └── train_custom_data.md ├── exps ├── default │ ├── nano.py │ ├── yolov3.py │ ├── yolox_l.py │ ├── yolox_m.py │ ├── yolox_s.py │ ├── yolox_tiny.py │ └── yolox_x.py └── example │ └── yolox_voc │ ├── __pycache__ │ └── yolox_yolo_s.cpython-38.pyc │ ├── yolox_voc_s.py │ └── yolox_yolo_s.py ├── requirements.txt ├── setup.cfg ├── setup.py ├── tools ├── demo.py ├── eval.py ├── export_onnx.py ├── train.py └── trt.py └── yolox ├── __init__.py ├── __pycache__ └── __init__.cpython-38.pyc ├── core ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── launch.cpython-38.pyc │ └── trainer.cpython-38.pyc ├── launch.py └── trainer.py ├── data ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── data_augment.cpython-38.pyc │ ├── data_prefetcher.cpython-38.pyc │ ├── dataloading.cpython-38.pyc │ └── samplers.cpython-38.pyc ├── data_augment.py ├── data_prefetcher.py ├── dataloading.py ├── datasets │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-38.pyc │ │ ├── coco.cpython-38.pyc │ │ ├── coco_classes.cpython-38.pyc │ │ ├── datasets_wrapper.cpython-38.pyc │ │ ├── mosaicdetection.cpython-38.pyc │ │ ├── voc.cpython-38.pyc │ │ ├── voc_classes.cpython-38.pyc │ │ ├── yolo.cpython-38.pyc │ │ └── yolo_classes.cpython-38.pyc │ ├── coco.py │ ├── coco_classes.py │ ├── datasets_wrapper.py │ ├── mosaicdetection.py │ ├── voc.py │ ├── voc_classes.py │ ├── yolo.py │ └── yolo_classes.py └── samplers.py ├── evaluators ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── coco_evaluator.cpython-38.pyc │ ├── voc_eval.cpython-38.pyc │ ├── voc_evaluator.cpython-38.pyc │ └── yolo_evaluator.cpython-38.pyc ├── coco_evaluator.py ├── voc_eval.py ├── voc_evaluator.py └── yolo_evaluator.py ├── exp ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── base_exp.cpython-38.pyc │ ├── build.cpython-38.pyc │ └── yolox_base.cpython-38.pyc ├── base_exp.py ├── build.py └── yolox_base.py ├── layers ├── __init__.py ├── csrc │ ├── cocoeval │ │ ├── cocoeval.cpp │ │ └── cocoeval.h │ └── vision.cpp └── fast_coco_eval_api.py ├── models ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── darknet.cpython-38.pyc │ ├── losses.cpython-38.pyc │ ├── network_blocks.cpython-38.pyc │ ├── yolo_fpn.cpython-38.pyc │ ├── yolo_head.cpython-38.pyc │ ├── yolo_pafpn.cpython-38.pyc │ └── yolox.cpython-38.pyc ├── darknet.py ├── losses.py ├── network_blocks.py ├── yolo_fpn.py ├── yolo_head.py ├── yolo_pafpn.py └── yolox.py └── utils ├── __init__.py ├── __pycache__ ├── __init__.cpython-38.pyc ├── allreduce_norm.cpython-38.pyc ├── boxes.cpython-38.pyc ├── checkpoint.cpython-38.pyc ├── demo_utils.cpython-38.pyc ├── dist.cpython-38.pyc ├── ema.cpython-38.pyc ├── logger.cpython-38.pyc ├── lr_scheduler.cpython-38.pyc ├── metric.cpython-38.pyc ├── model_utils.cpython-38.pyc ├── setup_env.cpython-38.pyc └── visualize.cpython-38.pyc ├── allreduce_norm.py ├── boxes.py ├── checkpoint.py ├── demo_utils.py ├── dist.py ├── ema.py ├── logger.py ├── lr_scheduler.py ├── metric.py ├── model_utils.py ├── setup_env.py └── visualize.py /README.md: -------------------------------------------------------------------------------- 1 | * 训练数据的格式请参考yolov5-face 2 | * 其它资料请参考yolox原始链接：https://github.com/Megvii-BaseDetection/YOLOX 3 | 4 | -------------------------------------------------------------------------------- /assets/100.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/100.jpg -------------------------------------------------------------------------------- /assets/89.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/89.jpg -------------------------------------------------------------------------------- /assets/99.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/99.jpg -------------------------------------------------------------------------------- /assets/demo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/demo.png -------------------------------------------------------------------------------- /assets/dog.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/dog.jpg -------------------------------------------------------------------------------- /assets/git_fig.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/git_fig.png -------------------------------------------------------------------------------- /assets/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/logo.png -------------------------------------------------------------------------------- /assets/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xialuxi/yolox-face-landmarks/71ff50bb35b6bf959cbe29b2ed93e17bf8a3a6fb/assets/zidane.jpg -------------------------------------------------------------------------------- /datasets/README.md: -------------------------------------------------------------------------------- 1 | # Prepare datasets 2 | 3 | If 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. 4 | ``` 5 | $YOLOX_DATADIR/ 6 | COCO/ 7 | ``` 8 | You can set the location for builtin datasets by 9 | ```shell 10 | export YOLOX_DATADIR=/path/to/your/datasets 11 | ``` 12 | If `YOLOX_DATADIR` is not set, the default value of dataset directory is `./datasets` relative to your current working directory. 13 | 14 | ## Expected dataset structure for [COCO detection](https://cocodataset.org/#download): 15 | 16 | ``` 17 | COCO/ 18 | annotations/ 19 | instances_{train,val}2017.json 20 | {train,val}2017/ 21 | # image files that are mentioned in the corresponding json 22 | ``` 23 | 24 | You can use the 2014 version of the dataset as well. 25 | -------------------------------------------------------------------------------- /demo/MegEngine/cpp/build.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -e 3 | 4 | if [ -z $CXX ];then 5 | echo "please export you c++ toolchain to CXX" 6 | echo "for example:" 7 | echo "build for host: export CXX=g++" 8 | echo "cross build for aarch64-android(always locate in NDK): export CXX=aarch64-linux-android21-clang++" 9 | echo "cross build for aarch64-linux: export CXX=aarch64-linux-gnu-g++" 10 | exit -1 11 | fi 12 | 13 | if [ -z $MGE_INSTALL_PATH ];then 14 | echo "please refsi ./README.md to init MGE_INSTALL_PATH env" 15 | exit -1 16 | fi 17 | 18 | if [ -z $OPENCV_INSTALL_INCLUDE_PATH ];then 19 | echo "please refs ./README.md to init OPENCV_INSTALL_INCLUDE_PATH env" 20 | exit -1 21 | fi 22 | 23 | if [ -z $OPENCV_INSTALL_LIB_PATH ];then 24 | echo "please refs ./README.md to init OPENCV_INSTALL_LIB_PATH env" 25 | exit -1 26 | fi 27 | 28 | INCLUDE_FLAG="-I$MGE_INSTALL_PATH/include -I$OPENCV_INSTALL_INCLUDE_PATH" 29 | LINK_FLAG="-L$MGE_INSTALL_PATH/lib/ -lmegengine -L$OPENCV_INSTALL_LIB_PATH -lopencv_core -lopencv_highgui -lopencv_imgproc -lopencv_imgcodecs" 30 | BUILD_FLAG="-static-libstdc++ -O3 -pie -fPIE -g" 31 | 32 | if [[ $CXX =~ "android" ]]; then 33 | LINK_FLAG="${LINK_FLAG} -llog -lz" 34 | fi 35 | 36 | echo "CXX: $CXX" 37 | echo "MGE_INSTALL_PATH: $MGE_INSTALL_PATH" 38 | echo "INCLUDE_FLAG: $INCLUDE_FLAG" 39 | echo "LINK_FLAG: $LINK_FLAG" 40 | echo "BUILD_FLAG: $BUILD_FLAG" 41 | 42 | echo "[" > compile_commands.json 43 | echo "{" >> compile_commands.json 44 | echo "\"directory\": \"$PWD\"," >> compile_commands.json 45 | echo "\"command\": \"$CXX yolox.cpp -o yolox ${INCLUDE_FLAG} ${LINK_FLAG}\"," >> compile_commands.json 46 | echo "\"file\": \"$PWD/yolox.cpp\"," >> compile_commands.json 47 | echo "}," >> compile_commands.json 48 | echo "]" >> compile_commands.json 49 | $CXX yolox.cpp -o yolox ${INCLUDE_FLAG} ${LINK_FLAG} ${BUILD_FLAG} 50 | 51 | echo "build success, output file: yolox" 52 | if [[ $CXX =~ "android" ]]; then 53 | echo "try command to run:" 54 | echo "adb push/scp $MGE_INSTALL_PATH/lib/libmegengine.so android_phone" 55 | echo "adb push/scp $OPENCV_INSTALL_LIB_PATH/*.so android_phone" 56 | echo "adb push/scp ./yolox yolox_s.mge android_phone" 57 | echo "adb push/scp ../../../assets/dog.jpg android_phone" 58 | echo "adb/ssh to android_phone, then run: LD_LIBRARY_PATH=. ./yolox yolox_s.mge dog.jpg cpu/multithread " 59 | else 60 | echo "try command to run: LD_LIBRARY_PATH=$MGE_INSTALL_PATH/lib/:$OPENCV_INSTALL_LIB_PATH ./yolox yolox_s.mge ../../../assets/dog.jpg cuda/cpu/multithread " 61 | fi 62 | -------------------------------------------------------------------------------- /demo/MegEngine/python/README.md: -------------------------------------------------------------------------------- 1 | # YOLOX-Python-MegEngine 2 | 3 | Python version of YOLOX object detection base on [MegEngine](https://github.com/MegEngine/MegEngine). 4 | 5 | ## Tutorial 6 | 7 | ### Step1: install requirements 8 | 9 | ``` 10 | python3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html 11 | ``` 12 | 13 | ### Step2: convert checkpoint weights from torch's path file 14 | 15 | ``` 16 | python3 convert_weights.py -w yolox_s.pth.tar -o yolox_s_mge.pkl 17 | ``` 18 | 19 | ### Step3: run demo 20 | 21 | This part is the same as torch's python demo, but no need to specify device. 22 | 23 | ``` 24 | python3 demo.py image -n yolox-s -c yolox_s_mge.pkl --path ../../../assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result 25 | ``` 26 | 27 | ### [Optional]Step4: dump model for cpp inference 28 | 29 | > **Note**: result model is dumped with `optimize_for_inference` and `enable_fuse_conv_bias_nonlinearity`. 30 | 31 | ``` 32 | python3 dump.py -n yolox-s -c yolox_s_mge.pkl --dump_path yolox_s.mge 33 | ``` 34 | -------------------------------------------------------------------------------- /demo/MegEngine/python/build.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | 4 | import megengine as mge 5 | import megengine.module as M 6 | from megengine import jit 7 | 8 | from models.yolo_fpn import YOLOFPN 9 | from models.yolo_head import YOLOXHead 10 | from models.yolo_pafpn import YOLOPAFPN 11 | from models.yolox import YOLOX 12 | 13 | 14 | def build_yolox(name="yolox-s"): 15 | num_classes = 80 16 | 17 | # value meaning: depth, width 18 | param_dict = { 19 | "yolox-nano": (0.33, 0.25), 20 | "yolox-tiny": (0.33, 0.375), 21 | "yolox-s": (0.33, 0.50), 22 | "yolox-m": (0.67, 0.75), 23 | "yolox-l": (1.0, 1.0), 24 | "yolox-x": (1.33, 1.25), 25 | } 26 | if name == "yolov3": 27 | depth = 1.0 28 | width = 1.0 29 | backbone = YOLOFPN() 30 | head = YOLOXHead(num_classes, width, in_channels=[128, 256, 512], act="lrelu") 31 | model = YOLOX(backbone, head) 32 | else: 33 | assert name in param_dict 34 | kwargs = {} 35 | depth, width = param_dict[name] 36 | if name == "yolox-nano": 37 | kwargs["depthwise"] = True 38 | in_channels = [256, 512, 1024] 39 | backbone = YOLOPAFPN(depth, width, in_channels=in_channels, **kwargs) 40 | head = YOLOXHead(num_classes, width, in_channels=in_channels, **kwargs) 41 | model = YOLOX(backbone, head) 42 | 43 | for m in model.modules(): 44 | if isinstance(m, M.BatchNorm2d): 45 | m.eps = 1e-3 46 | 47 | return model 48 | 49 | 50 | def build_and_load(weight_file, name="yolox-s"): 51 | model = build_yolox(name) 52 | model_weights = mge.load(weight_file) 53 | model.load_state_dict(model_weights, strict=False) 54 | return model 55 | -------------------------------------------------------------------------------- /demo/MegEngine/python/coco_classes.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | COCO_CLASSES = ( 6 | "person", 7 | "bicycle", 8 | "car", 9 | "motorcycle", 10 | "airplane", 11 | "bus", 12 | "train", 13 | "truck", 14 | "boat", 15 | "traffic light", 16 | "fire hydrant", 17 | "stop sign", 18 | "parking meter", 19 | "bench", 20 | "bird", 21 | "cat", 22 | "dog", 23 | "horse", 24 | "sheep", 25 | "cow", 26 | "elephant", 27 | "bear", 28 | "zebra", 29 | "giraffe", 30 | "backpack", 31 | "umbrella", 32 | "handbag", 33 | "tie", 34 | "suitcase", 35 | "frisbee", 36 | "skis", 37 | "snowboard", 38 | "sports ball", 39 | "kite", 40 | "baseball bat", 41 | "baseball glove", 42 | "skateboard", 43 | "surfboard", 44 | "tennis racket", 45 | "bottle", 46 | "wine glass", 47 | "cup", 48 | "fork", 49 | "knife", 50 | "spoon", 51 | "bowl", 52 | "banana", 53 | "apple", 54 | "sandwich", 55 | "orange", 56 | "broccoli", 57 | "carrot", 58 | "hot dog", 59 | "pizza", 60 | "donut", 61 | "cake", 62 | "chair", 63 | "couch", 64 | "potted plant", 65 | "bed", 66 | "dining table", 67 | "toilet", 68 | "tv", 69 | "laptop", 70 | "mouse", 71 | "remote", 72 | "keyboard", 73 | "cell phone", 74 | "microwave", 75 | "oven", 76 | "toaster", 77 | "sink", 78 | "refrigerator", 79 | "book", 80 | "clock", 81 | "vase", 82 | "scissors", 83 | "teddy bear", 84 | "hair drier", 85 | "toothbrush", 86 | ) 87 | -------------------------------------------------------------------------------- /demo/MegEngine/python/convert_weights.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | import argparse 4 | from collections import OrderedDict 5 | 6 | import megengine as mge 7 | import torch 8 | 9 | 10 | def make_parser(): 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument("-w", "--weights", type=str, help="path of weight file") 13 | parser.add_argument( 14 | "-o", 15 | "--output", 16 | default="weight_mge.pkl", 17 | type=str, 18 | help="path of weight file", 19 | ) 20 | return parser 21 | 22 | 23 | def numpy_weights(weight_file): 24 | torch_weights = torch.load(weight_file, map_location="cpu") 25 | if "model" in torch_weights: 26 | torch_weights = torch_weights["model"] 27 | new_dict = OrderedDict() 28 | for k, v in torch_weights.items(): 29 | new_dict[k] = v.cpu().numpy() 30 | return new_dict 31 | 32 | 33 | def map_weights(weight_file, output_file): 34 | torch_weights = numpy_weights(weight_file) 35 | 36 | new_dict = OrderedDict() 37 | for k, v in torch_weights.items(): 38 | if "num_batches_tracked" in k: 39 | print("drop: {}".format(k)) 40 | continue 41 | if k.endswith("bias"): 42 | print("bias key: {}".format(k)) 43 | v = v.reshape(1, -1, 1, 1) 44 | new_dict[k] = v 45 | elif "dconv" in k and "conv.weight" in k: 46 | print("depthwise conv key: {}".format(k)) 47 | cout, cin, k1, k2 = v.shape 48 | v = v.reshape(cout, 1, cin, k1, k2) 49 | new_dict[k] = v 50 | else: 51 | new_dict[k] = v 52 | 53 | mge.save(new_dict, output_file) 54 | print("save weights to {}".format(output_file)) 55 | 56 | 57 | def main(): 58 | parser = make_parser() 59 | args = parser.parse_args() 60 | map_weights(args.weights, args.output) 61 | 62 | 63 | if __name__ == "__main__": 64 | main() 65 | -------------------------------------------------------------------------------- /demo/MegEngine/python/dump.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | import argparse 6 | 7 | import megengine as mge 8 | import numpy as np 9 | from megengine import jit 10 | 11 | from build import build_and_load 12 | 13 | 14 | def make_parser(): 15 | parser = argparse.ArgumentParser("YOLOX Demo Dump") 16 | parser.add_argument("-n", "--name", type=str, default="yolox-s", help="model name") 17 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt for eval") 18 | parser.add_argument( 19 | "--dump_path", default="model.mge", help="path to save the dumped model" 20 | ) 21 | return parser 22 | 23 | 24 | def dump_static_graph(model, graph_name="model.mge"): 25 | model.eval() 26 | model.head.decode_in_inference = False 27 | 28 | data = mge.Tensor(np.random.random((1, 3, 640, 640))) 29 | 30 | @jit.trace(capture_as_const=True) 31 | def pred_func(data): 32 | outputs = model(data) 33 | return outputs 34 | 35 | pred_func(data) 36 | pred_func.dump( 37 | graph_name, 38 | arg_names=["data"], 39 | optimize_for_inference=True, 40 | enable_fuse_conv_bias_nonlinearity=True, 41 | ) 42 | 43 | 44 | def main(args): 45 | model = build_and_load(args.ckpt, name=args.name) 46 | dump_static_graph(model, args.dump_path) 47 | 48 | 49 | if __name__ == "__main__": 50 | args = make_parser().parse_args() 51 | main(args) 52 | -------------------------------------------------------------------------------- /demo/MegEngine/python/models/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | from .darknet import CSPDarknet, Darknet 6 | from .yolo_fpn import YOLOFPN 7 | from .yolo_head import YOLOXHead 8 | from .yolo_pafpn import YOLOPAFPN 9 | from .yolox import YOLOX 10 | -------------------------------------------------------------------------------- /demo/MegEngine/python/models/darknet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- encoding: utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import megengine.module as M 6 | 7 | from .network_blocks import BaseConv, CSPLayer, DWConv, Focus, ResLayer, SPPBottleneck 8 | 9 | 10 | class Darknet(M.Module): 11 | # number of blocks from dark2 to dark5. 12 | depth2blocks = {21: [1, 2, 2, 1], 53: [2, 8, 8, 4]} 13 | 14 | def __init__( 15 | self, depth, in_channels=3, stem_out_channels=32, out_features=("dark3", "dark4", "dark5"), 16 | ): 17 | """ 18 | Args: 19 | depth (int): depth of darknet used in model, usually use [21, 53] for this param. 20 | in_channels (int): number of input channels, for example, use 3 for RGB image. 21 | stem_out_channels (int): number of output chanels of darknet stem. 22 | It decides channels of darknet layer2 to layer5. 23 | out_features (Tuple[str]): desired output layer name. 24 | """ 25 | super().__init__() 26 | assert out_features, "please provide output features of Darknet" 27 | self.out_features = out_features 28 | self.stem = M.Sequential( 29 | BaseConv(in_channels, stem_out_channels, ksize=3, stride=1, act="lrelu"), 30 | *self.make_group_layer(stem_out_channels, num_blocks=1, stride=2), 31 | ) 32 | in_channels = stem_out_channels * 2 # 64 33 | 34 | num_blocks = Darknet.depth2blocks[depth] 35 | # create darknet with `stem_out_channels` and `num_blocks` layers. 36 | # to make model structure more clear, we don't use `for` statement in python. 37 | self.dark2 = M.Sequential(*self.make_group_layer(in_channels, num_blocks[0], stride=2)) 38 | in_channels *= 2 # 128 39 | self.dark3 = M.Sequential(*self.make_group_layer(in_channels, num_blocks[1], stride=2)) 40 | in_channels *= 2 # 256 41 | self.dark4 = M.Sequential(*self.make_group_layer(in_channels, num_blocks[2], stride=2)) 42 | in_channels *= 2 # 512 43 | 44 | self.dark5 = M.Sequential( 45 | *self.make_group_layer(in_channels, num_blocks[3], stride=2), 46 | *self.make_spp_block([in_channels, in_channels * 2], in_channels * 2), 47 | ) 48 | 49 | def make_group_layer(self, in_channels: int, num_blocks: int, stride: int = 1): 50 | "starts with conv layer then has `num_blocks` `ResLayer`" 51 | return [ 52 | BaseConv(in_channels, in_channels * 2, ksize=3, stride=stride, act="lrelu"), 53 | *[(ResLayer(in_channels * 2)) for _ in range(num_blocks)] 54 | ] 55 | 56 | def make_spp_block(self, filters_list, in_filters): 57 | m = M.Sequential( 58 | *[ 59 | BaseConv(in_filters, filters_list[0], 1, stride=1, act="lrelu"), 60 | BaseConv(filters_list[0], filters_list[1], 3, stride=1, act="lrelu"), 61 | SPPBottleneck( 62 | in_channels=filters_list[1], 63 | out_channels=filters_list[0], 64 | activation="lrelu" 65 | ), 66 | BaseConv(filters_list[0], filters_list[1], 3, stride=1, act="lrelu"), 67 | BaseConv(filters_list[1], filters_list[0], 1, stride=1, act="lrelu"), 68 | ] 69 | ) 70 | return m 71 | 72 | def forward(self, x): 73 | outputs = {} 74 | x = self.stem(x) 75 | outputs["stem"] = x 76 | x = self.dark2(x) 77 | outputs["dark2"] = x 78 | x = self.dark3(x) 79 | outputs["dark3"] = x 80 | x = self.dark4(x) 81 | outputs["dark4"] = x 82 | x = self.dark5(x) 83 | outputs["dark5"] = x 84 | return {k: v for k, v in outputs.items() if k in self.out_features} 85 | 86 | 87 | class CSPDarknet(M.Module): 88 | 89 | def __init__( 90 | self, dep_mul, wid_mul, 91 | out_features=("dark3", "dark4", "dark5"), 92 | depthwise=False, act="silu", 93 | ): 94 | super().__init__() 95 | assert out_features, "please provide output features of Darknet" 96 | self.out_features = out_features 97 | Conv = DWConv if depthwise else BaseConv 98 | 99 | base_channels = int(wid_mul * 64) # 64 100 | base_depth = max(round(dep_mul * 3), 1) # 3 101 | 102 | # stem 103 | self.stem = Focus(3, base_channels, ksize=3, act=act) 104 | 105 | # dark2 106 | self.dark2 = M.Sequential( 107 | Conv(base_channels, base_channels * 2, 3, 2, act=act), 108 | CSPLayer( 109 | base_channels * 2, base_channels * 2, 110 | n=base_depth, depthwise=depthwise, act=act 111 | ), 112 | ) 113 | 114 | # dark3 115 | self.dark3 = M.Sequential( 116 | Conv(base_channels * 2, base_channels * 4, 3, 2, act=act), 117 | CSPLayer( 118 | base_channels * 4, base_channels * 4, 119 | n=base_depth * 3, depthwise=depthwise, act=act, 120 | ), 121 | ) 122 | 123 | # dark4 124 | self.dark4 = M.Sequential( 125 | Conv(base_channels * 4, base_channels * 8, 3, 2, act=act), 126 | CSPLayer( 127 | base_channels * 8, base_channels * 8, 128 | n=base_depth * 3, depthwise=depthwise, act=act, 129 | ), 130 | ) 131 | 132 | # dark5 133 | self.dark5 = M.Sequential( 134 | Conv(base_channels * 8, base_channels * 16, 3, 2, act=act), 135 | SPPBottleneck(base_channels * 16, base_channels * 16, activation=act), 136 | CSPLayer( 137 | base_channels * 16, base_channels * 16, n=base_depth, 138 | shortcut=False, depthwise=depthwise, act=act, 139 | ), 140 | ) 141 | 142 | def forward(self, x): 143 | outputs = {} 144 | x = self.stem(x) 145 | outputs["stem"] = x 146 | x = self.dark2(x) 147 | outputs["dark2"] = x 148 | x = self.dark3(x) 149 | outputs["dark3"] = x 150 | x = self.dark4(x) 151 | outputs["dark4"] = x 152 | x = self.dark5(x) 153 | outputs["dark5"] = x 154 | return {k: v for k, v in outputs.items() if k in self.out_features} 155 | -------------------------------------------------------------------------------- /demo/MegEngine/python/models/yolo_fpn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- encoding: utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import megengine.functional as F 6 | import megengine.module as M 7 | 8 | from .darknet import Darknet 9 | from .network_blocks import BaseConv, UpSample 10 | 11 | 12 | class YOLOFPN(M.Module): 13 | """ 14 | YOLOFPN module. Darknet 53 is the default backbone of this model. 15 | """ 16 | 17 | def __init__( 18 | self, depth=53, in_features=["dark3", "dark4", "dark5"], 19 | ): 20 | super().__init__() 21 | 22 | self.backbone = Darknet(depth) 23 | self.in_features = in_features 24 | 25 | # out 1 26 | self.out1_cbl = self._make_cbl(512, 256, 1) 27 | self.out1 = self._make_embedding([256, 512], 512 + 256) 28 | 29 | # out 2 30 | self.out2_cbl = self._make_cbl(256, 128, 1) 31 | self.out2 = self._make_embedding([128, 256], 256 + 128) 32 | 33 | # upsample 34 | self.upsample = UpSample(scale_factor=2, mode="bilinear") 35 | 36 | def _make_cbl(self, _in, _out, ks): 37 | return BaseConv(_in, _out, ks, stride=1, act="lrelu") 38 | 39 | def _make_embedding(self, filters_list, in_filters): 40 | m = M.Sequential( 41 | *[ 42 | self._make_cbl(in_filters, filters_list[0], 1), 43 | self._make_cbl(filters_list[0], filters_list[1], 3), 44 | 45 | self._make_cbl(filters_list[1], filters_list[0], 1), 46 | 47 | self._make_cbl(filters_list[0], filters_list[1], 3), 48 | self._make_cbl(filters_list[1], filters_list[0], 1), 49 | ] 50 | ) 51 | return m 52 | 53 | def forward(self, inputs): 54 | """ 55 | Args: 56 | inputs (Tensor): input image. 57 | 58 | Returns: 59 | Tuple[Tensor]: FPN output features.. 60 | """ 61 | # backbone 62 | out_features = self.backbone(inputs) 63 | x2, x1, x0 = [out_features[f] for f in self.in_features] 64 | 65 | # yolo branch 1 66 | x1_in = self.out1_cbl(x0) 67 | x1_in = self.upsample(x1_in) 68 | x1_in = F.concat([x1_in, x1], 1) 69 | out_dark4 = self.out1(x1_in) 70 | 71 | # yolo branch 2 72 | x2_in = self.out2_cbl(out_dark4) 73 | x2_in = self.upsample(x2_in) 74 | x2_in = F.concat([x2_in, x2], 1) 75 | out_dark3 = self.out2(x2_in) 76 | 77 | outputs = (out_dark3, out_dark4, x0) 78 | return outputs 79 | -------------------------------------------------------------------------------- /demo/MegEngine/python/models/yolo_pafpn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- encoding: utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import megengine.module as M 6 | import megengine.functional as F 7 | 8 | from .darknet import CSPDarknet 9 | from .network_blocks import BaseConv, CSPLayer, DWConv, UpSample 10 | 11 | 12 | class YOLOPAFPN(M.Module): 13 | """ 14 | YOLOv3 model. Darknet 53 is the default backbone of this model. 15 | """ 16 | 17 | def __init__( 18 | self, depth=1.0, width=1.0, in_features=("dark3", "dark4", "dark5"), 19 | in_channels=[256, 512, 1024], depthwise=False, act="silu", 20 | ): 21 | super().__init__() 22 | self.backbone = CSPDarknet(depth, width, depthwise=depthwise, act=act) 23 | self.in_features = in_features 24 | self.in_channels = in_channels 25 | Conv = DWConv if depthwise else BaseConv 26 | 27 | self.upsample = UpSample(scale_factor=2, mode="bilinear") 28 | self.lateral_conv0 = BaseConv( 29 | int(in_channels[2] * width), int(in_channels[1] * width), 1, 1, act=act 30 | ) 31 | self.C3_p4 = CSPLayer( 32 | int(2 * in_channels[1] * width), 33 | int(in_channels[1] * width), 34 | round(3 * depth), 35 | False, 36 | depthwise=depthwise, 37 | act=act, 38 | ) # cat 39 | 40 | self.reduce_conv1 = BaseConv( 41 | int(in_channels[1] * width), int(in_channels[0] * width), 1, 1, act=act 42 | ) 43 | self.C3_p3 = CSPLayer( 44 | int(2 * in_channels[0] * width), 45 | int(in_channels[0] * width), 46 | round(3 * depth), 47 | False, 48 | depthwise=depthwise, 49 | act=act, 50 | ) 51 | 52 | # bottom-up conv 53 | self.bu_conv2 = Conv( 54 | int(in_channels[0] * width), int(in_channels[0] * width), 3, 2, act=act 55 | ) 56 | self.C3_n3 = CSPLayer( 57 | int(2 * in_channels[0] * width), 58 | int(in_channels[1] * width), 59 | round(3 * depth), 60 | False, 61 | depthwise=depthwise, 62 | act=act, 63 | ) 64 | 65 | # bottom-up conv 66 | self.bu_conv1 = Conv( 67 | int(in_channels[1] * width), int(in_channels[1] * width), 3, 2, act=act 68 | ) 69 | self.C3_n4 = CSPLayer( 70 | int(2 * in_channels[1] * width), 71 | int(in_channels[2] * width), 72 | round(3 * depth), 73 | False, 74 | depthwise=depthwise, 75 | act=act, 76 | ) 77 | 78 | def forward(self, input): 79 | """ 80 | Args: 81 | inputs: input images. 82 | 83 | Returns: 84 | Tuple[Tensor]: FPN feature. 85 | """ 86 | 87 | # backbone 88 | out_features = self.backbone(input) 89 | features = [out_features[f] for f in self.in_features] 90 | [x2, x1, x0] = features 91 | 92 | fpn_out0 = self.lateral_conv0(x0) # 1024->512/32 93 | f_out0 = self.upsample(fpn_out0) # 512/16 94 | f_out0 = F.concat([f_out0, x1], 1) # 512->1024/16 95 | f_out0 = self.C3_p4(f_out0) # 1024->512/16 96 | 97 | fpn_out1 = self.reduce_conv1(f_out0) # 512->256/16 98 | f_out1 = self.upsample(fpn_out1) # 256/8 99 | f_out1 = F.concat([f_out1, x2], 1) # 256->512/8 100 | pan_out2 = self.C3_p3(f_out1) # 512->256/8 101 | 102 | p_out1 = self.bu_conv2(pan_out2) # 256->256/16 103 | p_out1 = F.concat([p_out1, fpn_out1], 1) # 256->512/16 104 | pan_out1 = self.C3_n3(p_out1) # 512->512/16 105 | 106 | p_out0 = self.bu_conv1(pan_out1) # 512->512/32 107 | p_out0 = F.concat([p_out0, fpn_out0], 1) # 512->1024/32 108 | pan_out0 = self.C3_n4(p_out0) # 1024->1024/32 109 | 110 | outputs = (pan_out2, pan_out1, pan_out0) 111 | return outputs 112 | -------------------------------------------------------------------------------- /demo/MegEngine/python/models/yolox.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- encoding: utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import megengine.module as M 6 | 7 | from .yolo_head import YOLOXHead 8 | from .yolo_pafpn import YOLOPAFPN 9 | 10 | 11 | class YOLOX(M.Module): 12 | """ 13 | YOLOX model module. The module list is defined by create_yolov3_modules function. 14 | The network returns loss values from three YOLO layers during training 15 | and detection results during test. 16 | """ 17 | 18 | def __init__(self, backbone=None, head=None): 19 | super().__init__() 20 | if backbone is None: 21 | backbone = YOLOPAFPN() 22 | if head is None: 23 | head = YOLOXHead(80) 24 | 25 | self.backbone = backbone 26 | self.head = head 27 | 28 | def forward(self, x): 29 | # fpn output content features of [dark3, dark4, dark5] 30 | fpn_outs = self.backbone(x) 31 | assert not self.training 32 | outputs = self.head(fpn_outs) 33 | 34 | return outputs 35 | -------------------------------------------------------------------------------- /demo/MegEngine/python/process.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import cv2 6 | import megengine.functional as F 7 | import numpy as np 8 | 9 | __all__ = [ 10 | "preprocess", 11 | "postprocess", 12 | ] 13 | 14 | 15 | def preprocess(image, input_size, mean, std, swap=(2, 0, 1)): 16 | if len(image.shape) == 3: 17 | padded_img = np.ones((input_size[0], input_size[1], 3)) * 114.0 18 | else: 19 | padded_img = np.ones(input_size) * 114.0 20 | img = np.array(image) 21 | r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1]) 22 | resized_img = cv2.resize( 23 | img, 24 | (int(img.shape[1] * r), int(img.shape[0] * r)), 25 | interpolation=cv2.INTER_LINEAR, 26 | ).astype(np.float32) 27 | padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img 28 | image = padded_img 29 | 30 | image = image.astype(np.float32) 31 | image = image[:, :, ::-1] 32 | image /= 255.0 33 | if mean is not None: 34 | image -= mean 35 | if std is not None: 36 | image /= std 37 | image = image.transpose(swap) 38 | image = np.ascontiguousarray(image, dtype=np.float32) 39 | return image, r 40 | 41 | 42 | def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45): 43 | box_corner = F.zeros_like(prediction) 44 | box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2 45 | box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2 46 | box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2 47 | box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2 48 | prediction[:, :, :4] = box_corner[:, :, :4] 49 | 50 | output = [None for _ in range(len(prediction))] 51 | for i, image_pred in enumerate(prediction): 52 | 53 | # If none are remaining => process next image 54 | if not image_pred.shape[0]: 55 | continue 56 | # Get score and class with highest confidence 57 | class_conf = F.max(image_pred[:, 5 : 5 + num_classes], 1, keepdims=True) 58 | class_pred = F.argmax(image_pred[:, 5 : 5 + num_classes], 1, keepdims=True) 59 | 60 | class_conf_squeeze = F.squeeze(class_conf) 61 | conf_mask = image_pred[:, 4] * class_conf_squeeze >= conf_thre 62 | detections = F.concat((image_pred[:, :5], class_conf, class_pred), 1) 63 | detections = detections[conf_mask] 64 | if not detections.shape[0]: 65 | continue 66 | 67 | nms_out_index = F.vision.nms( 68 | detections[:, :4], detections[:, 4] * detections[:, 5], nms_thre, 69 | ) 70 | detections = detections[nms_out_index] 71 | if output[i] is None: 72 | output[i] = detections 73 | else: 74 | output[i] = F.concat((output[i], detections)) 75 | 76 | return output 77 | -------------------------------------------------------------------------------- /demo/MegEngine/python/visualize.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding:utf-8 -*- 3 | # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. 4 | 5 | import cv2 6 | import numpy as np 7 | 8 | __all__ = ["vis"] 9 | 10 | 11 | def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None): 12 | 13 | for i in range(len(boxes)): 14 | box = boxes[i] 15 | cls_id = int(cls_ids[i]) 16 | score = scores[i] 17 | if score < conf: 18 | continue 19 | x0 = int(box[0]) 20 | y0 = int(box[1]) 21 | x1 = int(box[2]) 22 | y1 = int(box[3]) 23 | 24 | color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist() 25 | text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100) 26 | txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255) 27 | font = cv2.FONT_HERSHEY_SIMPLEX 28 | 29 | txt_size = cv2.getTextSize(text, font, 0.4, 1)[0] 30 | cv2.rectangle(img, (x0, y0), (x1, y1), color, 2) 31 | 32 | txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist() 33 | cv2.rectangle( 34 | img, 35 | (x0, y0 + 1), 36 | (x0 + txt_size[0] + 1, y0 + int(1.5*txt_size[1])), 37 | txt_bk_color, 38 | -1 39 | ) 40 | cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1) 41 | 42 | return img 43 | 44 | 45 | _COLORS = np.array( 46 | [ 47 | 0.000, 0.447, 0.741, 48 | 0.850, 0.325, 0.098, 49 | 0.929, 0.694, 0.125, 50 | 0.494, 0.184, 0.556, 51 | 0.466, 0.674, 0.188, 52 | 0.301, 0.745, 0.933, 53 | 0.635, 0.078, 0.184, 54 | 0.300, 0.300, 0.300, 55 | 0.600, 0.600, 0.600, 56 | 1.000, 0.000, 0.000, 57 | 1.000, 0.500, 0.000, 58 | 0.749, 0.749, 0.000, 59 | 0.000, 1.000, 0.000, 60 | 0.000, 0.000, 1.000, 61 | 0.667, 0.000, 1.000, 62 | 0.333, 0.333, 0.000, 63 | 0.333, 0.667, 0.000, 64 | 0.333, 1.000, 0.000, 65 | 0.667, 0.333, 0.000, 66 | 0.667, 0.667, 0.000, 67 | 0.667, 1.000, 0.000, 68 | 1.000, 0.333, 0.000, 69 | 1.000, 0.667, 0.000, 70 | 1.000, 1.000, 0.000, 71 | 0.000, 0.333, 0.500, 72 | 0.000, 0.667, 0.500, 73 | 0.000, 1.000, 0.500, 74 | 0.333, 0.000, 0.500, 75 | 0.333, 0.333, 0.500, 76 | 0.333, 0.667, 0.500, 77 | 0.333, 1.000, 0.500, 78 | 0.667, 0.000, 0.500, 79 | 0.667, 0.333, 0.500, 80 | 0.667, 0.667, 0.500, 81 | 0.667, 1.000, 0.500, 82 | 1.000, 0.000, 0.500, 83 | 1.000, 0.333, 0.500, 84 | 1.000, 0.667, 0.500, 85 | 1.000, 1.000, 0.500, 86 | 0.000, 0.333, 1.000, 87 | 0.000, 0.667, 1.000, 88 | 0.000, 1.000, 1.000, 89 | 0.333, 0.000, 1.000, 90 | 0.333, 0.333, 1.000, 91 | 0.333, 0.667, 1.000, 92 | 0.333, 1.000, 1.000, 93 | 0.667, 0.000, 1.000, 94 | 0.667, 0.333, 1.000, 95 | 0.667, 0.667, 1.000, 96 | 0.667, 1.000, 1.000, 97 | 1.000, 0.000, 1.000, 98 | 1.000, 0.333, 1.000, 99 | 1.000, 0.667, 1.000, 100 | 0.333, 0.000, 0.000, 101 | 0.500, 0.000, 0.000, 102 | 0.667, 0.000, 0.000, 103 | 0.833, 0.000, 0.000, 104 | 1.000, 0.000, 0.000, 105 | 0.000, 0.167, 0.000, 106 | 0.000, 0.333, 0.000, 107 | 0.000, 0.500, 0.000, 108 | 0.000, 0.667, 0.000, 109 | 0.000, 0.833, 0.000, 110 | 0.000, 1.000, 0.000, 111 | 0.000, 0.000, 0.167, 112 | 0.000, 0.000, 0.333, 113 | 0.000, 0.000, 0.500, 114 | 0.000, 0.000, 0.667, 115 | 0.000, 0.000, 0.833, 116 | 0.000, 0.000, 1.000, 117 | 0.000, 0.000, 0.000, 118 | 0.143, 0.143, 0.143, 119 | 0.286, 0.286, 0.286, 120 | 0.429, 0.429, 0.429, 121 | 0.571, 0.571, 0.571, 122 | 0.714, 0.714, 0.714, 123 | 0.857, 0.857, 0.857, 124 | 0.000, 0.447, 0.741, 125 | 0.314, 0.717, 0.741, 126 | 0.50, 0.5, 0 127 | ] 128 | ).astype(np.float32).reshape(-1, 3) 129 | -------------------------------------------------------------------------------- /demo/ONNXRuntime/README.md: -------------------------------------------------------------------------------- 1 | ## YOLOX-ONNXRuntime in Python 2 | 3 | This doc introduces how to convert your pytorch model into onnx, and how to run an onnxruntime demo to verify your convertion. 4 | 5 | ### Download ONNX models. 6 | | Model | Parameters | GFLOPs | Test Size | mAP | Weights | 7 | |:------| :----: | :----: | :---: | :---: | :---: | 8 | | 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) | 9 | | 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) | 10 | | 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) | 11 | | 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) | 12 | | 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) | 13 | | 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) | 14 | | 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) | 15 | 16 | 17 | ### Convert Your Model to ONNX 18 | 19 | First, you should move to by: 20 | ```shell 21 | cd 22 | ``` 23 | Then, you can: 24 | 25 | 1. Convert a standard YOLOX model by -n: 26 | ```shell 27 | python3 tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth.tar 28 | ``` 29 | Notes: 30 | * -n: specify a model name. The model name must be one of the [yolox-s,m,l,x and yolox-nane, yolox-tiny, yolov3] 31 | * -c: the model you have trained 32 | * -o: opset version, default 11. **However, if you will further convert your onnx model to [OpenVINO](../OpenVINO/), please specify the opset version to 10.** 33 | * --no-onnxsim: disable onnxsim 34 | * To customize an input shape for onnx model, modify the following code in tools/export.py: 35 | 36 | ```python 37 | dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1]) 38 | ``` 39 | 40 | 2. Convert a standard YOLOX model by -f. When using -f, the above command is equivalent to: 41 | 42 | ```shell 43 | python3 tools/export_onnx.py --output-name yolox_s.onnx -f exps/default/yolox_s.py -c yolox_s.pth.tar 44 | ``` 45 | 46 | 3. To convert your customized model, please use -f: 47 | 48 | ```shell 49 | python3 tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth.tar 50 | ``` 51 | 52 | ### ONNXRuntime Demo 53 | 54 | Step1. 55 | ```shell 56 | cd /demo/ONNXRuntime 57 | ``` 58 | 59 | Step2. 60 | ```shell 61 | python3 onnx_inference.py -m -i -o -s 0.3 --input_shape 640,640 62 | ``` 63 | Notes: 64 | * -m: your converted onnx model 65 | * -i: input_image 66 | * -s: score threshold for visualization. 67 | * --input_shape: should be consistent with the shape you used for onnx convertion. 68 | -------------------------------------------------------------------------------- /demo/ONNXRuntime/onnx_inference.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | # Copyright (c) Megvii, Inc. and its affiliates. 4 | 5 | import argparse 6 | import os 7 | 8 | import cv2 9 | import numpy as np 10 | 11 | import onnxruntime 12 | 13 | from yolox.data.data_augment import preproc as preprocess 14 | from yolox.data.datasets import COCO_CLASSES 15 | from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis 16 | 17 | 18 | def make_parser(): 19 | parser = argparse.ArgumentParser("onnxruntime inference sample") 20 | parser.add_argument( 21 | "-m", 22 | "--model", 23 | type=str, 24 | default="yolox.onnx", 25 | help="Input your onnx model.", 26 | ) 27 | parser.add_argument( 28 | "-i", 29 | "--image_path", 30 | type=str, 31 | default='test_image.png', 32 | help="Path to your input image.", 33 | ) 34 | parser.add_argument( 35 | "-o", 36 | "--output_dir", 37 | type=str, 38 | default='demo_output', 39 | help="Path to your output directory.", 40 | ) 41 | parser.add_argument( 42 | "-s", 43 | "--score_thr", 44 | type=float, 45 | default=0.3, 46 | help="Score threshould to filter the result.", 47 | ) 48 | parser.add_argument( 49 | "--input_shape", 50 | type=str, 51 | default="640,640", 52 | help="Specify an input shape for inference.", 53 | ) 54 | parser.add_argument( 55 | "--with_p6", 56 | action="store_true", 57 | help="Whether your model uses p6 in FPN/PAN.", 58 | ) 59 | return parser 60 | 61 | 62 | if __name__ == '__main__': 63 | args = make_parser().parse_args() 64 | 65 | input_shape = tuple(map(int, args.input_shape.split(','))) 66 | origin_img = cv2.imread(args.image_path) 67 | mean = (0.485, 0.456, 0.406) 68 | std = (0.229, 0.224, 0.225) 69 | img, ratio = preprocess(origin_img, input_shape, mean, std) 70 | 71 | session = onnxruntime.InferenceSession(args.model) 72 | 73 | ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} 74 | output = session.run(None, ort_inputs) 75 | predictions = demo_postprocess(output[0], input_shape, p6=args.with_p6)[0] 76 | 77 | boxes = predictions[:, :4] 78 | scores = predictions[:, 4:5] * predictions[:, 5:] 79 | 80 | boxes_xyxy = np.ones_like(boxes) 81 | boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2. 82 | boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2. 83 | boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2. 84 | boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2. 85 | boxes_xyxy /= ratio 86 | dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1) 87 | if dets is not None: 88 | final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] 89 | origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds, 90 | conf=args.score_thr, class_names=COCO_CLASSES) 91 | 92 | mkdir(args.output_dir) 93 | output_path = os.path.join(args.output_dir, args.image_path.split("/")[-1]) 94 | cv2.imwrite(output_path, origin_img) 95 | -------------------------------------------------------------------------------- /demo/OpenVINO/README.md: -------------------------------------------------------------------------------- 1 | ## YOLOX for OpenVINO 2 | 3 | * [C++ Demo](./cpp) 4 | * [Python Demo](./python) -------------------------------------------------------------------------------- /demo/OpenVINO/cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.4.1) 2 | set(CMAKE_CXX_STANDARD 14) 3 | 4 | project(yolox_openvino_demo) 5 | 6 | find_package(OpenCV REQUIRED) 7 | find_package(InferenceEngine REQUIRED) 8 | find_package(ngraph REQUIRED) 9 | 10 | include_directories( 11 | ${OpenCV_INCLUDE_DIRS} 12 | ${CMAKE_CURRENT_SOURCE_DIR} 13 | ${CMAKE_CURRENT_BINARY_DIR} 14 | ) 15 | 16 | add_executable(yolox_openvino yolox_openvino.cpp) 17 | 18 | target_link_libraries( 19 | yolox_openvino 20 | ${InferenceEngine_LIBRARIES} 21 | ${NGRAPH_LIBRARIES} 22 | ${OpenCV_LIBS} 23 | ) -------------------------------------------------------------------------------- /demo/OpenVINO/cpp/README.md: -------------------------------------------------------------------------------- 1 | # YOLOX-OpenVINO in C++ 2 | 3 | This toturial includes a C++ demo for OpenVINO, as well as some converted models. 4 | 5 | ### Download OpenVINO models. 6 | | Model | Parameters | GFLOPs | Test Size | mAP | Weights | 7 | |:------| :----: | :----: | :---: | :---: | :---: | 8 | | [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) | 9 | | [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) | 10 | | [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) | 11 | | [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) | 12 | | [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) | 13 | | [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) | 14 | | [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) | 15 | 16 | ## Install OpenVINO Toolkit 17 | 18 | Please visit [Openvino Homepage](https://docs.openvinotoolkit.org/latest/get_started_guides.html) for more details. 19 | 20 | ## Set up the Environment 21 | 22 | ### For Linux 23 | 24 | **Option1. Set up the environment tempororally. You need to run this command everytime you start a new shell window.** 25 | 26 | ```shell 27 | source /opt/intel/openvino_2021/bin/setupvars.sh 28 | ``` 29 | 30 | **Option2. Set up the environment permenantly.** 31 | 32 | *Step1.* For Linux: 33 | ```shell 34 | vim ~/.bashrc 35 | ``` 36 | 37 | *Step2.* Add the following line into your file: 38 | 39 | ```shell 40 | source /opt/intel/openvino_2021/bin/setupvars.sh 41 | ``` 42 | 43 | *Step3.* Save and exit the file, then run: 44 | 45 | ```shell 46 | source ~/.bashrc 47 | ``` 48 | 49 | 50 | ## Convert model 51 | 52 | 1. Export ONNX model 53 | 54 | Please refer to the [ONNX toturial](../../ONNXRuntime). **Note that you should set --opset to 10, otherwise your next step will fail.** 55 | 56 | 2. Convert ONNX to OpenVINO 57 | 58 | ``` shell 59 | cd /openvino_2021/deployment_tools/model_optimizer 60 | ``` 61 | 62 | Install requirements for convert tool 63 | 64 | ```shell 65 | sudo ./install_prerequisites/install_prerequisites_onnx.sh 66 | ``` 67 | 68 | Then convert model. 69 | ```shell 70 | python3 mo.py --input_model --input_shape [--data_type FP16] 71 | ``` 72 | For example: 73 | ```shell 74 | python3 mo.py --input_model yolox.onnx --input_shape (1,3,640,640) --data_type FP16 75 | ``` 76 | 77 | ## Build 78 | 79 | ### Linux 80 | ```shell 81 | source /opt/intel/openvino_2021/bin/setupvars.sh 82 | mkdir build 83 | cd build 84 | cmake .. 85 | make 86 | ``` 87 | 88 | ## Demo 89 | 90 | ### c++ 91 | 92 | ```shell 93 | ./yolox_openvino 94 | ``` 95 | -------------------------------------------------------------------------------- /demo/OpenVINO/python/README.md: -------------------------------------------------------------------------------- 1 | # YOLOX-OpenVINO in Python 2 | 3 | This toturial includes a Python demo for OpenVINO, as well as some converted models. 4 | 5 | ### Download OpenVINO models. 6 | | Model | Parameters | GFLOPs | Test Size | mAP | Weights | 7 | |:------| :----: | :----: | :---: | :---: | :---: | 8 | | [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) | 9 | | [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) | 10 | | [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) | 11 | | [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) | 12 | | [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) | 13 | | [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) | 14 | | [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) | 15 | 16 | ## Install OpenVINO Toolkit 17 | 18 | Please visit [Openvino Homepage](https://docs.openvinotoolkit.org/latest/get_started_guides.html) for more details. 19 | 20 | ## Set up the Environment 21 | 22 | ### For Linux 23 | 24 | **Option1. Set up the environment tempororally. You need to run this command everytime you start a new shell window.** 25 | 26 | ```shell 27 | source /opt/intel/openvino_2021/bin/setupvars.sh 28 | ``` 29 | 30 | **Option2. Set up the environment permenantly.** 31 | 32 | *Step1.* For Linux: 33 | ```shell 34 | vim ~/.bashrc 35 | ``` 36 | 37 | *Step2.* Add the following line into your file: 38 | 39 | ```shell 40 | source /opt/intel/openvino_2021/bin/setupvars.sh 41 | ``` 42 | 43 | *Step3.* Save and exit the file, then run: 44 | 45 | ```shell 46 | source ~/.bashrc 47 | ``` 48 | 49 | 50 | ## Convert model 51 | 52 | 1. Export ONNX model 53 | 54 | Please refer to the [ONNX toturial](../../ONNXRuntime). **Note that you should set --opset to 10, otherwise your next step will fail.** 55 | 56 | 2. Convert ONNX to OpenVINO 57 | 58 | ``` shell 59 | cd /openvino_2021/deployment_tools/model_optimizer 60 | ``` 61 | 62 | Install requirements for convert tool 63 | 64 | ```shell 65 | sudo ./install_prerequisites/install_prerequisites_onnx.sh 66 | ``` 67 | 68 | Then convert model. 69 | ```shell 70 | python3 mo.py --input_model --input_shape [--data_type FP16] 71 | ``` 72 | For example: 73 | ```shell 74 | python3 mo.py --input_model yolox.onnx --input_shape [1,3,640,640] --data_type FP16 --output_dir converted_output 75 | ``` 76 | 77 | ## Demo 78 | 79 | ### python 80 | 81 | ```shell 82 | python openvino_inference.py -m -i 83 | ``` 84 | or 85 | ```shell 86 | python openvino_inference.py -m -i -o -s -d 87 | ``` 88 | 89 | -------------------------------------------------------------------------------- /demo/TensorRT/cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | 3 | project(yolox) 4 | 5 | add_definitions(-std=c++11) 6 | 7 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 8 | set(CMAKE_CXX_STANDARD 11) 9 | set(CMAKE_BUILD_TYPE Debug) 10 | 11 | find_package(CUDA REQUIRED) 12 | 13 | include_directories(${PROJECT_SOURCE_DIR}/include) 14 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 15 | # cuda 16 | include_directories(/data/cuda/cuda-10.2/cuda/include) 17 | link_directories(/data/cuda/cuda-10.2/cuda/lib64) 18 | # cudnn 19 | include_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/include) 20 | link_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/lib64) 21 | # tensorrt 22 | include_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/include) 23 | link_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/lib) 24 | 25 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") 26 | 27 | find_package(OpenCV) 28 | include_directories(${OpenCV_INCLUDE_DIRS}) 29 | 30 | add_executable(yolox ${PROJECT_SOURCE_DIR}/yolox.cpp) 31 | target_link_libraries(yolox nvinfer) 32 | target_link_libraries(yolox cudart) 33 | target_link_libraries(yolox ${OpenCV_LIBS}) 34 | 35 | add_definitions(-O2 -pthread) 36 | 37 | -------------------------------------------------------------------------------- /demo/TensorRT/cpp/README.md: -------------------------------------------------------------------------------- 1 | # YOLOX-TensorRT in C++ 2 | 3 | As YOLOX models is easy to converted to tensorrt using [torch2trt gitrepo](https://github.com/NVIDIA-AI-IOT/torch2trt), 4 | our C++ demo will not include the model converting or constructing like other tenorrt demos. 5 | 6 | 7 | ## Step 1: Prepare serialized engine file 8 | 9 | Follow the trt [python demo README](../python/README.md) to convert and save the serialized engine file. 10 | 11 | Check the 'model_trt.engine' file generated from Step 1, which will automatically saved at the current demo dir. 12 | 13 | 14 | ## Step 2: build the demo 15 | 16 | Please follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) to install TensorRT. 17 | 18 | Install opencv with ```sudo apt-get install libopencv-dev```. 19 | 20 | build the demo: 21 | 22 | ```shell 23 | mkdir build 24 | cd build 25 | cmake .. 26 | make 27 | ``` 28 | 29 | Then run the demo: 30 | 31 | ```shell 32 | ./yolox ../model_trt.engine -i ../../../../assets/dog.jpg 33 | ``` 34 | 35 | or 36 | 37 | ```shell 38 | ./yolox -i 39 | ``` 40 | -------------------------------------------------------------------------------- /demo/TensorRT/python/README.md: -------------------------------------------------------------------------------- 1 | # YOLOX-TensorRT in Python 2 | 3 | This toturial includes a Python demo for TensorRT. 4 | 5 | ## Install TensorRT Toolkit 6 | 7 | Please 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. 8 | 9 | ## Convert model 10 | 11 | YOLOX models can be easily conveted to TensorRT models using torch2trt 12 | 13 | If you want to convert our model, use the flag -n to specify a model name: 14 | ```shell 15 | python tools/trt.py -n -c 16 | ``` 17 | For example: 18 | ```shell 19 | python tools/trt.py -n yolox-s -c your_ckpt.pth.tar 20 | ``` 21 | can be: yolox-nano, yolox-tiny. yolox-s, yolox-m, yolox-l, yolox-x. 22 | 23 | If you want to convert your customized model, use the flag -f to specify you exp file: 24 | ```shell 25 | python tools/trt.py -f -c 26 | ``` 27 | For example: 28 | ```shell 29 | python tools/trt.py -f /path/to/your/yolox/exps/yolox_s.py -c your_ckpt.pth.tar 30 | ``` 31 | *yolox_s.py* can be any exp file modified by you. 32 | 33 | The converted model and the serialized engine file (for C++ demo) will be saved on your experiment output dir. 34 | 35 | ## Demo 36 | 37 | The TensorRT python demo is merged on our pytorch demo file, so you can run the pytorch demo command with ```--trt```. 38 | 39 | ```shell 40 | python tools/demo.py image -n yolox-s --trt --save_result 41 | ``` 42 | or 43 | ```shell 44 | python tools/demo.py image -f exps/default/yolox_s.py --trt --save_result 45 | ``` 46 | 47 | -------------------------------------------------------------------------------- /demo/ncnn/android/README.md: -------------------------------------------------------------------------------- 1 | # YOLOX-Android-ncnn 2 | 3 | Andoird app of YOLOX object detection base on [ncnn](https://github.com/Tencent/ncnn) 4 | 5 | 6 | ## Tutorial 7 | 8 | ### Step1 9 | 10 | Download ncnn-android-vulkan.zip from [releases of ncnn](https://github.com/Tencent/ncnn/releases). This repo uses 11 | [20210525 release](https://github.com/Tencent/ncnn/releases/download/20210525/ncnn-20210525-android-vulkan.zip) for building. 12 | 13 | ### Step2 14 | 15 | After downloading, please extract your zip file. Then, there are two ways to finish this step: 16 | * put your extracted directory into **app/src/main/jni** 17 | * change the **ncnn_DIR** path in **app/src/main/jni/CMakeLists.txt** to your extracted directory 18 | 19 | ### Step3 20 | Download 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**. 21 | 22 | ### Step4 23 | Open this project with Android Studio, build it and enjoy! 24 | 25 | ## Reference 26 | 27 | * [ncnn-android-yolov5](https://github.com/nihui/ncnn-android-yolov5) 28 | -------------------------------------------------------------------------------- /demo/ncnn/android/app/build.gradle: -------------------------------------------------------------------------------- 1 | apply plugin: 'com.android.application' 2 | 3 | android { 4 | compileSdkVersion 24 5 | buildToolsVersion "29.0.2" 6 | 7 | defaultConfig { 8 | applicationId "com.megvii.yoloXncnn" 9 | archivesBaseName = "$applicationId" 10 | 11 | ndk { 12 | moduleName "ncnn" 13 | abiFilters "armeabi-v7a", "arm64-v8a" 14 | } 15 | minSdkVersion 24 16 | } 17 | 18 | externalNativeBuild { 19 | cmake { 20 | version "3.10.2" 21 | path file('src/main/jni/CMakeLists.txt') 22 | } 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /demo/ncnn/android/app/src/main/AndroidManifest.xml: -------------------------------------------------------------------------------- 1 | 2 | 6 | 7 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | -------------------------------------------------------------------------------- /demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/YOLOXncnn.java: -------------------------------------------------------------------------------- 1 | // Copyright (C) Megvii, Inc. and its affiliates. All rights reserved. 2 | 3 | package com.megvii.yoloXncnn; 4 | 5 | import android.content.res.AssetManager; 6 | import android.graphics.Bitmap; 7 | 8 | public class YOLOXncnn 9 | { 10 | public native boolean Init(AssetManager mgr); 11 | 12 | public class Obj 13 | { 14 | public float x; 15 | public float y; 16 | public float w; 17 | public float h; 18 | public String label; 19 | public float prob; 20 | } 21 | 22 | public native Obj[] Detect(Bitmap bitmap, boolean use_gpu); 23 | 24 | static { 25 | System.loadLibrary("yoloXncnn"); 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/yoloXncnn.java: -------------------------------------------------------------------------------- 1 | // Copyright (C) Megvii, Inc. and its affiliates. All rights reserved. 2 | 3 | package com.megvii.yoloXncnn; 4 | 5 | import android.content.res.AssetManager; 6 | import android.graphics.Bitmap; 7 | 8 | public class YOLOXncnn 9 | { 10 | public native boolean Init(AssetManager mgr); 11 | 12 | public class Obj 13 | { 14 | public float x; 15 | public float y; 16 | public float w; 17 | public float h; 18 | public String label; 19 | public float prob; 20 | } 21 | 22 | public native Obj[] Detect(Bitmap bitmap, boolean use_gpu); 23 | 24 | static { 25 | System.loadLibrary("yoloXncnn"); 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /demo/ncnn/android/app/src/main/jni/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | project(yoloXncnn) 2 | 3 | cmake_minimum_required(VERSION 3.4.1) 4 | 5 | set(ncnn_DIR ${CMAKE_SOURCE_DIR}/ncnn-20210525-android-vulkan/${ANDROID_ABI}/lib/cmake/ncnn) 6 | find_package(ncnn REQUIRED) 7 | 8 | add_library(yoloXncnn SHARED yoloXncnn_jni.cpp) 9 | 10 | target_link_libraries(yoloXncnn 11 | ncnn 12 | 13 | jnigraphics 14 | ) 15 | -------------------------------------------------------------------------------- /demo/ncnn/android/app/src/main/res/layout/main.xml: -------------------------------------------------------------------------------- 1 | 2 | 6 | 7 | 11 | 12 |