Update app.py

app.py

@@ -1,6 +1,5 @@
 import torch
 from transformers import ViTImageProcessor, ViTForImageClassification
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 from fastai.learner import load_learner
 from fastai.vision.core import PILImage
 from PIL import Image
@@ -9,84 +8,72 @@ import numpy as np
 import gradio as gr
 import io
 import base64
-import os
-import zipfile
 
-# --- Cargar modelo ViT ---
+# --- Cargar modelo ViT preentrenado fine‑tuned HAM10000 ---
+TF_MODEL_NAME = "Anwarkh1/Skin_Cancer-Image_Classification"
+feature_extractor_tf = ViTImageProcessor.from_pretrained(TF_MODEL_NAME)
+model_tf_vit = ViTForImageClassification.from_pretrained(TF_MODEL_NAME)
+model_tf_vit.eval()
+
+# 🔹 Cargar modelo ViT base
 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()
 
-# --- Cargar modelos Fast.ai ---
+# 🔹 Cargar modelos Fast.ai locales
 model_malignancy = load_learner("ada_learn_malben.pkl")
 model_norm2000 = load_learner("ada_learn_skin_norm2000.pkl")
 
-# --- Cargar modelo EfficientNetB3 desde Hugging Face ---
-model_effnet = AutoModelForImageClassification.from_pretrained("syaha/skin_cancer_detection_model")
-extractor_effnet = AutoFeatureExtractor.from_pretrained("syaha/skin_cancer_detection_model")
-model_effnet.eval()
-
+# Clases estándar de HAM10000
 CLASSES = [
     "Queratosis actínica / Bowen", "Carcinoma células basales",
     "Lesión queratósica benigna", "Dermatofibroma",
     "Melanoma maligno", "Nevus melanocítico", "Lesión vascular"
 ]
-
 RISK_LEVELS = {
     0: {'level': 'Moderado', 'color': '#ffaa00', 'weight': 0.6},
-    1: {'level': 'Alto', 'color': '#ff4444', 'weight': 0.8},
-    2: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1},
-    3: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1},
-    4: {'level': 'Crítico', 'color': '#cc0000', 'weight': 1.0},
-    5: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1},
-    6: {'level': 'Bajo', 'color': '#44ff44', 'weight': 0.1}
+    1: {'level': 'Alto',     'color': '#ff4444', 'weight': 0.8},
+    2: {'level': 'Bajo',     'color': '#44ff44', 'weight': 0.1},
+    3: {'level': 'Bajo',     'color': '#44ff44', 'weight': 0.1},
+    4: {'level': 'Crítico',  'color': '#cc0000', 'weight': 1.0},
+    5: {'level': 'Bajo',     'color': '#44ff44', 'weight': 0.1},
+    6: {'level': 'Bajo',     'color': '#44ff44', 'weight': 0.1}
 }
-
-MALIGNANT_INDICES = [0, 1, 4]  # clases de riesgo alto/crítico
+MALIGNANT_INDICES = [0, 1, 4]  # akiec, bcc, melanoma
 
 def analizar_lesion_combined(img):
