Intel/delivery-package-verification

1#!/usr/bin/env bash2# SPDX-License-Identifier: MIT3# Copyright (C) Intel Corporation4#5# Export a YOLO26 detector to OpenVINO IR for the delivery-package-verification6# use case (COCO backpack/handbag/suitcase relabeled to "package" at runtime).7# Usage: ./export_and_quantize.sh [MODEL_VARIANT] [PRECISION]8# Example: ./export_and_quantize.sh yolo26n FP169 10set -euo pipefail11 12MODEL_NAME="${1:-yolo26n}"13PRECISION="${2:-FP16}"14PRECISION="$(echo "${PRECISION}" | tr '[:lower:]' '[:upper:]')"15 16if [[ "${PRECISION}" != "FP32" && "${PRECISION}" != "FP16" && "${PRECISION}" != "INT8" ]]; then17    echo "ERROR: unsupported precision '${PRECISION}'. Choose FP32, FP16, or INT8." >&218    exit 119fi20 21echo "--- Installing dependencies ---"22if [[ "${PRECISION}" == "INT8" ]]; then23    pip install -qU openvino nncf ultralytics24else25    pip install -qU openvino ultralytics26fi27 28# Ask for approval before downloading models and sample files29echo ""30echo "This script will download:"31echo "  - Model weights and/or sample files"32echo ""33read -p "Continue with downloads? (yes/no): " APPROVAL34if [[ "${APPROVAL}" != "yes" ]]; then35    echo "Download cancelled by user."36    exit 037fi38 39# Ping the HuggingFace repo to register a tracked download of config.json.40# This is best-effort: a failure (offline, or before the repo is published)41# must not stop the export.42echo "--- Registering HuggingFace download (tracking ping) ---"43HF_REPO_ID="Intel/delivery-package-verification"44HF_CONFIG_URL="https://huggingface.co/${HF_REPO_ID}/resolve/main/config.json"45if curl -fsSL -o /dev/null "${HF_CONFIG_URL}"; then46    echo "Registered HuggingFace download for ${HF_REPO_ID}"47else48    echo "WARNING: HuggingFace tracking ping failed (offline?); continuing." >&249fi50echo ""51echo "--- Downloading sample test image ---"52if [[ ! -f test.jpg ]]; then53    wget -q -O test.jpg https://ultralytics.com/images/bus.jpg54    echo "Downloaded: test.jpg"55else56    echo "Already present: test.jpg"57fi58echo ""59echo "--- Downloading sample test video ---"60if [[ ! -f test_video.mp4 ]]; then61    wget -q -O test_video.mp4 \62        "https://www.pexels.com/download/video/6170052/?fps=25&w=540&h=960"63    echo "Downloaded: test_video.mp4"64else65    echo "Already present: test_video.mp4"66fi67 68if [[ "${PRECISION}" == "FP32" ]]; then69    HALF_FLAG="False"70    EXPORT_LABEL="FP32"71else72    HALF_FLAG="True"73    EXPORT_LABEL="FP16"74fi75 76echo "--- Exporting ${MODEL_NAME} to OpenVINO IR (${EXPORT_LABEL}) ---"77python3 -c "78from ultralytics import YOLO79 80model = YOLO('${MODEL_NAME}.pt')81model.export(format='openvino', half=${HALF_FLAG}, dynamic=False, imgsz=640)82print('Export complete: ${MODEL_NAME}_openvino_model/')83"84 85echo "--- Writing package label map (relabels backpack/handbag/suitcase -> package) ---"86python3 - "${MODEL_NAME}" <<'PY'87import sys88import yaml89 90name = sys.argv[1]91with open(f"{name}_openvino_model/metadata.yaml") as f:92    meta = yaml.safe_load(f)93names = meta["names"]94labels = [names[i] for i in range(len(names))]95for i in (24, 26, 28):  # backpack, handbag, suitcase -> package96    labels[i] = "package"97with open("coco_package_labels.txt", "w") as f:98    f.write("\n".join(labels) + "\n")99print(f"Wrote coco_package_labels.txt ({len(labels)} labels)")100PY101 102if [[ "${PRECISION}" == "INT8" ]]; then103    echo "--- Quantizing to INT8 with NNCF ---"104    python3 -c "105import nncf106import openvino as ov107import numpy as np108import cv2109 110core = ov.Core()111model = core.read_model('${MODEL_NAME}_openvino_model/${MODEL_NAME}.xml')112 113img = cv2.imread('test.jpg')114img = cv2.resize(img, (640, 640))115img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0116img = img.transpose(2, 0, 1)[np.newaxis, ...]117 118def transform_fn(data_item):119    return img120 121calibration_dataset = nncf.Dataset(list(range(300)), transform_fn)122 123quantized = nncf.quantize(124    model,125    calibration_dataset,126    preset=nncf.QuantizationPreset.MIXED,127    subset_size=300,128)129 130ov.save_model(quantized, '${MODEL_NAME}_package_int8.xml')131print('Quantization complete: ${MODEL_NAME}_package_int8.xml')132"133fi134echo "--- Done ---"135