ayzeksalimli/tram-dynamic-roi-tracker

1from __future__ import annotations2 3import shutil4import sys5from pathlib import Path6 7_ROOT = Path(__file__).resolve().parent.parent8MODELS_DIR = _ROOT / "models"9 10SEGMENTOR_REPO_ID = "ayzeksalimli/tram-dynamic-roi-tracker-yolo11s"11SEGMENTOR_FILENAME = "tram-dynamic-roi-tracker-yolo11s.pt"12SEGMENTOR_DIR = MODELS_DIR / "segmentor"13 14DETECTOR_FILENAME = "yolov8s.pt"15DETECTOR_DIR = MODELS_DIR / "detector"16 17TRACKER_FILENAME = "bytetrack.yaml"18TRACKER_DIR = MODELS_DIR / "trackers"19 20 21def download_segmentor() -> Path:22    from huggingface_hub import hf_hub_download23 24    SEGMENTOR_DIR.mkdir(parents=True, exist_ok=True)25    dest = SEGMENTOR_DIR / SEGMENTOR_FILENAME26    if dest.exists():27        print(f"[segmentor] already present: {dest}")28        return dest29 30    print(f"[segmentor] downloading {SEGMENTOR_REPO_ID} ...")31    local_path = hf_hub_download(repo_id=SEGMENTOR_REPO_ID, filename=SEGMENTOR_FILENAME)32    shutil.copy2(local_path, dest)33    print(f"[segmentor] saved to {dest}")34    return dest35 36 37def download_detector() -> Path:38    from ultralytics.utils.downloads import attempt_download_asset39 40    DETECTOR_DIR.mkdir(parents=True, exist_ok=True)41    dest = DETECTOR_DIR / DETECTOR_FILENAME42    if dest.exists():43        print(f"[detector] already present: {dest}")44        return dest45 46    print(f"[detector] downloading {DETECTOR_FILENAME} ...")47    attempt_download_asset(str(dest))48    print(f"[detector] saved to {dest}")49    return dest50 51 52def download_tracker() -> Path:53    import ultralytics54 55    TRACKER_DIR.mkdir(parents=True, exist_ok=True)56    dest = TRACKER_DIR / TRACKER_FILENAME57    if dest.exists():58        print(f"[tracker] already present: {dest}")59        return dest60 61    src = Path(ultralytics.__file__).resolve().parent / "cfg" / "trackers" / TRACKER_FILENAME62    if not src.exists():63        print(f"[tracker] ERROR: bundled config not found at {src}", file=sys.stderr)64        sys.exit(1)65    shutil.copy2(src, dest)66    print(f"[tracker] saved to {dest}")67    return dest68 69 70def main() -> None:71    download_segmentor()72    download_detector()73    download_tracker()74    print("\nAll models are ready in", MODELS_DIR)75 76 77if __name__ == "__main__":78    main()79