Update app.py

app.py

@@ -21,6 +21,7 @@ if not os.path.exists(extract_dir):
         zip_ref.extractall(extract_dir)
 
 model_tf = tf.saved_model.load(extract_dir)
+TF_NUM_CLASSES = 7  # asumimos que son las mismas que CLASSES
 
 # Función helper para inferencia TensorFlow
 def predict_tf(img: Image.Image):
@@ -35,17 +36,17 @@ def predict_tf(img: Image.Image):
         infer = model_tf.signatures["serving_default"]
         output = infer(img_tf)
         pred = list(output.values())[0].numpy()[0]
-        probs = tf.nn.softmax(pred).numpy()
+        probs = tf.nn.softmax(pred[:TF_NUM_CLASSES]).numpy()
         return probs
     except Exception as e:
         print(f"Error en predict_tf: {e}")
-        return np.zeros(2)
+        return np.zeros(TF_NUM_CLASSES)
 
+# --- Cargar modelos ---
 MODEL_NAME = "ahishamm/vit-base-HAM-10000-sharpened-patch-32"
 feature_extractor = ViTImageProcessor.from_pretrained(MODEL_NAME)
 model_vit = ViTForImageClassification.from_pretrained(MODEL_NAME)
 model_vit.eval()
-
 model_malignancy = load_learner("ada_learn_malben.pkl")
 model_norm2000 = load_learner("ada_learn_skin_norm2000.pkl")
 
@@ -64,6 +65,9 @@ RISK_LEVELS = {
     6: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1}
 }
 
+MALIGNANT_INDICES = [0, 1, 4]  # clases de riesgo alto/crítico
+
+
 def analizar_lesion_combined(img):
     try:
         img_fastai = PILImage.create(img)
@@ -74,7 +78,7 @@ def analizar_lesion_combined(img):
         pred_idx_vit = int(np.argmax(probs_vit))
         pred_class_vit = CLASSES[pred_idx_vit]
         confidence_vit = probs_vit[pred_idx_vit]
-    except:
+    except Exception as e:
         pred_class_vit = "Error"
         confidence_vit = 0.0
         probs_vit = np.zeros(len(CLASSES))
@@ -92,13 +96,12 @@ def analizar_lesion_combined(img):
 
     try:
         probs_tf = predict_tf(img)
-        if len(probs_tf) == 2:
-            benign_prob, malignant_prob = probs_tf
-            pred_class_tf = "Maligno" if malignant_prob > benign_prob else "Benigno"
-            confidence_tf = max(probs_tf)
+        pred_idx_tf = int(np.argmax(probs_tf))
+        confidence_tf = probs_tf[pred_idx_tf]
+        if pred_idx_tf < len(CLASSES):
+            pred_class_tf = "Maligno" if pred_idx_tf in MALIGNANT_INDICES else "Benigno"
         else:
-            pred_class_tf = "Modelo no binario"
-            confidence_tf = 0.0
+            pred_class_tf = f"Desconocido"
     except:
         pred_class_tf = "Error"
         confidence_tf = 0.0
@@ -115,13 +118,12 @@ def analizar_lesion_combined(img):
     buf = io.BytesIO()
     plt.savefig(buf, format="png")
     plt.close(fig)
-    img_bytes = buf.getvalue()
-    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
+    img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
     html_chart = f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
 
     informe = f"""
     <div style="font-family:sans-serif; max-width:800px; margin:auto">
-    <h2>🦢 Diagnóstico por 4 modelos de IA</h2>
+    <h2>🧪 Diagnóstico por 4 modelos de IA</h2>
     <table style="border-collapse: collapse; width:100%; font-size:16px">
         <tr><th style="text-align:left">🔍 Modelo</th><th>Resultado</th><th>Confianza</th></tr>
         <tr><td>🧠 ViT (transformer)</td><td><b>{pred_class_vit}</b></td><td>{confidence_vit:.1%}</td></tr>
@@ -130,7 +132,7 @@ def analizar_lesion_combined(img):
         <tr><td>🔬 TensorFlow (saved_model)</td><td><b>{pred_class_tf}</b></td><td>{confidence_tf:.1%}</td></tr>
     </table>
     <br>
-    <b>🦥 Recomendación automática:</b><br>
+    <b>🩺 Recomendación automática:</b><br>
     """
 
     cancer_risk_score = sum(probs_vit[i] * RISK_LEVELS[i]['weight'] for i in range(7))
@@ -146,6 +148,7 @@ def analizar_lesion_combined(img):
     informe += "</div>"
     return informe, html_chart
 
+# Interfaz Gradio
 demo = gr.Interface(
     fn=analizar_lesion_combined,
     inputs=gr.Image(type="pil", label="Sube una imagen de la lesión"),
@@ -158,3 +161,4 @@ demo = gr.Interface(
 if __name__ == "__main__":
     demo.launch()
 
+