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()