import gradio as gr
from ultralytics import YOLO

import cv2

#info Gradio
#- https://www.gradio.app/docs/gradio/image
#- https://www.gradio.app/docs/gradio/file

# Función que recibe una imagen y la devuelve
def cargar_imagen(opcion, iou, conf, imagen, color):
    if opcion == "Tanques":
        model = YOLO(rf"\\HELIOS\myc\LabIA\Proyecto RINT 1\MyDatasets\Z. modelos\YOLO Open Imagen\prueba11 - tanks\prueba2\weights\best.pt") # 1 📁
    elif opcion == "Tanques o vehiculos":
        model = YOLO(rf"\\HELIOS\myc\LabIA\Proyecto RINT 1\MyDatasets\Z. modelos\YOLO Open Imagen\pruebas31-tankstrucks\modelo3\weights\best.pt") # 2 📁
    elif opcion == "Veh. españoles v1":
        model = YOLO(rf"\\HELIOS\myc\LabIA\Proyecto RINT 1\MyDatasets\Z. modelos\YOLO Open Imagen\prueba22\prueba102\weights\best.pt") #muchos 📁
    elif opcion == "Veh. españoles v2 (sin VERT)":
        model = YOLO(rf"\\HELIOS\myc\LabIA\Proyecto RINT 1\MyDatasets\Z. modelos\YOLO Open Imagen\pruebas41\modelov2(6B)+VEC\weights\best.pt") #muchos 📁 --> igual que arriba
    elif opcion == "Veh. militares españoles agrupados":
        model = YOLO(rf"\\HELIOS\myc\LabIA\Proyecto RINT 1\MyDatasets\Z. modelos\YOLO Open Imagen\pruebas51-9A\modelo9A\weights\best.pt") #1 📁 -> creo que no hay ninguno nuevo en 18.3
    else: #como si fueran "Vehiculos militares españoles"
        model = YOLO(rf"\\HELIOS\myc\LabIA\Proyecto RINT 1\MyDatasets\Z. modelos\YOLO Open Imagen\prueba22\prueba102\weights\best.pt")
        
    if color == "Blanco y Negro":
        from PIL import ImageOps
        imagen = ImageOps.grayscale(imagen) 
        
    results = model(imagen, iou=iou, conf=conf, visualize=False, augment=True, verbose = False)
    
    text = "" 
    for result in results[0].summary(): 
        r = result["name"] if not "VAMTAC" in result["name"] or opcion != "Veh. militares españoles agrupados"  else "VAMTAC+VERT"
        text = text + r + " - " + str(result["confidence"]) + "\n"
    im_rgb = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
    
    return im_rgb, text
    #results[0].verbose(), results[0].summary()

# Interfaz Gradio
iface = gr.Interface(
    fn=cargar_imagen,  # La función que ejecuta
    inputs=[gr.Image(type="pil"),
            gr.Slider(0, 1, value=0.7, label="iou"), gr.Slider(0, 1, value=0.3, label="confianza", scale = 0.1, min_width=20),             
            gr.Radio(choices=["Tanques", "Tanques o vehiculos", "Veh. militares españoles v1",
                              "Veh. militares españoles v2 (falta VERT)", "Veh. militares españoles agrupados"], label="Modelo de predicción"),  
            gr.Radio(choices=["Color", "Blanco y Negro"], label="Color"),
            ],  # Entrada de imagen:  gr.File(label="Elige el modelo")
    outputs=[gr.Image(type="pil", label="Predicción"), gr.Textbox(label="Confianza")],  # Salida de imagen
    title="Predicción Vehículos Militares",  # Título de la app
         # Descripción:
    description='''Sube una imagen y te devolverá los vehículos en la imagen.
    
    **Veh. militares españoles agrupados** incluye VERT pero no se muestra en la imagen.    
    
    *iou* es una medida para evitar el solapamiento de detecciones y *confianza* es el mínimo necesario para mostrarse.
    ''' 
)

# Ejecutar la aplicación
iface.launch(share=True)