Spaces:
Running
Running
File size: 1,763 Bytes
b87a24a a869123 b87a24a 1e649ad b87a24a a869123 b87a24a a869123 b87a24a a869123 b87a24a a869123 1e649ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | 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")
|