scripts/download_yolo_weights.py

"""
Pull Matt Nudi's honey-bee/drone/queen YOLO weights from Roboflow once,
then copy the .pt file into weights/ so the Space can load it locally
with ultralytics, no Roboflow auth needed at runtime.

Usage:
    py scripts/download_yolo_weights.py
"""

from __future__ import annotations

import os
import shutil
from pathlib import Path

from dotenv import load_dotenv


WORKSPACE = "matt-nudi"
PROJECT = "honey-bee-detection-model-zgjnb"
WEIGHTS_DIR = Path("weights")
TARGET_PT = WEIGHTS_DIR / "honey_bee_detector.pt"


def main() -> None:
    load_dotenv()
    api_key = os.environ.get("ROBOFLOW_API_KEY")
    if not api_key:
        raise SystemExit(
            "Missing ROBOFLOW_API_KEY. Add it to .env at the project root."
        )

    WEIGHTS_DIR.mkdir(exist_ok=True)

    print(f"Connecting to Roboflow as workspace={WORKSPACE!r} ...")
    from roboflow import Roboflow

    rf = Roboflow(api_key=api_key)
    project = rf.workspace(WORKSPACE).project(PROJECT)

    versions = project.versions()
    if not versions:
        raise SystemExit(f"No trained versions found for {PROJECT!r}.")

    # Pick the highest version number with a trained model
    versions_sorted = sorted(versions, key=lambda v: v.version, reverse=True)
    print(f"Available versions: {[v.version for v in versions_sorted]}")

    version_id = versions_sorted[0].version
    print(f"Using latest version: v{version_id}")
    version = project.version(version_id)

    # Approach 1: try to grab the trained weights via the inference package.
    # It downloads to ~/.cache/inference (or similar) on first use.
    print("\nTriggering weight download via roboflow.inference ...")
    try:
        from inference import get_model

        model = get_model(
            model_id=f"{PROJECT}/{version_id}",
            api_key=api_key,
        )
        # Best-effort: tell us where it landed
        cache_root = Path.home() / ".inference"
        pt_files = list(cache_root.rglob("*.pt"))
        if not pt_files:
            cache_root = Path.home() / ".cache" / "inference"
            pt_files = list(cache_root.rglob("*.pt"))
        if pt_files:
            # Pick the most recently modified
            latest = max(pt_files, key=lambda p: p.stat().st_mtime)
            shutil.copy(latest, TARGET_PT)
            print(f"\n Copied weights from {latest}")
            print(f" to {TARGET_PT.resolve()}")
            print(f" size: {TARGET_PT.stat().st_size / 1024 / 1024:.1f} MB")
            return
        print("inference cache had no .pt files yet; falling back to dataset export.")
    except Exception as e:
        print(f"inference path failed: {e}\nFalling back to dataset export ...")

    # Approach 2: download the dataset export, which sometimes ships weights.
    download_dir = Path("data") / "raw" / f"roboflow_{PROJECT}_v{version_id}"
    download_dir.parent.mkdir(parents=True, exist_ok=True)
    dataset = version.download("yolov8", location=str(download_dir))
    print(f"\nDataset downloaded to: {dataset.location}")
    pt_files = list(Path(dataset.location).rglob("*.pt"))
    if pt_files:
        shutil.copy(pt_files[0], TARGET_PT)
        print(f"\n Copied weights to {TARGET_PT.resolve()}")
        print(f" size: {TARGET_PT.stat().st_size / 1024 / 1024:.1f} MB")
        return

    raise SystemExit(
        "\nCould not locate trained .pt weights via either path. "
        "Check the project's available versions and whether the author "
        "published a trained model (some only ship the dataset)."
    )


if __name__ == "__main__":
    main()