translate

app.py

@@ -62,8 +62,10 @@ transforms_val = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# --- 3. Función de Predicción ---
+
+# --- 3. Función de Predicción (CON LA TRADUCCIÓN A NOMBRES) ---
 def predict(image):
+    # La primera parte de la función no cambia
     if image is None:
         return None
     pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
@@ -73,7 +75,23 @@ def predict(image):
         probabilities = torch.nn.functional.softmax(output[0], dim=0)
 
     top5_prob, top5_catid = torch.topk(probabilities, 5)
-    confidences = {labels_map.get(str(cat_id.item()), f"ID {cat_id.item()}"): prob.item() for prob, cat_id in zip(top5_prob, top5_catid)}
+    
+    # --- ¡AQUÍ ESTÁ LA ÚNICA CORRECCIÓN QUE NECESITAS! ---
+    confidences = {}
+    for i in range(top5_prob.size(0)):
+        # Obtenemos el ID numérico predicho
+        species_id = top5_catid[i].item()
+        # Obtenemos la probabilidad
+        prob = top5_prob[i].item()
+        
+        # Usamos el mapa de etiquetas para "traducir" el ID a un nombre.
+        # Lo convertimos a string (str(species_id)) para que coincida con las claves del JSON.
+        species_name = labels_map.get(str(species_id), f"Desconocido (ID: {species_id})")
+        
+        # Añadimos al diccionario el NOMBRE como clave y la probabilidad como valor.
+        confidences[species_name] = prob
+    # --------------------------------------------------------
+        
     return confidences
 
 # --- 4. Crear la Interfaz de Gradio ---