Spaces:
Sleeping
Sleeping
File size: 2,480 Bytes
5d2a593 | 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 | import gradio as gr
import torch
from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
ORIGINAL_MODEL_NAME = "facebook/dinov2-base"
LOCAL_MODEL_PATH = "./model-files"
PRETTY_NAMES_MAP = {
"seg_neutrophil": "Segmented neutrophil",
"lymphocyte": "Lymphocyte",
"band_neutrophil": "Banded neutrophil",
"eosinophil": "Eosinophil",
"monocyte": "Monocyte",
"basophil": "Basophil",
"blast": "Blast",
"immature_wbc": "Immature WBCs",
"myelocyte": "Myelocyte",
"promyelocyte": "Promyelocyte",
"abnormal_lymphocyte": "Abnormal lymphocyte",
"smudge": "Smudge",
"metamyelocyte": "Metamyelocyte",
"agg_plt": "Aggregated platelet",
"n_rbc": "Nucleated Red Blood Cell",
"g_plt": "Giant platelet",
"artifact": "Artifact",
"unk_wbc": "Unknown WBC",
}
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")
try:
print(f"Loading processor from: '{ORIGINAL_MODEL_NAME}'")
processor = AutoImageProcessor.from_pretrained(ORIGINAL_MODEL_NAME)
print(f"Loading fine-tuned model from: '{LOCAL_MODEL_PATH}'")
model = AutoModelForImageClassification.from_pretrained(LOCAL_MODEL_PATH)
model.to(DEVICE)
model.eval()
print("Model and processor loaded successfully!")
except Exception as e:
print(f"Error loading model: {e}")
raise
def predict(image: Image.Image):
if image is None:
return {}
image = image.convert("RGB")
inputs = processor(images=image, return_tensors="pt").to(DEVICE)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
top3_probs, top3_indices = torch.topk(probabilities, 3)
results = {}
for i in range(top3_probs.size(0)):
technical_name = model.config.id2label[top3_indices[i].item()]
pretty_name = PRETTY_NAMES_MAP.get(technical_name, technical_name)
results[pretty_name] = top3_probs[i].item()
return results
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload White Blood Cell Image"),
outputs=gr.Label(num_top_classes=3, label="Top 3 Predictions"),
title="White Blood Cell Classifier",
description="Upload a microscopic image of a white blood cell to get a prediction of its type. Model based on DinoV2.",
allow_flagging="never",
)
if __name__ == "__main__":
iface.launch()
|