import gradio as gr
import torch
from transformers import AutoModelForZeroShotObjectDetection, AutoProcessor
from PIL import Image
import numpy as np
import cv2

model_id = "iSEE-Laboratory/llmdet_tiny"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

prompts = ["E. coli rods", "Klebsiella rods", "yeast cells", "WBC", "RBC"]

def predict(input_img):
    img = Image.fromarray(input_img.astype('uint8'), 'RGB')
    
    overlay = np.array(input_img)
    overlay = cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)
    
    summary = "Detected elements:\n"
    rod_count = 0
    yeast_count = 0
    wbc_count = 0
    rbc_count = 0
    
    for prompt in prompts:
        text_labels = [prompt]
        inputs = processor(images=img, text=text_labels, return_tensors="pt").to(device)
        
        with torch.no_grad():
            outputs = model(**inputs)
        
        results = processor.post_process_grounded_object_detection(
            outputs,
            threshold=0.3,
            target_sizes=[img.size[::-1]]
        )[0]
        
        count = len(results["boxes"])
        if "rods" in prompt:
            rod_count += count
            color = (255, 0, 0)
        elif "yeast" in prompt:
            yeast_count += count
            color = (0, 255, 0)
        elif "WBC" in prompt:
            wbc_count += count
            color = (255, 255, 0)
        elif "RBC" in prompt:
            rbc_count += count
            color = (0, 0, 255)
        
        if count > 0:
            summary += f"- {prompt}: {count} detections\n"
            
            for box in results["boxes"]:
                box = [int(i) for i in box.tolist()]
                cv2.rectangle(overlay, (box[0], box[1]), (box[2], box[3]), color, thickness=2)
    
    overlay = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
    
    total = rod_count + yeast_count + wbc_count + rbc_count
    if total == 0:
        summary += "No significant elements detected."
    else:
        if total > 50:
            cfu = ">10^6 CFU/ml"
        elif total > 20:
            cfu = "10^5–10^6 CFU/ml"
        else:
            cfu = "<10^5 CFU/ml"
        summary += f"\nEstimated load: {cfu}"
    
    return overlay, summary

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Urine Microscopy Image"),
    outputs=[
        gr.Image(label="Image with Detection Boxes"),
        gr.Textbox(label="Preliminary Findings")
    ],
    title="UTI-Vision Demo",
    description="Upload a wet-mount urine image. Red = rods (bacteria), Green = yeast, Yellow = WBC, Blue = RBC. Rough CFU estimate."
)

iface.launch()