#!/usr/bin/env bash
# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Download the pre-trained FP32 Worker Safety Gear Detection model
# from Intel Edge AI Resources.
# Usage: ./export_and_quantize.sh [OUTPUT_DIR]
# Example: ./export_and_quantize.sh ./models

set -euo pipefail

OUTPUT_DIR="${1:-./models}"
MODEL_DIR="${OUTPUT_DIR}/worker-safety-gear-detection"
ZIP_URL="https://github.com/open-edge-platform/edge-ai-resources/raw/main/models/FP32/worker-safety-gear-detection.zip"
ZIP_FILE="${OUTPUT_DIR}/worker-safety-gear-detection.zip"

echo "--- Installing dependencies ---"
pip install -qU "openvino>=2026.0.0"

echo "--- Downloading sample test video ---"
if [[ ! -f test_video.avi ]]; then
    wget -q -O test_video.avi \
        https://github.com/open-edge-platform/edge-ai-resources/raw/refs/heads/main/videos/Safety_Full_Hat_and_Vest.avi
    echo "Downloaded: test_video.avi"
else
    echo "Already present: test_video.avi"
fi

echo "--- Downloading worker-safety-gear-detection FP32 model ---"
mkdir -p "${OUTPUT_DIR}"
if [[ ! -f "${ZIP_FILE}" ]]; then
    wget -q -O "${ZIP_FILE}" "${ZIP_URL}"
    echo "Downloaded: ${ZIP_FILE}"
else
    echo "Already downloaded: ${ZIP_FILE}"
fi

echo "--- Extracting model ---"
unzip -qo "${ZIP_FILE}" -d "${MODEL_DIR}"
echo "Extracted to: ${MODEL_DIR}"

# Locate the model XML.
MODEL_XML=$(find "${MODEL_DIR}" -name "model.xml" -type f | head -1)
if [[ -z "${MODEL_XML}" ]]; then
    echo "ERROR: model.xml not found in extracted archive." >&2
    exit 1
fi
echo "Model IR found: ${MODEL_XML}"
echo "--- Done ---"