|
|
import json |
|
|
import os |
|
|
from functools import lru_cache |
|
|
from pathlib import Path |
|
|
from typing import Dict, List, Tuple |
|
|
|
|
|
import psutil |
|
|
import torch |
|
|
import gradio as gr |
|
|
from transformers import AutoModelForZeroShotImageClassification, AutoProcessor |
|
|
|
|
|
from utils.cache_manager import cached_inference, configure_cache |
|
|
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") |
|
|
|
|
|
physical_cores = psutil.cpu_count(logical=False) or psutil.cpu_count() or 1 |
|
|
torch.set_num_threads(min(physical_cores, 4)) |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
processor = AutoProcessor.from_pretrained( |
|
|
MODEL_ID, |
|
|
token=HF_TOKEN, |
|
|
use_fast=True, |
|
|
) |
|
|
model = AutoModelForZeroShotImageClassification.from_pretrained( |
|
|
MODEL_ID, |
|
|
token=HF_TOKEN, |
|
|
dtype=torch.float32, |
|
|
).to(device) |
|
|
model.eval() |
|
|
|
|
|
configure_cache(model, processor) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
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.", |
|
|
) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
server_name = os.getenv("SERVER_NAME", "0.0.0.0") |
|
|
port_env = os.getenv("SERVER_PORT") or os.getenv("PORT") or "7860" |
|
|
share_env = os.getenv("GRADIO_SHARE", "false").lower() |
|
|
queue_env = os.getenv("GRADIO_QUEUE", "false").lower() |
|
|
|
|
|
share_enabled = share_env in {"1", "true", "yes"} |
|
|
queue_enabled = queue_env in {"1", "true", "yes"} |
|
|
|
|
|
app_to_launch = demo.queue() if queue_enabled else demo |
|
|
|
|
|
app_to_launch.launch( |
|
|
server_name=server_name, |
|
|
server_port=int(port_env), |
|
|
share=share_enabled, |
|
|
show_api=False, |
|
|
) |
|
|
|
