Spaces:
Runtime error
Runtime error
File size: 1,592 Bytes
0ddd5ad 39ea0a2 5322bcf 39ea0a2 8ad1547 ccfcbd4 39ea0a2 c2882e7 8ad1547 c2882e7 8ad1547 c2882e7 8ad1547 39ea0a2 8ad1547 929016b 8ad1547 |
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 |
from transformers import AutoModelForImageClassification, AutoImageProcessor
import torch
from PIL import Image
import gradio as gr
model_name = "codewithdark/vit-chest-xray"
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)
model.eval()
labels = ['Cardiomegaly', 'Edema', 'Consolidation', 'No Finding', 'Pneumonia']
target_labels = ['Pneumonia', 'Consolidation', 'Edema']
target_idxs = [labels.index(lbl) for lbl in target_labels]
def predict(image):
# Make sure image is RGB
if image.mode != "RGB":
image = image.convert("RGB")
```
# Process the image properly
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
# Keep batch dimension for safety
probs = torch.sigmoid(logits)[0] # [batch, num_labels] -> [num_labels]
detected = []
results = []
for idx, lbl in zip(target_idxs, target_labels):
prob = probs[idx].item()
status = "YES" if prob > 0.5 else "NO"
results.append(f"{lbl}: {status} ({prob:.2f})")
if status == "YES":
detected.append(lbl)
if detected:
summary = f"⚠️ Patient shows signs of: {', '.join(detected)}."
else:
summary = "✅ Patient appears healthy — no major lung issues detected."
return "\n".join(results + ["\n" + summary])
```
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs="text",
title="Chest X-ray Disease Detector",
description="Upload a chest X-ray to detect Pneumonia, Consolidation, and Edema. Gives clear patient health summary."
)
iface.launch()
|