abhishek-gola's picture
Yolo Tflite Models (#14)
4bbaf21
|
Raw
History Blame Contribute Delete
2.78 kB

YOLOv5nu TFLite Conversion

YOLOv5nu is the anchor-free ("u") variant of YOLOv5 from Ultralytics. It uses the YOLOv8-style decoupled, anchor-free detection head, so its output is a single [1, 84, 8400] tensor (84 = 4 box coordinates + 80 COCO class scores, no separate objectness channel).

The exported TFLite outputs box coordinates normalized to [0, 1] (the demo's post-processing accounts for this).

Prerequisites

Python environment

The .pt → TFLite conversion uses Ultralytics. Python 3.10 was used.

conda create -n <env_name> python=3.10 -y
conda activate <env_name>
pip install "ultralytics==8.4.90" "torch==2.12.1" "torchvision==0.27.1" \
            "tensorflow==2.21.0" "ai-edge-litert==2.1.5"

These are the exact versions used to produce the provided yolov5nu.tflite (Python 3.10.20). tensorflow is required for the TFLite export, and torchvision must match torch (0.27.1 ↔ 2.12.1).


Conversion of YOLOv5nu to TFLite

The yolov5nu.pt checkpoint is downloaded automatically by Ultralytics on first use.

# CLI
yolo export model=yolov5nu.pt format=tflite imgsz=640

or equivalently with the Python API:

from ultralytics import YOLO
YOLO("yolov5nu.pt").export(format="tflite", imgsz=640)

Ultralytics writes the export artifacts to yolov5nu_saved_model/; the float32 model is yolov5nu_saved_model/yolov5nu_float32.tflite. Copy/rename it to yolov5nu.tflite (the name the demo expects).


Usage

A demo script is provided to run inference using OpenCV DNN:

python demo.py --model yolov5nu.tflite --image example_outputs/input.jpg --output example_outputs/yolov5nu_output.jpg

demo.py uses only OpenCV's dnn module (cv2.dnn.readNet). It letterboxes the input to 640x640 (aspect-preserving resize + centered 114-pad, matching ultralytics.data.augment.LetterBox), runs the network, decodes the [1, 84, 8400] head (argmax over the 80 class scores; boxes are normalized [0, 1], scaled to the letterboxed input and mapped back through the inverse letterbox), applies class-aware NMS (NMSBoxesBatched), prints the detected COCO classes / confidences / bounding boxes, and saves an annotated output image. It also handles the NMS-free (v10/26-style) [1, 300, 6+] head, so the same script works for other exported YOLO .tflite models.

On the provided example_outputs/input.jpg it detects dog, truck and bicycle.


License

YOLOv5nu weights and the Ultralytics exporter are released under AGPL-3.0 — see the Ultralytics LICENSE. Any redistribution of the converted yolov5nu.tflite is subject to those terms.