100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
# Add current directory to path
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.append(current_dir)
|
|
|
|
# Import the conversion function from detection_openvino.py
|
|
from detection_openvino import convert_yolo_to_openvino
|
|
|
|
def main():
|
|
"""
|
|
Convert yolo11n.pt model to OpenVINO IR format.
|
|
Usage: python convert_yolo11n.py
|
|
"""
|
|
print("\n" + "="*80)
|
|
print("YOLO11n Model Converter - PyTorch to OpenVINO IR")
|
|
print("="*80)
|
|
# Check if the model exists
|
|
model_path = Path("yolo11n.pt")
|
|
if not model_path.exists():
|
|
print(f"❌ Error: Model file {model_path} not found!")
|
|
print(f" Please ensure '{model_path}' is in the current directory.")
|
|
return
|
|
|
|
print(f"✅ Found model: {model_path}")
|
|
|
|
# Check for OpenVINO and other dependencies
|
|
try:
|
|
import openvino as ov
|
|
print(f"✅ OpenVINO version: {ov.__version__}")
|
|
except ImportError:
|
|
print("⚠️ OpenVINO not installed. Installing now...")
|
|
os.system('pip install --quiet "openvino>=2024.0.0"')
|
|
import openvino as ov
|
|
print(f"✅ OpenVINO installed: {ov.__version__}")
|
|
|
|
try:
|
|
from ultralytics import YOLO
|
|
except ImportError:
|
|
print("⚠️ Ultralytics not installed. Installing now...")
|
|
os.system('pip install --quiet "ultralytics>=8.0.0"')
|
|
from ultralytics import YOLO
|
|
print("✅ Ultralytics installed")
|
|
|
|
# Create destination directory for the models
|
|
openvino_dir = Path("openvino_models")
|
|
if not openvino_dir.exists():
|
|
openvino_dir.mkdir(exist_ok=True)
|
|
print(f"✅ Created directory: {openvino_dir}")
|
|
|
|
try:
|
|
# Convert model to OpenVINO IR format
|
|
print("\n📦 Converting model to OpenVINO IR format...")
|
|
start_time = time.time()
|
|
output_path = convert_yolo_to_openvino("yolo11n", half=True)
|
|
conversion_time = time.time() - start_time
|
|
|
|
print(f"✅ Conversion completed in {conversion_time:.2f} seconds!")
|
|
print(f"✅ Output model: {output_path}")
|
|
|
|
# Verify output files
|
|
if output_path and Path(output_path).exists():
|
|
xml_path = Path(output_path)
|
|
bin_path = xml_path.with_suffix('.bin')
|
|
xml_size = xml_path.stat().st_size / (1024 * 1024) # in MB
|
|
bin_size = bin_path.stat().st_size / (1024 * 1024) # in MB
|
|
|
|
print(f"✅ XML file: {xml_path} ({xml_size:.2f} MB)")
|
|
print(f"✅ BIN file: {bin_path} ({bin_size:.2f} MB)")
|
|
|
|
# Copy to openvino_models directory for easier access by the Qt app
|
|
dst_xml = openvino_dir / xml_path.name
|
|
dst_bin = openvino_dir / bin_path.name
|
|
|
|
shutil.copy2(xml_path, dst_xml)
|
|
shutil.copy2(bin_path, dst_bin)
|
|
|
|
print(f"✅ Copied models to: {openvino_dir}")
|
|
print("\n🚀 Model conversion and setup complete!")
|
|
print("\n📋 Instructions:")
|
|
print(f" 1. The model files are available at: {openvino_dir}")
|
|
print(" 2. In the Qt app, you can now select this model from the dropdown")
|
|
print(" 3. Use the device selection dropdown to choose between CPU and GPU")
|
|
else:
|
|
print("❌ Failed to verify output files.")
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error converting model: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|