convert_to_onnx.py

#!/usr/bin/env python3
"""
Convert super-gradients model to ONNX format for lightweight inference.

This script requires super-gradients to be installed. Run it ONCE to convert the model,
then you can use the ONNX model without needing super-gradients.

Run this on your host machine or in a temporary container with super-gradients installed:
    pip install super-gradients
    python3 convert_to_onnx.py
"""

import os
import torch
import yaml
from pathlib import Path
from super_gradients.training import models

def convert_model():
    script_dir = Path(__file__).parent
    checkpoints_dir = script_dir / 'models' / 'checkpoints'
    model_dirs = list(checkpoints_dir.glob('*_custom'))
    
    if not model_dirs:
        print(f"ERROR: No model directories found in {checkpoints_dir}")
        return False
    
    model_dir = sorted(model_dirs, key=lambda x: x.stat().st_mtime)[-1]
    model_path = model_dir / 'latest_model.pth'
    model_name = model_dir.name.replace('_custom', '')
    config_path = script_dir / 'configs' / 'ontology.yaml'
    output_path = script_dir / 'models' / 'model.onnx'
    
    if not model_path.exists():
        print(f"ERROR: Model file not found: {model_path}")
        return False
    
    if not config_path.exists():
        print(f"ERROR: Config file not found: {config_path}")
        return False
    
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    num_classes = len(config['ontology']['classes']) + 1
    class_names = [cls['name'] for cls in config['ontology']['classes']]
    print(f"Classes: {', '.join(class_names)}")
    
    model = models.get(model_name, num_classes=num_classes)
    checkpoint = torch.load(model_path, map_location='cpu')
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    dummy_input = torch.randn(1, 3, 480, 640)
    
    with torch.no_grad():
        output = model(dummy_input)
        if isinstance(output, (list, tuple)):
            output = output[0]
    
    with torch.no_grad():
        torch.onnx.export(
            model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=11,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size', 2: 'height', 3: 'width'},
                'output': {0: 'batch_size', 2: 'height', 3: 'width'}
            },
            verbose=False,
            dynamo=False  # Use legacy exporter (no onnxscript needed)
        )
    
    print(f"Model exported to: {output_path}")
    return True

if __name__ == '__main__':
    try:
        success = convert_model()
        if not success:
            exit(1)
    except Exception as e:
        print(f"ERROR: {e}")
        exit(1)