Add application file and dependencies

app.py

@@ -5,22 +5,22 @@ 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}
+# Define class to index mapping dictionary
+class_to_idx = {'Monkeypox': 0, 'Melanoma': 1, 'Herpes': 2, 'Measles': 3, 'Chickenpox': 4}
 
-# Transformación de datos
+# Data transformation
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
 ])
 
-# Función para cargar y preprocesar una imagen
+# Function to load and preprocess an image
 def load_image(image_path):
     image = Image.open(image_path).convert('RGB')
-    image = transform(image).unsqueeze(0)  # Añadir dimensión del batch
+    image = transform(image).unsqueeze(0)  # Add batch dimension
     return image
 
-# Cargar el modelo
+# Load the model
 model_name = 'vit_base_patch16_224'
 pretrained = True
 num_classes = len(class_to_idx)
@@ -28,33 +28,29 @@ 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)
+# Define the prediction function
+def predict_image(image_path):
+    # Load and transform the image from the file path
+    img_tensor = load_image(image_path)
     
-    # Convertir la imagen a tensor y añadir dimensión del batch
-    img_tensor = transform(img).unsqueeze(0)
-    
-    # Realizar la predicción
+    # Perform the prediction
     with torch.no_grad():
         output = model(img_tensor)
         _, predicted = torch.max(output, 1)
 
-   
+    # Get the predicted label
     predicted_label = list(class_to_idx.keys())[predicted.item()]
     
     return predicted_label
 
-# Crear la interfaz de Gradio
+# Create the Gradio interface
 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."
+    inputs=gr.Image(type="filepath", label="Upload an image"),
+    outputs=gr.Label(label="Prediction"),
+    title="Skin Lesion Image Classification",
+    description="Upload an image of a skin lesion to get a prediction."
 )
 
-# Lanzar la interfaz de Gradio
+# Launch the Gradio interface
 iface.launch()