-    try:
-        img_fastai = PILImage.create(img)
-        inputs = feature_extractor(img, return_tensors="pt")
-        with torch.no_grad():
-            outputs = model_vit(**inputs)
-            probs_vit = outputs.logits.softmax(dim=-1).cpu().numpy()[0]
-        pred_idx_vit = int(np.argmax(probs_vit))
-        pred_class_vit = CLASSES[pred_idx_vit]
-        confidence_vit = probs_vit[pred_idx_vit]
-    except Exception as e:
-        pred_class_vit = "Error"
-        confidence_vit = 0.0
-        probs_vit = np.zeros(len(CLASSES))
-
-    try:
-        pred_fast_malignant, _, probs_fast_mal = model_malignancy.predict(img_fastai)
-        prob_malignant = float(probs_fast_mal[1])
-    except:
-        prob_malignant = 0.0
-
-    try:
-        pred_fast_type, _, _ = model_norm2000.predict(img_fastai)
-    except:
-        pred_fast_type = "Error"
-
-    try:
-        inputs_eff = extractor_effnet(images=img, return_tensors="pt")
-        with torch.no_grad():
-            outputs_eff = model_effnet(**inputs_eff)
-            probs_eff = outputs_eff.logits.softmax(dim=-1).cpu().numpy()[0]
-        pred_idx_eff = int(np.argmax(probs_eff))
-        confidence_eff = probs_eff[pred_idx_eff]
-        pred_class_eff = model_effnet.config.id2label[str(pred_idx_eff)]
-    except Exception as e:
-        pred_class_eff = "Error"
-        confidence_eff = 0.0
-
-    colors_bars = [RISK_LEVELS[i]['color'] for i in range(7)]
+    img_fastai = PILImage.create(img)
+
+    # ViT base
+    inputs = feature_extractor(img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model_vit(**inputs)
+        probs_vit = outputs.logits.softmax(dim=-1).cpu().numpy()[0]
+    idx_vit = int(np.argmax(probs_vit))
+    class_vit = CLASSES[idx_vit]
+    conf_vit = probs_vit[idx_vit]
+
+    # Fast.ai modelos
+    _, _, probs_mal = model_malignancy.predict(img_fastai)
+    prob_malign = float(probs_mal[1])
+    pred_fast_type, _, _ = model_norm2000.predict(img_fastai)
+
+    # ViT pre-trained fine-tuned (último modelo recomendado)
+    inputs_tf = feature_extractor_tf(img, return_tensors="pt")
+    with torch.no_grad():
+        outputs_tf = model_tf_vit(**inputs_tf)
+        probs_tf = outputs_tf.logits.softmax(dim=-1).cpu().numpy()[0]
+    idx_tf = int(np.argmax(probs_tf))
+    class_tf_model = CLASSES[idx_tf]
+    conf_tf = probs_tf[idx_tf]
+    mal_tf = "Maligno" if idx_tf in MALIGNANT_INDICES else "Benigno"
+
+    # Gráfico ViT base
+    colors = [RISK_LEVELS[i]['color'] for i in range(7)]
     fig, ax = plt.subplots(figsize=(8, 3))
-    ax.bar(CLASSES, probs_vit*100, color=colors_bars)
-    ax.set_title("Probabilidad ViT por tipo de lesión")
+    ax.bar(CLASSES, probs_vit*100, color=colors)
+    ax.set_title("Probabilidad ViT base por tipo de lesión")
     ax.set_ylabel("Probabilidad (%)")
     ax.set_xticks(np.arange(len(CLASSES)))
     ax.set_xticklabels(CLASSES, rotation=45, ha='right')
@@ -95,45 +82,39 @@ def analizar_lesion_combined(img):
     buf = io.BytesIO()
     plt.savefig(buf, format="png")
     plt.close(fig)
-    img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
-    html_chart = f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
+    html_chart = f'<img src="data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}" 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>
-    <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>
-        <tr><td>🧬 Fast.ai (clasificación)</td><td><b>{pred_fast_type}</b></td><td>N/A</td></tr>
-        <tr><td>⚠️ Fast.ai (malignidad)</td><td><b>{"Maligno" if prob_malignant > 0.5 else "Benigno"}</b></td><td>{prob_malignant:.1%}</td></tr>
-        <tr><td>🔬 EfficientNetB3 (HAM10000)</td><td><b>{pred_class_eff}</b></td><td>{confidence_eff:.1%}</td></tr>
-    </table>
-    <br>
-    <b>🧪 Recomendación automática:</b><br>
+    <h2>🧪 Diagnóstico por múltiples modelos de IA</h2>
+    <table style="width:100%; font-size:16px; border-collapse:collapse">
+      <tr><th>Modelo</th><th>Resultado</th><th>Confianza</th></tr>
+      <tr><td>🧠 ViT base</td><td><b>{class_vit}</b></td><td>{conf_vit:.1%}</td></tr>
+      <tr><td>🧬 Fast.ai (tipo)</td><td><b>{pred_fast_type}</b></td><td>N/A</td></tr>
+      <tr><td>⚠️ Fast.ai (malignidad)</td><td><b>{'Maligno' if prob_malign > 0.5 else 'Benigno'}</b></td><td>{prob_malign:.1%}</td></tr>
+      <tr><td>🌟 ViT fined‑tuned (HAM10000)</td><td><b>{mal_tf} ({class_tf_model})</b></td><td>{conf_tf:.1%}</td></tr>
+    </table><br>
+    <b>🩺 Recomendación automática:</b><br>
     """
-
-    cancer_risk_score = sum(probs_vit[i] * RISK_LEVELS[i]['weight'] for i in range(7))
-    if prob_malignant > 0.7 or cancer_risk_score > 0.6:
+    risk = sum(probs_vit[i] * RISK_LEVELS[i]['weight'] for i in range(7))
+    if prob_malign > 0.7 or risk > 0.6:
         informe += "🚨 <b>CRÍTICO</b> – Derivación urgente a oncología dermatológica"
-    elif prob_malignant > 0.4 or cancer_risk_score > 0.4:
+    elif prob_malign > 0.4 or risk > 0.4:
         informe += "⚠️ <b>ALTO RIESGO</b> – Consulta con dermatólogo en 7 días"
-    elif cancer_risk_score > 0.2:
-        informe += "📋 <b>RIESGO MODERADO</b> – Evaluación programada (2-4 semanas)"
+    elif risk > 0.2:
+        informe += "📋 <b>RIESGO MODERADO</b> – Evaluación programada en 2-4 semanas"
     else:
         informe += "✅ <b>BAJO RIESGO</b> – Seguimiento de rutina (3-6 meses)"
+    informe += "</div>"""
 
-    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"),
-    outputs=[gr.HTML(label="Informe combinado"), gr.HTML(label="Gráfico ViT")],
-    title="Detector de Lesiones Cutáneas (ViT + Fast.ai + EfficientNetB3)",
-    description="Comparación entre ViT transformer (HAM10000), dos modelos Fast.ai y un modelo EfficientNetB3.",
-    flagging_mode="never"
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.HTML(label="Informe"), gr.HTML(label="Gráfico ViT base")],
+    title="Detector de Lesiones Cutáneas (ViT + Fast.ai)",
 )
-
 if __name__ == "__main__":
     demo.launch()
+