Update app.py

app.py

@@ -1,5 +1,6 @@
 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
@@ -10,51 +11,28 @@ import io
 import base64
 import os
 import zipfile
-import tensorflow as tf
 
-# --- 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)
-TF_NUM_CLASSES = 7  # asumimos que son las mismas que CLASSES
-
-# Función helper para inferencia TensorFlow
-def predict_tf(img: Image.Image):
-    try:
-        img_resized = img.resize((224,224))
-        img_np = np.array(img_resized) / 255.0
-        if img_np.shape[-1] == 4:
-            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)
-
-        infer = model_tf.signatures["serving_default"]
-        output = infer(img_tf)
-        pred = list(output.values())[0].numpy()[0]
-        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(TF_NUM_CLASSES)
-
-# --- Cargar modelos ---
+# --- Cargar modelo ViT ---
 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 ---
 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()
+
 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},
@@ -94,16 +72,16 @@ def analizar_lesion_combined(img):
         pred_fast_type = "Error"
 
     try:
-        probs_tf = predict_tf(img)
-        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 = f"Desconocido"
-    except:
-        pred_class_tf = "Error"
-        confidence_tf = 0.0
+        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)]
     fig, ax = plt.subplots(figsize=(8, 3))
@@ -128,7 +106,7 @@ def analizar_lesion_combined(img):
         <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>
+        <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>
@@ -152,12 +130,10 @@ 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 + TensorFlow)",
-    description="Comparación entre ViT transformer (HAM10000), dos modelos Fast.ai y un modelo TensorFlow.",
+    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"
 )
 
 if __name__ == "__main__":
     demo.launch()
-
-