#!/usr/bin/env python3
"""Export all .pt models to TensorRT .engine format for Jetson."""

from pathlib import Path
from ultralytics import YOLO

MODELS_DIR = Path(__file__).resolve().parent.parent / "models"

def main():
    pt_files = sorted(MODELS_DIR.glob("*.pt"))
    if not pt_files:
        print("No .pt files found in", MODELS_DIR)
        return

    success = 0
    failed = 0
    skipped = 0

    for pt_path in pt_files:
        engine_path = pt_path.with_suffix(".engine")
        if engine_path.exists():
            print(f"[SKIP] {pt_path.name} — .engine already exists")
            skipped += 1
            continue

        print(f"[INFO] Exporting {pt_path.name} to TensorRT engine...")
        try:
            model = YOLO(str(pt_path))
            engine_path_str = model.export(format="engine", device=0, half=True, imgsz=640)
            print(f"[OK]   Exported: {engine_path_str}")
            success += 1
        except Exception as e:
            print(f"[FAIL] {pt_path.name}: {e}")
            failed += 1

    print(f"\nSummary: {success} exported, {skipped} skipped, {failed} failed")

if __name__ == "__main__":
    main()