Awesome
YOLOX-ti-lite Object Detector in TFLite
This repository is a fork of TexasInstruments/edgeai-yolox.
Based on the TexasInstruments' repository, following new features and minor modification were added:
- ONNX -> TFLite export with PINTO0309/onnx2tf
- TFLite int8 quantization with PINTO0309/onnx2tf
- Minor fix of YOLOX head (see here)
Pretrained models
TFLite models exported with this repository are available from my GoogleDrive.
Requirements
- Docker Compose
Setup
Download COCO dataset:
./scripts/get_coco.sh
Download YOLOX-nanoti-lite pretrained weight:
wget http://software-dl.ti.com/jacinto7/esd/modelzoo/latest/models/vision/detection/coco/edgeai-yolox/yolox_nano_ti_lite_26p1_41p8_checkpoint.pth
Export PyTorch to ONNX
Build and run torch2onnx
Docker container:
docker compose build torch2onnx
docker compose run --rm torch2onnx bash
Note that (torch2onnx)
in the code blocks below means you have to run the command in torch2onnx
container.
Evaluate PyTorch model (optional):
(torch2onnx) python tools/eval.py -n yolox_nano_ti_lite -c yolox_nano_ti_lite_26p1_41p8_checkpoint.pth --conf 0.001
Expected result:
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.262
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.421
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.273
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.095
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.280
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.405
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.245
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.391
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.426
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.174
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.478
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.620
Export PyTorch to ONNX:
(torch2onnx) python tools/export_onnx.py -n yolox_nano_ti_lite -c yolox_nano_ti_lite_26p1_41p8_checkpoint.pth --output-name yolox_nano_ti_lite.onnx
Run inference on a sample image with ONNX (optional):
(torch2onnx) python demo/ONNXRuntime/onnx_inference.py -m yolox_nano_ti_lite.onnx -i assets/dog.jpg -o tmp/onnx/ -s 0.6 --input_shape 416,416
Export ONNX to TFLite
Build and run onnx2tf
Docker container:
docker compose build onnx2tf
docker compose run --rm onnx2tf bash
Note that (onnx2tf)
in the code blocks below means you have to run the command in onnx2tf
container.
Generate calibration data from COCO train-set (can be skipped if you don't need quantized TFLite models):
(onnx2tf) python demo/TFLite/generate_calib_data.py --img-size 416 416 --n-img 200 -o calib_data_416x416_n200.npy
Export ONNX to TFLite:
# If you need quantized TFlite models:
(onnx2tf) onnx2tf -i yolox_nano_ti_lite.onnx -oiqt -qcind images calib_data_416x416_n200.npy "[[[[0,0,0]]]]" "[[[[1,1,1]]]]"
# if you don't need quantized TFLite models:
(onnx2tf) onnx2tf -i yolox_nano_ti_lite.onnx
Run inference on a sample image with TFLite:
# fp32:
(onnx2tf) TFLITE_PATH=saved_model/yolox_nano_ti_lite_float32.tflite
# fp16:
(onnx2tf) TFLITE_PATH=saved_model/yolox_nano_ti_lite_float16.tflite
# int8 static quantized:
(onnx2tf) TFLITE_PATH=saved_model/yolox_nano_ti_lite_integer_quant.tflite
# check under `saved_model/` for some other models.
(onnx2tf) python demo/TFLite/tflite_inference.py -m $TFLITE_PATH -i assets/dog.jpg -o tmp/tflite/ -s 0.6
TODO
- evaluate mAP of ONNX models
- evaluate mAP of TFLite models
- compare inference speed of ONNX, TFLite fp32, and int8 models
- ONNX/TFlite export of the model decoder and NMS
- add webcam demo