|
|
|
|
|
""" |
|
|
🔬 Clasificador de Insectos Polinizadores - Versión de Producción |
|
|
Precisión alcanzada: 92.07% |
|
|
Modelo: YOLOv8 Nano |
|
|
""" |
|
|
|
|
|
from ultralytics import YOLO |
|
|
import sys |
|
|
import os |
|
|
from pathlib import Path |
|
|
|
|
|
class PollinatorClassifier: |
|
|
def __init__(self, model_path="pollinator_results/nano_quick/weights/best.pt"): |
|
|
"""Inicializar el clasificador""" |
|
|
try: |
|
|
self.model = YOLO(model_path) |
|
|
self.classes = [ |
|
|
'Acmaeodera Flavomarginata', 'Acromyrmex Octospinosus', |
|
|
'Adelpha Basiloides', 'Adelpha Iphicleola', 'Aedes Aegypti', |
|
|
'Agrius Cingulata', 'Anaea Aidea', 'Anartia fatima', |
|
|
'Anartia jatrophae', 'Anoplolepis Gracilipes' |
|
|
] |
|
|
print("🔬 Clasificador de Insectos Polinizadores v1.0") |
|
|
print(f"✅ Modelo cargado con 92.07% de precisión") |
|
|
print(f"🏷️ {len(self.classes)} clases disponibles") |
|
|
|
|
|
except Exception as e: |
|
|
print(f"❌ Error cargando modelo: {e}") |
|
|
sys.exit(1) |
|
|
|
|
|
def classify(self, image_path): |
|
|
"""Clasificar una imagen de insecto""" |
|
|
|
|
|
if not os.path.exists(image_path): |
|
|
print(f"❌ Imagen no encontrada: {image_path}") |
|
|
return None |
|
|
|
|
|
|
|
|
results = self.model(image_path, verbose=False) |
|
|
probs = results[0].probs |
|
|
|
|
|
|
|
|
top_class_idx = probs.top1 |
|
|
confidence = probs.top1conf.item() * 100 |
|
|
predicted_class = self.classes[top_class_idx] |
|
|
|
|
|
print(f"\n🔍 Imagen: {os.path.basename(image_path)}") |
|
|
print(f"🎯 Predicción: {predicted_class}") |
|
|
print(f"📊 Confianza: {confidence:.1f}%") |
|
|
|
|
|
|
|
|
print(f"\n📋 Top 3 predicciones:") |
|
|
for i in range(min(3, len(probs.top5))): |
|
|
idx = probs.top5[i] |
|
|
conf = probs.top5conf[i].item() * 100 |
|
|
class_name = self.classes[idx] |
|
|
emoji = "🥇" if i == 0 else "🥈" if i == 1 else "🥉" |
|
|
print(f" {emoji} {class_name}: {conf:.1f}%") |
|
|
|
|
|
return predicted_class, confidence |
|
|
|
|
|
def classify_batch(self, folder_path): |
|
|
"""Clasificar múltiples imágenes en una carpeta""" |
|
|
|
|
|
folder = Path(folder_path) |
|
|
if not folder.exists(): |
|
|
print(f"❌ Carpeta no encontrada: {folder_path}") |
|
|
return |
|
|
|
|
|
|
|
|
image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.JPG', '*.JPEG', '*.PNG'] |
|
|
images = [] |
|
|
for ext in image_extensions: |
|
|
images.extend(list(folder.glob(ext))) |
|
|
|
|
|
if not images: |
|
|
print("❌ No se encontraron imágenes") |
|
|
return |
|
|
|
|
|
print(f"🔍 Clasificando {len(images)} imágenes...") |
|
|
print("-" * 60) |
|
|
|
|
|
results = [] |
|
|
for img_path in images: |
|
|
pred_class, confidence = self.classify(str(img_path)) |
|
|
if pred_class: |
|
|
results.append({ |
|
|
'imagen': img_path.name, |
|
|
'prediccion': pred_class, |
|
|
'confianza': confidence |
|
|
}) |
|
|
|
|
|
return results |
|
|
|
|
|
def main(): |
|
|
"""Función principal""" |
|
|
classifier = PollinatorClassifier() |
|
|
|
|
|
if len(sys.argv) < 2: |
|
|
|
|
|
print("\n🎯 MODO INTERACTIVO") |
|
|
print("Opciones:") |
|
|
print("1. Clasificar una imagen") |
|
|
print("2. Clasificar carpeta de imágenes") |
|
|
|
|
|
choice = input("\nSelecciona opción (1 o 2): ") |
|
|
|
|
|
if choice == "1": |
|
|
image_path = input("Ruta de la imagen: ") |
|
|
classifier.classify(image_path) |
|
|
elif choice == "2": |
|
|
folder_path = input("Ruta de la carpeta: ") |
|
|
classifier.classify_batch(folder_path) |
|
|
else: |
|
|
print("Opción inválida") |
|
|
else: |
|
|
|
|
|
path = sys.argv[1] |
|
|
if os.path.isfile(path): |
|
|
classifier.classify(path) |
|
|
elif os.path.isdir(path): |
|
|
classifier.classify_batch(path) |
|
|
else: |
|
|
print(f"❌ Ruta inválida: {path}") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
