export_benchmark.py

import time
import cv2
import numpy as np
from ultralytics import YOLO

# ── 1. Export to ONNX ──────────────────────────────────────────────
print("Exporting YOLOv8-pose to ONNX...")
model = YOLO("yolov8n-pose.pt")
model.export(format="onnx", opset=12)
print("Export done → yolov8n-pose.onnx\n")

# ── 2. Load both models ────────────────────────────────────────────
pytorch_model = YOLO("yolov8n-pose.pt")
onnx_model    = YOLO("yolov8n-pose.onnx")

# ── 3. Grab 100 frames from your test video ────────────────────────
cap = cv2.VideoCapture("test_video.mp4")
frames = []
while len(frames) < 100:
    ret, frame = cap.read()
    if not ret:
        break
    frames.append(frame)
cap.release()

if len(frames) == 0:
    raise RuntimeError("Could not read frames from test_video.mp4")

print(f"Benchmarking on {len(frames)} frames...\n")

# ── 4. Benchmark PyTorch ───────────────────────────────────────────
start = time.time()
for frame in frames:
    pytorch_model(frame, verbose=False)
pytorch_time = time.time() - start
pytorch_fps  = len(frames) / pytorch_time

# ── 5. Benchmark ONNX ─────────────────────────────────────────────
start = time.time()
for frame in frames:
    onnx_model(frame, verbose=False)
onnx_time = time.time() - start
onnx_fps  = len(frames) / onnx_time

# ── 6. Results ────────────────────────────────────────────────────
speedup = onnx_fps / pytorch_fps

print("=" * 40)
print(f"  PyTorch  → {pytorch_fps:.1f} FPS  ({pytorch_time:.2f}s)")
print(f"  ONNX     → {onnx_fps:.1f} FPS  ({onnx_time:.2f}s)")
print(f"  Speedup  → {speedup:.2f}x faster with ONNX")
print("=" * 40)

# ── 7. Save results to file ───────────────────────────────────────
with open("benchmark_results.txt", "w") as f:
    f.write("ONNX Export Benchmark Results\n")
    f.write("=" * 40 + "\n")
    f.write(f"Frames tested : {len(frames)}\n")
    f.write(f"PyTorch FPS   : {pytorch_fps:.1f}\n")
    f.write(f"ONNX FPS      : {onnx_fps:.1f}\n")
    f.write(f"Speedup       : {speedup:.2f}x\n")

print("\nSaved to benchmark_results.txt")
f.write(f"\nNote: ONNX speedup is most significant on GPU/edge hardware (Jetson, etc.).\n")
f.write(f"On CPU, PyTorch runtime is already well-optimized.\n")
f.write(f"ONNX export enables TensorRT acceleration for production edge deployment.\n")