Spaces:
Build error
Build error
File size: 3,108 Bytes
31447a6 9f5d1e5 31447a6 9f5d1e5 31447a6 9f5d1e5 31447a6 1a1572a 9f5d1e5 1a1572a 9f5d1e5 1a1572a 9f5d1e5 98e5a10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import gradio as gr
import torch
import torch.nn as nn
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModel
# --- 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 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__()
self.bert = bert_model
self.lstm = nn.LSTM(768, hidden_dim, batch_first=True, bidirectional=True)
self.dropout = nn.Dropout(dropout)
self.fc = nn.Linear(hidden_dim*2, num_classes)
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
embeddings = outputs.last_hidden_state
lstm_out, _ = self.lstm(embeddings)
pooled = torch.mean(lstm_out, dim=1)
pooled = self.dropout(pooled)
return self.fc(pooled)
# --- Load classifier weights from Hugging Face Hub ---
classifier = BiLSTMClassifier(bert_model).to("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 ---
classes = [
"actinic keratosis",
"basal cell carcinoma",
"dermatofibroma",
"nevus",
"pigmented benign keratosis",
"squamous cell carcinoma",
"vascular lesion"
]
# --- Prediction function ---
def predict(image):
# Step 1: BLIP caption
inputs = processor(images=image, return_tensors="pt").to("cpu")
generated_ids = blip_model.generate(**inputs, max_length=50)
caption = processor.decode(generated_ids[0], skip_special_tokens=True)
# Step 2: Classifier prediction
text_inputs = tokenizer(caption, max_length=50, padding="max_length",
truncation=True, return_tensors="pt")
input_ids = text_inputs["input_ids"].to("cpu")
attention_mask = text_inputs["attention_mask"].to("cpu")
with torch.no_grad():
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()
|