backend/model.py

import os
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 _auth_header():
    token = os.getenv("HF_TOKEN", "")
    return f'-H "Authorization: Bearer {token}"' if token else ""


def _download_ov_model():
    """Download the pre-built OV INT8 model files directly from HF."""
    OV_DIR.mkdir(parents=True, exist_ok=True)
    auth = _auth_header()
    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} ...")
        ret = os.system(f'curl -L --fail {auth} -o "{dest}" "{url}"')
        if ret != 0 or not dest.exists() or dest.stat().st_size < 100:
            raise RuntimeError(
                f"[model] Failed to download {filename} from HF. "
                "Check HF_TOKEN and repo visibility."
            )
    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")