feat: add TensorRT engine export script
This commit is contained in:
1 parent
06b013a778
commit
d32045b433
1 file changed
+39
Executable
+39
@@ -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()
|
||||||
Reference in new issue
Block a user