Update app.py

app.py

@@ -1,135 +1,155 @@
-import gradio as gr
 import torch
-import numpy as np
+from transformers import ViTImageProcessor, ViTForImageClassification
+from fastai.learner import load_learner
+from fastai.vision.core import PILImage
+from PIL import Image
 import matplotlib.pyplot as plt
-import base64
+import numpy as np
+import gradio as gr
 import io
-from fastai.vision.all import *
-import tensorflow as tf
-import zipfile
+import base64
 import os
-import traceback
-
-# Descomprimir el modelo si no se ha descomprimido aún
-if not os.path.exists("saved_model"):
-    with zipfile.ZipFile("saved_model.zip", "r") as zip_ref:
-        zip_ref.extractall("saved_model")
-
-# Cargar modelo ISIC con TensorFlow desde la carpeta correcta
-try:
-    model_isic = tf.keras.models.load_model("saved_model")
-except Exception as e:
-    print("\U0001F534 Error al cargar el modelo ISIC con tf.keras.models.load_model:", e)
-    raise
+import zipfile
+import tensorflow as tf
 
-# Cargar modelos fastai
+# --- Extraer y cargar modelo TensorFlow desde zip ---
+zip_path = "saved_model.zip"
+extract_dir = "saved_model"
+if not os.path.exists(extract_dir):
+    os.makedirs(extract_dir)
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(extract_dir)
+
+model_tf = tf.saved_model.load(extract_dir)
+
+# Función helper para inferencia TensorFlow
+def predict_tf(img: Image.Image):
+    # Preprocesar imagen para TF: convertir a tensor float32, normalizar, añadir batch
+    img_resized = img.resize((224,224))  # ajusta según modelo
+    img_np = np.array(img_resized) / 255.0
+    if img_np.shape[-1] == 4:  # eliminar canal alfa si existe
+        img_np = img_np[..., :3]
+    img_tf = tf.convert_to_tensor(img_np, dtype=tf.float32)
+    img_tf = tf.expand_dims(img_tf, axis=0)  # batch dimension
+
+    # Ejecutar modelo (suponiendo firma default)
+    infer = model_tf.signatures["serving_default"]
+    output = infer(img_tf)
+    # Extraemos el primer tensor de salida (puede cambiar según modelo)
+    pred = list(output.values())[0].numpy()[0]
+    probs = tf.nn.softmax(pred).numpy()
+    return probs
+
+# 🔹 Cargar modelo ViT desde Hugging Face
+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 desde archivos locales
 model_malignancy = load_learner("ada_learn_malben.pkl")
 model_norm2000 = load_learner("ada_learn_skin_norm2000.pkl")
 
-# Cargar modelo ViT
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-feature_extractor = AutoImageProcessor.from_pretrained("nateraw/vit-skin-cancer")
-model_vit = AutoModelForImageClassification.from_pretrained("nateraw/vit-skin-cancer")
-
-# Clases y colores
-CLASSES = ['akiec', 'bcc', 'bkl', 'df', 'mel', 'nv', 'vasc']
+# 🔹 Clases y niveles de riesgo
+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: {"label": "akiec", "color": "#FF6F61", "weight": 0.9},
-    1: {"label": "bcc", "color": "#FF8C42", "weight": 0.7},
-    2: {"label": "bkl", "color": "#FFD166", "weight": 0.3},
-    3: {"label": "df", "color": "#06D6A0", "weight": 0.1},
-    4: {"label": "mel", "color": "#EF476F", "weight": 1.0},
-    5: {"label": "nv", "color": "#118AB2", "weight": 0.2},
-    6: {"label": "vasc", "color": "#073B4C", "weight": 0.4},
+    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}
 }
 
-# Preprocesado para TensorFlow ISIC
-def preprocess_image_isic(pil_image):
-    image = pil_image.resize((224, 224))
-    array = np.array(image) / 255.0
-    return np.expand_dims(array, axis=0)
-
-# Función de análisis
 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]
