YOLO

"""YOLO 모델을 ONNX 형식으로 변환하는 스크립트."""

from pathlib import Path

from ultralytics import YOLO


def export_to_onnx(weights_path: str, opset: int = 12, simplify: bool = True, dynamic: bool = True) -> str:
    """
    YOLO 모델을 ONNX 형식으로 변환합니다.

    Args:
        weights_path: .pt 모델 파일 경로
        opset: ONNX opset 버전 (기본값: 12)
        simplify: onnx-simplifier 적용 여부 (기본값: True)
        dynamic: dynamic batch size 지원 여부 (기본값: True)

    Returns:
        변환된 ONNX 파일 경로
    """
    model = YOLO(weights_path)
    export_path = model.export(format="onnx", opset=opset, simplify=simplify, dynamic=dynamic)
    return export_path


def main():
    weights_dir = Path(__file__).parent / "weights"

    models = [
        weights_dir / "sh_weld_defect-20251113.pt",
        weights_dir / "sh_weld_zone-20251222.pt",
    ]

    for model_path in models:
        if not model_path.exists():
            print(f"모델 파일을 찾을 수 없습니다: {model_path}")
            continue

        print(f"변환 중: {model_path.name}")
        onnx_path = export_to_onnx(str(model_path))
        print(f"완료: {onnx_path}")


if __name__ == "__main__":
    main()