pollinator-classifier / utils /fix_dataset_structure.py
leonelgv's picture
Add Dataset preparation utilities
43940ae verified
#!/usr/bin/env python3
"""
Script simplificado de entrenamiento YOLOv8 clasificación
"""
import torch
from ultralytics import YOLO
import os
def main():
print("🚀 ENTRENAMIENTO YOLO CLASIFICACIÓN")
print("=" * 50)
# Verificar CUDA
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"💻 Dispositivo: {device}")
print(f"🎯 Objetivo: >90% precisión")
# Dataset path
dataset_path = "/home/leonel/sistema_polinizador/Dataset/Classification_YOLO"
if not os.path.exists(dataset_path):
print(f"❌ Dataset no encontrado: {dataset_path}")
print("💡 Ejecuta primero: python fix_structure.py")
return
# Configuraciones de entrenamiento
configs = [
{
"name": "nano_quick",
"model": "yolov8n-cls.pt",
"epochs": 30,
"imgsz": 224,
"batch": 32
},
{
"name": "small_balanced",
"model": "yolov8s-cls.pt",
"epochs": 60,
"imgsz": 256,
"batch": 16
},
{
"name": "medium_accurate",
"model": "yolov8m-cls.pt",
"epochs": 100,
"imgsz": 320,
"batch": 8
}
]
best_accuracy = 0
best_model = None
for i, config in enumerate(configs, 1):
print(f"\n{i}️⃣ MODELO: {config['name']}")
print("=" * 40)
try:
# Cargar modelo
model = YOLO(config["model"])
print(f"📥 Modelo cargado: {config['model']}")
# Entrenar
print(f"⏰ Iniciando entrenamiento...")
results = model.train(
data=dataset_path,
epochs=config["epochs"],
imgsz=config["imgsz"],
batch=config["batch"],
device=device,
project="pollinator_final",
name=config["name"],
patience=20,
save=True,
verbose=False,
plots=True
)
# Evaluar
print(f"📊 Evaluando en test set...")
test_results = model.val(split='test')
accuracy = float(test_results.top1) * 100
print(f"✅ Entrenamiento completado")
print(f"🎯 Precisión: {accuracy:.2f}%")
if accuracy > best_accuracy:
best_accuracy = accuracy
best_model = f"pollinator_final/{config['name']}/weights/best.pt"
# Verificar objetivo
if accuracy >= 90:
print(f"🎉 ¡OBJETIVO ALCANZADO! {accuracy:.2f}% ≥ 90%")
break
else:
print(f"⚠️ Faltan {90-accuracy:.2f}% para objetivo")
except Exception as e:
print(f"❌ Error: {e}")
continue
# Resultados finales
print(f"\n" + "=" * 50)
print("📊 RESULTADOS FINALES")
print("=" * 50)
print(f"🏆 Mejor precisión: {best_accuracy:.2f}%")
if best_accuracy >= 90:
print(f"✅ OBJETIVO ALCANZADO!")
else:
print(f"❌ Objetivo no alcanzado")
print(f"💡 Recomendación: Entrenar modelo YOLOv8l o YOLOv8x")
if best_model:
print(f"📁 Mejor modelo: {best_model}")
# Crear script de predicción simple
pred_script = f'''#!/usr/bin/env python3
from ultralytics import YOLO
# Cargar modelo entrenado
model = YOLO('{best_model}')
# Función para clasificar
def classify_insect(image_path):
results = model(image_path, verbose=False)
probs = results[0].probs
classes = [
'Acmaeodera Flavomarginata', 'Acromyrmex Octospinosus',
'Adelpha Basiloides', 'Adelpha Iphicleola', 'Aedes Aegypti',
'Agrius Cingulata', 'Anaea Aidea', 'Anartia fatima',
'Anartia jatrophae', 'Anoplolepis Gracilipes'
]
top_class = classes[probs.top1]
confidence = probs.top1conf.item() * 100
print(f"🔍 Predicción: {{top_class}}")
print(f"📊 Confianza: {{confidence:.1f}}%")
return top_class, confidence
# Ejemplo de uso
if __name__ == "__main__":
image_path = input("Ruta de imagen: ")
if image_path:
classify_insect(image_path)
'''
with open('predict_final.py', 'w') as f:
f.write(pred_script)
print(f"✅ Script de predicción: predict_final.py")
return best_accuracy
if __name__ == "__main__":
final_accuracy = main()
print(f"\n🎯 Entrenamiento completado. Precisión final: {final_accuracy:.2f}%")