backend/model.py

import os
import shutil
import urllib.request
from pathlib import Path
from dotenv import load_dotenv
from ultralytics import YOLO

load_dotenv()

# Silence Ultralytics config dir warning on read-only containers
os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics")

MODEL_DIR = Path(__file__).parent / "weights"
OV_DIR = MODEL_DIR / "best_int8_openvino_model"

_HF_BASE = (
    "https://huggingface.co/Perception365/VehicleNet-Y26s"
    "/resolve/main/weights/best_int8_openvino_model"
)

# The three files that make up the OpenVINO INT8 model
_OV_FILES = ["best.bin", "best.xml", "metadata.yaml"]


def _download_ov_model():
    """Download the pre-built OV INT8 model files directly from HF."""
    OV_DIR.mkdir(parents=True, exist_ok=True)
    token = os.getenv("HF_TOKEN", "")
    headers = {"Authorization": f"Bearer {token}"} if token else {}

    for filename in _OV_FILES:
        dest = OV_DIR / filename
        if dest.exists():
            print(f"[model] {filename} already present, skipping.")
            continue
        url = f"{_HF_BASE}/{filename}"
        print(f"[model] Downloading {filename} ...")
        try:
            # urlopen, not `os.system("curl ...")` — keeps HF_TOKEN out of the
            # shell command line (visible in /proc) and out of shell quoting.
            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req) as resp, open(dest, "wb") as f:
                shutil.copyfileobj(resp, f)
        except Exception as e:
            dest.unlink(missing_ok=True)
            raise RuntimeError(
                f"[model] Failed to download {filename} from HF ({e}). "
                "Check HF_TOKEN and repo visibility."
            )
        if dest.stat().st_size < 100:
            dest.unlink(missing_ok=True)
            raise RuntimeError(f"[model] {filename} downloaded but is truncated.")
    print("[model] OV model files ready ✅")


def load_model():
    if not OV_DIR.exists() or not all(
        (OV_DIR / f).exists() for f in _OV_FILES
    ):
        _download_ov_model()
    else:
        print("[model] OV model already present — skipping download.")
    return YOLO(str(OV_DIR), task="detect")