Update app.py

app.py

@@ -1,19 +1,18 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-import pandas as pd
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModel
 
-# --- Load BLIP model ---
-blip_path = "Models/blip_skin_cancer"
-processor = BlipProcessor.from_pretrained(blip_path)
-blip_model = BlipForConditionalGeneration.from_pretrained(blip_path).to("cpu")
+# --- Load BLIP model from Hugging Face Hub ---
+processor = BlipProcessor.from_pretrained("Mewish/blip_medical")
+blip_model = BlipForConditionalGeneration.from_pretrained("Mewish/blip_medical").to("cpu")
 
-# --- Load BioClinicalBERT + BiLSTM classifier ---
+# --- Load BioClinicalBERT backbone ---
 tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 bert_model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
 
+# --- Define BiLSTM classifier ---
 class BiLSTMClassifier(nn.Module):
     def __init__(self, bert_model, hidden_dim=256, num_classes=7, dropout=0.5):
         super().__init__()
@@ -30,9 +29,10 @@ class BiLSTMClassifier(nn.Module):
         pooled = self.dropout(pooled)
         return self.fc(pooled)
 
-# --- Load trained classifier weights ---
+# --- Load classifier weights from Hugging Face Hub ---
 classifier = BiLSTMClassifier(bert_model).to("cpu")
-classifier.load_state_dict(torch.load("Models/classifier.pt", map_location="cpu"))
+state_dict_url = "https://huggingface.co/Mewish/skin_cancer_classifier/resolve/main/classifier.pt"
+classifier.load_state_dict(torch.hub.load_state_dict_from_url(state_dict_url, map_location="cpu"))
 classifier.eval()
 
 # --- Classes ---
@@ -60,4 +60,25 @@ def predict(image):
     attention_mask = text_inputs["attention_mask"].to("cpu")
 
     with torch.no_grad():
-        outputs
+        outputs = classifier(input_ids, attention_mask)
+        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+        pred_class = classes[probs.argmax()]
+        confidence = probs.max()
+
+    return caption, pred_class, float(confidence)
+
+# --- Gradio Interface ---
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Textbox(label="Generated Caption"),
+        gr.Textbox(label="Predicted Class"),
+        gr.Number(label="Confidence Score")
+    ],
+    title="Skin Cancer AI Agent",
+    description="Upload a lesion image to generate a caption and predict the cancer type."
+)
+
+if __name__ == "__main__":
+    iface.launch