app.py

import gradio as gr
import torch
import numpy as np
from torchvision import transforms
from PIL import Image
from timm import create_model

# Definir el diccionario de mapeo de clases a índices
class_to_idx = {'Monkeypox': 0, 'Melanoma': 1, 'Herpes': 2, 'Sarampion': 3, 'Varicela': 4}

# Transformación de datos
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# Función para cargar y preprocesar una imagen
def load_image(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)  # Añadir dimensión del batch
    return image

# Cargar el modelo
model_name = 'vit_base_patch16_224'
pretrained = True
num_classes = len(class_to_idx)
model = create_model(model_name, pretrained=pretrained, num_classes=num_classes)
model.load_state_dict(torch.load('ARTmodelo5ns_vit_weights_epoch6.pth', map_location='cpu', weights_only=True))
model.eval()

# Definir la función de predicción
def predict_image(img):
    # Convertir la imagen a PIL.Image si es un numpy array
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img)
    
    # Convertir la imagen a tensor y añadir dimensión del batch
    img_tensor = transform(img).unsqueeze(0)
    
    # Realizar la predicción
    with torch.no_grad():
        output = model(img_tensor)
        _, predicted = torch.max(output, 1)

   
    predicted_label = list(class_to_idx.keys())[predicted.item()]
    
    return predicted_label

# Crear la interfaz de Gradio
iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="filepath", label="Sube una imagen"),
    outputs=gr.Label(label="Predicción"),
    title="Clasificación de Imágenes de Lesiones Cutáneas",
    description="Carga una imagen de una lesión cutánea para obtener una predicción."
)

# Lanzar la interfaz de Gradio
iface.launch()