Intel/object-detection

1#!/usr/bin/env bash2# SPDX-License-Identifier: MIT3# Copyright (C) Intel Corporation4#5# Export a YOLO26 general-purpose object detector to OpenVINO IR.6# Usage: ./export_and_quantize.sh [MODEL_VARIANT] [PRECISION]7# Example: ./export_and_quantize.sh yolo26n FP168#9# Supported precisions:10#   FP32  -- Full-precision floating-point weights11#   FP16  -- Half-precision floating-point weights (default)12#   INT8  -- Quantized 8-bit integer weights (requires NNCF)13#14# Precision / device compatibility:15#   | Precision | CPU | GPU | NPU |16#   |-----------|-----|-----|-----|17#   | FP32      | Yes | Yes | No  |18#   | FP16      | Yes | Yes | Yes |19#   | INT8      | Yes | Yes | Yes |20 21set -euo pipefail22 23MODEL_NAME="${1:-yolo26n}"24PRECISION="${2:-FP16}"25PRECISION="$(echo "${PRECISION}" | tr '[:lower:]' '[:upper:]')"26 27if [[ "${PRECISION}" != "FP32" && "${PRECISION}" != "FP16" && "${PRECISION}" != "INT8" ]]; then28    echo "ERROR: unsupported precision '${PRECISION}'. Choose FP32, FP16, or INT8." >&229    exit 130fi31 32echo "--- Installing dependencies ---"33if [[ "${PRECISION}" == "INT8" ]]; then34    pip install -qU openvino nncf ultralytics35else36    pip install -qU openvino ultralytics37fi38 39# Ask for approval before downloading models and sample files40echo ""41echo "This script will download:"42echo "  - Model weights and/or sample files"43echo ""44read -p "Continue with downloads? (yes/no): " APPROVAL45if [[ "${APPROVAL}" != "yes" ]]; then46    echo "Download cancelled by user."47    exit 048fi49 50# Ping the HuggingFace repo to register a tracked download of config.json.51# This is best-effort: a failure (offline, or before the repo is published)52# must not stop the export.53echo "--- Registering HuggingFace download (tracking ping) ---"54HF_REPO_ID="Intel/object-detection"55HF_CONFIG_URL="https://huggingface.co/${HF_REPO_ID}/resolve/main/config.json"56if curl -fsSL -o /dev/null "${HF_CONFIG_URL}"; then57    echo "Registered HuggingFace download for ${HF_REPO_ID}"58else59    echo "WARNING: HuggingFace tracking ping failed (offline?); continuing." >&260fi61echo ""62echo "--- Downloading sample test image ---"63if [[ ! -f test.jpg ]]; then64    wget -q -O test.jpg https://ultralytics.com/images/bus.jpg65    echo "Downloaded: test.jpg"66else67    echo "Already present: test.jpg"68fi69echo ""70echo "--- Downloading sample test video ---"71if [[ ! -f test_video.mp4 ]]; then72    wget -q -O test_video.mp4 \73        "https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Coco%20Walking%20in%20Berkeley.mp4"74    echo "Downloaded: test_video.mp4"75else76    echo "Already present: test_video.mp4"77fi78 79if [[ "${PRECISION}" == "FP32" ]]; then80    HALF_FLAG="False"81    EXPORT_LABEL="FP32"82else83    HALF_FLAG="True"84    EXPORT_LABEL="FP16"85fi86 87echo "--- Exporting ${MODEL_NAME} to OpenVINO IR (${EXPORT_LABEL}) ---"88python3 -c "89from ultralytics import YOLO90 91model = YOLO('${MODEL_NAME}.pt')92model.export(format='openvino', half=${HALF_FLAG}, dynamic=False, imgsz=640)93print('Export complete: ${MODEL_NAME}_openvino_model/')94"95 96if [[ "${PRECISION}" == "INT8" ]]; then97    echo "--- Quantizing to INT8 with NNCF ---"98    python3 -c "99import nncf100import openvino as ov101import numpy as np102import cv2103 104core = ov.Core()105model = core.read_model('${MODEL_NAME}_openvino_model/${MODEL_NAME}.xml')106 107# Use the downloaded test image for calibration instead of random noise.108img = cv2.imread('test.jpg')109img = cv2.resize(img, (640, 640))110img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0111img = img.transpose(2, 0, 1)[np.newaxis, ...]  # NCHW112 113def transform_fn(data_item):114    return img115 116calibration_dataset = nncf.Dataset(list(range(300)), transform_fn)117 118quantized = nncf.quantize(119    model,120    calibration_dataset,121    preset=nncf.QuantizationPreset.MIXED,122    subset_size=300,123)124 125ov.save_model(quantized, '${MODEL_NAME}_objdet_int8.xml')126print('Quantization complete: ${MODEL_NAME}_objdet_int8.xml')127"128fi129echo "--- Done ---"130