You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.6 KiB
78 lines
2.6 KiB
#!/usr/bin/env python3
|
|
"""
|
|
Export YOLO model variants with different NMS settings for shiny icon debugging
|
|
"""
|
|
|
|
from ultralytics import YOLO
|
|
import os
|
|
|
|
def export_model_variants():
|
|
model_path = "./raw_models/best.pt"
|
|
output_dir = "./raw_models/exports"
|
|
|
|
# Create output directory
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
print(f"Loading model from: {model_path}")
|
|
model = YOLO(model_path)
|
|
|
|
# Export configurations to test
|
|
configs = [
|
|
{
|
|
"name": "no_nms",
|
|
"nms": False,
|
|
"simplify": True,
|
|
"description": "Raw model output without NMS - for debugging shiny detection"
|
|
},
|
|
{
|
|
"name": "nms_relaxed",
|
|
"nms": True,
|
|
"max_det": 500, # Increase from default 300
|
|
"conf": 0.1, # Lower confidence threshold
|
|
"simplify": True,
|
|
"description": "NMS with more detections and lower confidence"
|
|
},
|
|
{
|
|
"name": "nms_very_relaxed",
|
|
"nms": True,
|
|
"max_det": 1000, # Even more detections
|
|
"conf": 0.05, # Very low confidence
|
|
"simplify": True,
|
|
"description": "NMS with maximum detections for rare classes"
|
|
}
|
|
]
|
|
|
|
for config in configs:
|
|
try:
|
|
print(f"\n🚀 Exporting {config['name']}: {config['description']}")
|
|
|
|
# Extract export parameters
|
|
export_params = {k: v for k, v in config.items()
|
|
if k not in ['name', 'description']}
|
|
|
|
# Export model
|
|
exported_path = model.export(
|
|
format='onnx',
|
|
**export_params
|
|
)
|
|
|
|
# Move to organized location
|
|
output_file = os.path.join(output_dir, f"best_{config['name']}.onnx")
|
|
if os.path.exists(exported_path):
|
|
os.rename(exported_path, output_file)
|
|
print(f"✅ Exported: {output_file}")
|
|
else:
|
|
print(f"❌ Export failed for {config['name']}")
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error exporting {config['name']}: {e}")
|
|
|
|
print(f"\n📁 All exports saved to: {output_dir}")
|
|
print("\n📋 Summary:")
|
|
print("- best_no_nms.onnx: Raw 8400x99 output for debugging")
|
|
print("- best_nms_relaxed.onnx: NMS with 500 max detections")
|
|
print("- best_nms_very_relaxed.onnx: NMS with 1000 max detections")
|
|
print("\nNext: Copy desired model to app/src/main/assets/ as best.onnx")
|
|
|
|
if __name__ == "__main__":
|
|
export_model_variants()
|