108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
import os
|
|
import sys
|
|
from ultralytics import YOLO
|
|
import yaml
|
|
import shutil
|
|
from pathlib import Path
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
|
|
dir = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
def verify_images(dataset_path):
|
|
"""Überprüfe alle Bilder auf Lesbarkeit und korrekte Dimensionen."""
|
|
print("\nÜberprüfe Bilder...")
|
|
corrupted = []
|
|
for split in ['train', 'val']:
|
|
img_dir = dataset_path / split / 'images'
|
|
if not img_dir.exists():
|
|
continue
|
|
|
|
for img_path in img_dir.glob('*'):
|
|
try:
|
|
img = cv2.imread(str(img_path))
|
|
if img is None:
|
|
raise ValueError("Bild konnte nicht gelesen werden")
|
|
if img.shape[0] < 10 or img.shape[1] < 10:
|
|
raise ValueError("Bild ist zu klein")
|
|
except Exception as e:
|
|
print(f"Warnung: Problem mit {img_path}: {str(e)}")
|
|
corrupted.append(img_path)
|
|
|
|
if corrupted:
|
|
print(f"\nGefundene problematische Bilder: {len(corrupted)}")
|
|
return False
|
|
return True
|
|
|
|
def create_yolo_dataset():
|
|
"""Erstelle die YOLO-Dateistruktur und konvertiere die Annotationen."""
|
|
try:
|
|
# Erstelle Verzeichnisse
|
|
dataset_path = Path('yolo_dataset')
|
|
if not dataset_path.exists():
|
|
print("Fehler: YOLO Dataset nicht gefunden! Bitte führen Sie zuerst convert_annotations.py aus.")
|
|
sys.exit(1)
|
|
|
|
# Überprüfe Bilder
|
|
if not verify_images(dataset_path):
|
|
print("Warnung: Einige Bilder könnten problematisch sein!")
|
|
|
|
# Erstelle dataset.yaml
|
|
dataset_yaml = {
|
|
'path': str(dataset_path.absolute()),
|
|
'train': 'train/images',
|
|
'val': 'val/images',
|
|
'names': {
|
|
0: 'NAO-Roboter'
|
|
},
|
|
'nc': 1 # Anzahl der Klassen
|
|
}
|
|
|
|
yaml_path = dataset_path / 'dataset.yaml'
|
|
with open(yaml_path, 'w') as f:
|
|
yaml.dump(dataset_yaml, f)
|
|
|
|
return str(yaml_path)
|
|
except Exception as e:
|
|
print(f"Fehler beim Erstellen des Datasets: {str(e)}")
|
|
sys.exit(1)
|
|
|
|
def train_yolo():
|
|
"""Trainiere das YOLOv8 Modell."""
|
|
# Erstelle und konfiguriere das Dataset
|
|
dataset_yaml = create_yolo_dataset()
|
|
|
|
# Initialisiere das Modell
|
|
model = YOLO('yolov8n.pt') # Verwende das kleine YOLOv8 Modell
|
|
|
|
# Trainiere das Modell
|
|
results = model.train(
|
|
data=dataset_yaml,
|
|
epochs=50,
|
|
imgsz=320, # Bildgröße
|
|
batch=16, # Batch-Größe
|
|
device=0, # Verwende CPU (oder 'cuda' für GPU)
|
|
patience=5, # Early Stopping
|
|
save=True, # Speichere die besten Gewichte
|
|
project='yolo_training', # Projektname
|
|
name='NAO_detector', # Experimentname
|
|
)
|
|
|
|
# Exportiere das Modell für den NAO
|
|
print("\nExportiere Modell in verschiedene Formate...")
|
|
|
|
# ONNX Format (gut für eingebettete Systeme)
|
|
model.export(format='onnx', imgsz=320)
|
|
|
|
# OpenVINO Format (optimiert für Intel Hardware)
|
|
model.export(format='openvino', imgsz=320)
|
|
|
|
# tflite Format
|
|
model.export(format="tflite",imgsz=320,)
|
|
|
|
print("\nTraining abgeschlossen. Die Modelle wurden im 'yolo_training/NAO_detector' Ordner gespeichert.")
|
|
|
|
if __name__ == '__main__':
|
|
train_yolo() |