Spaces:
Running
Running
File size: 2,209 Bytes
b87a24a f048b9a b87a24a a869123 b87a24a 1e649ad b87a24a a869123 b87a24a a869123 f048b9a a869123 f048b9a a869123 f048b9a a869123 f048b9a 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 57 58 59 60 61 62 63 64 | 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")
|