#!/usr/bin/env python3
"""Export YOLOv5s / YOLOv8s to ONNX for OpenCV DNN inference."""
import sys
import argparse

def export_yolov5(weights, output, input_size=640):
    import torch
    model = torch.hub.load('ultralytics/yolov5', 'custom', path=weights, force_reload=False)
    model.eval()
    img = torch.zeros(1, 3, input_size, input_size)
    torch.onnx.export(model, img, output, opset_version=12,
                      input_names=['images'], output_names=['output'],
                      dynamic_axes={'images': {0: 'batch'}, 'output': {0: 'batch'}})
    print(f"YOLOv5 exported -> {output}")

def export_yolov8(weights, output, input_size=640):
    from ultralytics import YOLO
    model = YOLO(weights)
    model.export(format='onnx', imgsz=input_size, simplify=True)
    print(f"YOLOv8 exported -> {output}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', required=True, help='yolov5s.pt or yolov8s.pt')
    parser.add_argument('--output', default='models/yolov8s.onnx')
    parser.add_argument('--input-size', type=int, default=640)
    args = parser.parse_args()

    if 'yolov8' in args.weights.lower():
        export_yolov8(args.weights, args.output, args.input_size)
    else:
        export_yolov5(args.weights, args.output, args.input_size)