Spaces:
Sleeping
Sleeping
| 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() | |