Intel/object-classification

1#!/usr/bin/env bash2# SPDX-License-Identifier: MIT3# Copyright (C) Intel Corporation4#5# Export a YOLO26 detector to OpenVINO IR for the object-classification use case.6# The accompanying samples group detected objects into higher-level city7# categories (People, Vehicles) for traffic situational awareness.8# Usage: ./export_and_quantize.sh [MODEL_VARIANT] [PRECISION]9# Example: ./export_and_quantize.sh yolo26n FP1610#11# Supported precisions:12#   FP32  -- Full-precision floating-point weights13#   FP16  -- Half-precision floating-point weights (default)14#   INT8  -- Quantized 8-bit integer weights (requires NNCF)15#16# Precision / device compatibility:17#   | Precision | CPU | GPU | NPU |18#   |-----------|-----|-----|-----|19#   | FP32      | Yes | Yes | No  |20#   | FP16      | Yes | Yes | Yes |21#   | INT8      | Yes | Yes | Yes |22 23set -euo pipefail24 25MODEL_NAME="${1:-yolo26n}"26PRECISION="${2:-FP16}"27PRECISION="$(echo "${PRECISION}" | tr '[:lower:]' '[:upper:]')"28 29if [[ "${PRECISION}" != "FP32" && "${PRECISION}" != "FP16" && "${PRECISION}" != "INT8" ]]; then30    echo "ERROR: unsupported precision '${PRECISION}'. Choose FP32, FP16, or INT8." >&231    exit 132fi33 34echo "--- Installing dependencies ---"35if [[ "${PRECISION}" == "INT8" ]]; then36    pip install -qU openvino nncf ultralytics37else38    pip install -qU openvino ultralytics39fi40 41# Ask for approval before downloading models and sample files42echo ""43echo "This script will download:"44echo "  - Model weights and a sample traffic video"45echo ""46read -p "Continue with downloads? (yes/no): " APPROVAL47if [[ "${APPROVAL}" != "yes" ]]; then48    echo "Download cancelled by user."49    exit 050fi51 52# Ping the HuggingFace repo to register a tracked download of config.json.53# This is best-effort: a failure (offline, or before the repo is published)54# must not stop the export.55echo "--- Registering HuggingFace download (tracking ping) ---"56HF_REPO_ID="Intel/object-classification"57HF_CONFIG_URL="https://huggingface.co/${HF_REPO_ID}/resolve/main/config.json"58if curl -fsSL -o /dev/null "${HF_CONFIG_URL}"; then59    echo "Registered HuggingFace download for ${HF_REPO_ID}"60else61    echo "WARNING: HuggingFace tracking ping failed (offline?); continuing." >&262fi63 64echo ""65echo "--- Downloading sample test video (urban roundabout) ---"66if [[ ! -f test_video.mp4 ]]; then67    wget -q -O test_video.mp4 \68        "https://www.pexels.com/download/video/30119018/?fps=59.94&h=720&w=1280"69    echo "Downloaded: test_video.mp4"70else71    echo "Already present: test_video.mp4"72fi73 74if [[ "${PRECISION}" == "FP32" ]]; then75    HALF_FLAG="False"76    EXPORT_LABEL="FP32"77else78    HALF_FLAG="True"79    EXPORT_LABEL="FP16"80fi81 82echo "--- Exporting ${MODEL_NAME} to OpenVINO IR (${EXPORT_LABEL}) ---"83python3 -c "84from ultralytics import YOLO85 86model = YOLO('${MODEL_NAME}.pt')87model.export(format='openvino', half=${HALF_FLAG}, dynamic=False, imgsz=640)88print('Export complete: ${MODEL_NAME}_openvino_model/')89"90 91if [[ "${PRECISION}" == "INT8" ]]; then92    echo "--- Quantizing to INT8 with NNCF ---"93    python3 -c "94import nncf95import openvino as ov96import numpy as np97import cv298 99core = ov.Core()100model = core.read_model('${MODEL_NAME}_openvino_model/${MODEL_NAME}.xml')101 102# Calibrate on a representative frame from the sample traffic video.103cap = cv2.VideoCapture('test_video.mp4')104ok, frame = cap.read()105cap.release()106if not ok:107    raise SystemExit('Could not read a calibration frame from test_video.mp4')108img = cv2.resize(frame, (640, 640))109img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0110img = img.transpose(2, 0, 1)[np.newaxis, ...]  # NCHW111 112def transform_fn(data_item):113    return img114 115calibration_dataset = nncf.Dataset(list(range(300)), transform_fn)116 117quantized = nncf.quantize(118    model,119    calibration_dataset,120    preset=nncf.QuantizationPreset.MIXED,121    subset_size=300,122)123 124ov.save_model(quantized, '${MODEL_NAME}_objcls_int8.xml')125print('Quantization complete: ${MODEL_NAME}_objcls_int8.xml')126"127fi128echo "--- Done ---"129