|
|
import json |
|
|
import os |
|
|
from functools import lru_cache |
|
|
from pathlib import Path |
|
|
from typing import Dict, List, Tuple |
|
|
|
|
|
import torch |
|
|
import gradio as gr |
|
|
from transformers import AutoModelForZeroShotImageClassification, AutoProcessor |
|
|
|
|
|
from utils.cache_manager import cached_inference |
|
|
from utils.modality_router import detect_modality |
|
|
|
|
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent |
|
|
LABEL_DIR = BASE_DIR / "labels" |
|
|
MODEL_ID = "google/medsiglip-448" |
|
|
|
|
|
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN") |
|
|
|
|
|
torch.set_num_threads(1) |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
model_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 |
|
|
|
|
|
processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN) |
|
|
model = AutoModelForZeroShotImageClassification.from_pretrained( |
|
|
MODEL_ID, |
|
|
token=HF_TOKEN, |
|
|
torch_dtype=model_dtype, |
|
|
).to(device) |
|
|
model.eval() |
|
|
|
|
|
|
|
|
LABEL_OVERRIDES = { |
|
|
"xray": "chest_labels.json", |
|
|
"mri": "brain_labels.json", |
|
|
} |
|
|
|
|
|
|
|
|
@lru_cache(maxsize=None) |
|
|
def load_labels(file_name: str) -> List[str]: |
|
|
label_path = LABEL_DIR / file_name |
|
|
with label_path.open("r", encoding="utf-8") as handle: |
|
|
return json.load(handle) |
|
|
|
|
|
|
|
|
def get_candidate_labels(image_path: str) -> Tuple[str, ...]: |
|
|
modality = detect_modality(image_path) |
|
|
candidate_path = LABEL_DIR / f"{modality}_labels.json" |
|
|
if not candidate_path.exists(): |
|
|
override = LABEL_OVERRIDES.get(modality) |
|
|
if override: |
|
|
candidate_path = LABEL_DIR / override |
|
|
if not candidate_path.exists(): |
|
|
candidate_path = LABEL_DIR / "general_labels.json" |
|
|
|
|
|
return tuple(load_labels(candidate_path.name)) |
|
|
|
|
|
|
|
|
def classify_medical_image(image_path: str) -> Dict[str, float]: |
|
|
if not image_path: |
|
|
return {} |
|
|
|
|
|
candidate_labels = get_candidate_labels(image_path) |
|
|
scores = cached_inference(image_path, candidate_labels, model, processor) |
|
|
|
|
|
if not scores: |
|
|
return {} |
|
|
|
|
|
results = sorted(zip(candidate_labels, scores), key=lambda x: x[1], reverse=True) |
|
|
top_results = results[:5] |
|
|
|
|
|
return {label: float(score) for label, score in top_results} |
|
|
|
|
|
|
|
|
demo = gr.Interface( |
|
|
fn=classify_medical_image, |
|
|
inputs=gr.Image(type="filepath", label="📤 Upload Medical Image"), |
|
|
outputs=gr.Label(num_top_classes=5, label="🧠 Top Predictions"), |
|
|
title="🩻 MedSigLIP Smart Medical Classifier", |
|
|
description="Zero-shot model with automatic label filtering for different modalities.", |
|
|
flagging_mode="never", |
|
|
) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False) |
|
|
|
