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")