Initial commit: Wood knot detection model and GUI

export_rtdetr_oak.py

@@ -0,0 +1,90 @@
+"""
+Export trained RT-DETR model for OAK-D camera deployment.
+
+This script exports your trained model to OpenVINO format, which can then be
+converted to a blob for the OAK-D camera.
+
+Usage:
+    python export_rtdetr_oak.py --weights runs/rtdetr_training/training/weights/best.pt
+"""
+
+import argparse
+from pathlib import Path
+
+
+def export_for_oak(weights_path: Path, img_size: int = 640):
+    """
+    Export RT-DETR model for OAK-D deployment.
+    
+    Args:
+        weights_path: Path to trained .pt weights
+        img_size: Input image size (should match training)
+    """
+    from ultralytics import RTDETR
+    
+    if not weights_path.exists():
+        raise ValueError(f"Weights not found: {weights_path}")
+    
+    print(f"\n{'='*60}")
+    print(f"Exporting RT-DETR for OAK-D camera")
+    print(f"Weights: {weights_path}")
+    print(f"Image size: {img_size}x{img_size}")
+    print(f"{'='*60}\n")
+    
+    # Load model
+    model = RTDETR(weights_path)
+    
+    # Export to ONNX first (intermediate format)
+    print("Step 1/2: Exporting to ONNX...")
+    onnx_path = model.export(
+        format="onnx",
+        imgsz=img_size,
+        simplify=True,
+        opset=11,  # OAK-compatible opset
+    )
+    print(f"✓ ONNX exported: {onnx_path}")
+    
+    # Export to OpenVINO (for OAK)
+    print("\nStep 2/2: Exporting to OpenVINO...")
+    openvino_path = model.export(
+        format="openvino",
+        imgsz=img_size,
+        half=False,  # Use FP32 for better compatibility
+    )
+    print(f"✓ OpenVINO exported: {openvino_path}")
+    
+    print(f"\n{'='*60}")
+    print(f"✓ Export complete!")
+    print(f"\nNext steps:")
+    print(f"1. Test OpenVINO model: python test_openvino.py --model {openvino_path}")
+    print(f"2. Convert to blob for OAK:")
+    print(f"   Online: https://blobconverter.luxonis.com/")
+    print(f"   Or use: blobconverter --openvino-xml {openvino_path}/model.xml")
+    print(f"3. Deploy blob to OAK-D camera with DepthAI")
+    print(f"{'='*60}\n")
+    
+    return openvino_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export RT-DETR for OAK-D deployment")
+    parser.add_argument(
+        "--weights",
+        type=Path,
+        required=True,
+        help="Path to trained .pt weights file"
+    )
+    parser.add_argument(
+        "--img-size",
+        type=int,
+        default=640,
+        help="Input image size (should match training)"
+    )
+    
+    args = parser.parse_args()
+    
+    export_for_oak(args.weights, args.img_size)
+
+
+if __name__ == "__main__":
+    main()