91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
|
"""
|
||
|
|
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()
|