#!/usr/bin/env python3
"""
Convert trained models for OAK-D deployment.

Supports conversion to ONNX and OpenVINO formats for edge deployment.

Usage:
    python convert_for_deployment.py --model runs/training/weights/best.pt --output oak_d_deployment
"""

import argparse
import shutil
from pathlib import Path


def convert_rfdetr_for_oak(model_path: Path, output_dir: Path, img_size: int = 640):
    """Convert RF-DETR model for OAK-D deployment."""
    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        from export_rtdetr_oak import export_rfdetr_onnx

        # Export to ONNX
        onnx_path = output_dir / "model.onnx"
        export_rfdetr_onnx(str(model_path), str(onnx_path), img_size)

        # Convert ONNX to OpenVINO
        import subprocess
        xml_path = output_dir / "model.xml"
        bin_path = output_dir / "model.bin"

        cmd = [
            "mo", "--input_model", str(onnx_path),
            "--output_dir", str(output_dir),
            "--model_name", "model",
            "--input_shape", f"[1,3,{img_size},{img_size}]",
            "--data_type", "FP16"
        ]

        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            return f"❌ OpenVINO conversion failed: {result.stderr}"

        return f"✓ RF-DETR exported for OAK-D!\n📁 Output: {output_dir}\n🔗 Next: Convert ONNX → RVC using HubAI (online) or ModelConverter (offline)."

    except Exception as e:
        return f"❌ RF-DETR conversion failed: {e}"


def convert_rtdetr_for_oak(model_path: Path, output_dir: Path, img_size: int = 640):
    """Convert RT-DETR model for OAK-D deployment."""
    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        from export_rtdetr_oak import export_rtdetr_openvino

        # Export to OpenVINO directly
        xml_path, bin_path = export_rtdetr_openvino(str(model_path), str(output_dir), img_size)

        return f"✓ RT-DETR exported for OAK-D!\n📁 Output: {output_dir}\n🔗 Next: Convert .xml/.bin → blob using blobconverter.luxonis.com"

    except Exception as e:
        return f"❌ RT-DETR conversion failed: {e}"


def convert_yolov6_for_oak(model_path: Path, output_dir: Path, img_size: int = 640):
    """Convert YOLOv6 model for OAK-D deployment."""
    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        from export_onnx import export_yolov6_onnx

        # Export to ONNX
        onnx_path = export_yolov6_onnx(str(model_path), str(output_dir), img_size)

        # Convert ONNX to OpenVINO
        import subprocess
        xml_path = output_dir / "model.xml"
        bin_path = output_dir / "model.bin"

        cmd = [
            "mo", "--input_model", str(onnx_path),
            "--output_dir", str(output_dir),
            "--model_name", "model",
            "--input_shape", f"[1,3,{img_size},{img_size}]",
            "--data_type", "FP16",
            "--reverse_input_channels"
        ]

        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            return f"❌ OpenVINO conversion failed: {result.stderr}"

        return f"✓ YOLOv6 exported for OAK-D!\n📁 Output: {output_dir}\n🔗 Next: Convert .xml/.bin → blob using blobconverter.luxonis.com"

    except Exception as e:
        return f"❌ YOLOv6 conversion failed: {e}"


def convert_yolox_for_oak(model_path: Path, output_dir: Path, img_size: int = 640):
    """Convert YOLOX model for OAK-D deployment."""
    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        from export_onnx import export_yolox_onnx

        # Export to ONNX
        onnx_path = export_yolox_onnx(str(model_path), str(output_dir), img_size)

        # Convert ONNX to OpenVINO
        import subprocess
        xml_path = output_dir / "model.xml"
        bin_path = output_dir / "model.bin"

        cmd = [
            "mo", "--input_model", str(onnx_path),
            "--output_dir", str(output_dir),
            "--model_name", "model",
            "--input_shape", f"[1,3,{img_size},{img_size}]",
            "--data_type", "FP16",
            "--reverse_input_channels"
        ]

        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            return f"❌ OpenVINO conversion failed: {result.stderr}"

        return f"✓ YOLOX exported for OAK-D!\n📁 Output: {output_dir}\n🔗 Next: Convert .xml/.bin → blob using blobconverter.luxonis.com"

    except Exception as e:
        return f"❌ YOLOX conversion failed: {e}"