-
-        pred_fast_malignant, _, probs_fast_mal = model_malignancy.predict(img_fastai)
-        prob_malignant = float(probs_fast_mal[1])
-        pred_fast_type, _, probs_fast_type = model_norm2000.predict(img_fastai)
-
-        x_isic = preprocess_image_isic(img)
-        preds_isic = model_isic.predict(x_isic)[0]
-        pred_idx_isic = int(np.argmax(preds_isic))
-        pred_class_isic = CLASSES[pred_idx_isic]
-        confidence_isic = preds_isic[pred_idx_isic]
-
-        colors_bars = [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.set_ylabel("Probabilidad (%)")
-        ax.set_xticks(np.arange(len(CLASSES)))
-        ax.set_xticklabels(CLASSES, rotation=45, ha='right')
-        ax.grid(axis='y', alpha=0.2)
-        plt.tight_layout()
-        buf = io.BytesIO()
-        plt.savefig(buf, format="png")
-        plt.close(fig)
-        img_bytes = buf.getvalue()
-        img_b64 = base64.b64encode(img_bytes).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>
-        <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>🔬 ISIC TensorFlow</td><td><b>{pred_class_isic}</b></td><td>{confidence_isic:.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:
-            informe += "🚨 <b>CRÍTICO</b> – Derivación urgente a oncología dermatológica"
-        elif prob_malignant > 0.4 or cancer_risk_score > 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)"
-        else:
-            informe += "✅ <b>BAJO RIESGO</b> – Seguimiento de rutina (3-6 meses)"
-
-        informe += "</div>"
-        return informe, html_chart
-
-    except Exception as e:
-        print("\U0001F534 ERROR en analizar_lesion_combined:")
-        print(str(e))
-        traceback.print_exc()
-        return f"<b>Error interno:</b> {str(e)}", ""
-
-# INTERFAZ
+    # Convertir imagen para Fastai
+    img_fastai = PILImage.create(img)
+
+    # ViT prediction
+    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]
+
+    # Fast.ai models
+    pred_fast_malignant, _, probs_fast_mal = model_malignancy.predict(img_fastai)
+    prob_malignant = float(probs_fast_mal[1])  # 1 = maligno
+
+    pred_fast_type, _, probs_fast_type = model_norm2000.predict(img_fastai)
+
+    # TensorFlow model prediction
+    probs_tf = predict_tf(img)
+    pred_idx_tf = int(np.argmax(probs_tf))
+    pred_class_tf = CLASSES[pred_idx_tf]
+    confidence_tf = probs_tf[pred_idx_tf]
+
+    # Gráfico ViT
+    colors_bars = [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.set_ylabel("Probabilidad (%)")
+    ax.set_xticks(np.arange(len(CLASSES)))  # evita warning
+    ax.set_xticklabels(CLASSES, rotation=45, ha='right')
+    ax.grid(axis='y', alpha=0.2)
+    plt.tight_layout()
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    plt.close(fig)
+    img_bytes = buf.getvalue()
+    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
+    html_chart = f'<img src="data:image/png;base64,{img_b64}" style="max-width:100%"/>'
+
+    # Informe HTML
+    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>🔬 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>
+    """
+
+    # Recomendación basada en ViT + malignidad (podrías adaptar aquí según TF)
+    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:
+        informe += "🚨 <b>CRÍTICO</b> – Derivación urgente a oncología dermatológica"
+    elif prob_malignant > 0.4 or cancer_risk_score > 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)"
+    else:
+        informe += "✅ <b>BAJO RIESGO</b> – Seguimiento de rutina (3-6 meses)"
+
+    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 + ISIC TensorFlow)",
-    description="Comparación entre ViT transformer (HAM10000), dos modelos Fast.ai y el modelo ISIC TensorFlow.",
+    title="Detector de Lesiones Cutáneas (ViT + Fast.ai + TensorFlow)",
+    description="Comparación entre ViT transformer (HAM10000), dos modelos Fast.ai y un modelo TensorFlow.",
     flagging_mode="never"
 )
 
-# LANZAMIENTO
 if __name__ == "__main__":
     demo.launch()
 
 
+