diff --git a/scripts/export_engines.py b/scripts/export_engines.py new file mode 100755 index 0000000..76868c6 --- /dev/null +++ b/scripts/export_engines.py @@ -0,0 +1,39 @@ +#!/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() \ No newline at end of file