def detect_model_type(model_path: Path) -> str:
    """Detect model type from file path or contents."""
    path_str = str(model_path).lower()

    # Check path patterns
    if 'rf-detr' in path_str or 'rfdetr' in path_str:
        return 'rf-detr'
    elif 'rt-detr' in path_str or 'rtdetr' in path_str:
        return 'rt-detr'
    elif 'yolov6' in path_str:
        return 'yolov6'
    elif 'yolox' in path_str:
        return 'yolox'

    # Try to detect from file contents
    try:
        import torch
        checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)

        if 'state_dict' in checkpoint:
            state_dict = checkpoint['state_dict']

            # Check for RT-DETR patterns
            if any('rtdetr' in key.lower() for key in state_dict.keys()):
                return 'rt-detr'

            # Check for RF-DETR patterns
            if any('rf_detr' in key.lower() for key in state_dict.keys()):
                return 'rf-detr'

            # Check for YOLO patterns
            if any('yolo' in key.lower() for key in state_dict.keys()):
                if any('v6' in key.lower() for key in state_dict.keys()):
                    return 'yolov6'
                else:
                    return 'yolox'

    except Exception:
        pass

    # Default fallback
    return 'yolox'


def main():
    parser = argparse.ArgumentParser(description="Convert trained models for OAK-D deployment")
    parser.add_argument(
        '--model',
        type=Path,
        required=True,
        help='Path to trained model weights (.pt file)'
    )
    parser.add_argument(
        '--output',
        type=Path,
        default=Path('oak_d_deployment'),
        help='Output directory for converted models'
    )
    parser.add_argument(
        '--img-size',
        type=int,
        default=640,
        choices=[320, 416, 512, 640, 800, 1024],
        help='Input image size for the model'
    )
    parser.add_argument(
        '--framework',
        choices=['auto', 'rf-detr', 'rt-detr', 'yolov6', 'yolox'],
        default='auto',
        help='Model framework (auto-detect if not specified)'
    )

    args = parser.parse_args()

    # Auto-detect framework if not specified
    if args.framework == 'auto':
        framework = detect_model_type(args.model)
        print(f"Auto-detected framework: {framework}")
    else:
        framework = args.framework

    print(f"Converting {framework.upper()} model for OAK-D deployment...")
    print(f"Model: {args.model}")
    print(f"Output: {args.output}")
    print(f"Image size: {args.img_size}")

    # Convert based on framework
    if framework == 'rf-detr':
        result = convert_rfdetr_for_oak(args.model, args.output, args.img_size)
    elif framework == 'rt-detr':
        result = convert_rtdetr_for_oak(args.model, args.output, args.img_size)
    elif framework == 'yolov6':
        result = convert_yolov6_for_oak(args.model, args.output, args.img_size)
    elif framework == 'yolox':
        result = convert_yolox_for_oak(args.model, args.output, args.img_size)
    else:
        result = f"❌ Unsupported framework: {framework}"

    print(result)

    if "✓" in result:
        print("\n📋 Next steps:")
        print("1. Test OpenVINO model (optional):")
        print("   python -c \"from openvino.runtime import Core; core = Core(); model = core.read_model('model.xml'); print('✓ Model loaded')\"")
        print("2. Convert to RVC compiled format:")
        print("   - Online: HubAI conversion (fastest setup)")
        print("   - Offline: ModelConverter (requires Docker)")
        print("   - Docs: https://docs.luxonis.com/software-v3/ai-inference/conversion/")
        print("3. Deploy to OAK-D using DepthAI Python API")


if __name__ == "__main__":
    main()</content>
<parameter name="filePath">/home/dillon/_code/saw_mill_knot_detection/convert_for_deployment.py