import gradio as gr
import torch
import numpy as np
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import segmentation_models_pytorch as smp
import cv2 


MODEL_FILENAME = "brain_tumor_unet_model.pth"


DEVICE = "cpu"


val_augs = A.Compose([
    A.Resize(256, 256),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])



model = smp.Unet(
    encoder_name="resnet34",
    encoder_weights=None,
    in_channels=3,
    classes=1,
)
model.load_state_dict(torch.load(MODEL_FILENAME, map_location=DEVICE))
model.to(DEVICE)
model.eval()
print("Model loaded successfully.")





def predict(input_image):

    image_np = np.array(input_image)
    

    augmented = val_augs(image=image_np)
    image_tensor = augmented['image'].unsqueeze(0).to(DEVICE)
    

    with torch.no_grad():
        logits = model(image_tensor)
        probs = torch.sigmoid(logits)
        mask = (probs > 0.5).float().squeeze().cpu().numpy()


    resized_image = cv2.resize(image_np, (256, 256))
    

    mask_colored = np.zeros((256, 256, 3), dtype=np.uint8)
    mask_colored[mask == 1] = [0, 255, 0]

    

    overlay = resized_image.copy()
    overlay[mask == 1] = [0, 255, 0]
    
    return overlay


title = "Brain Tumor Segmentation Demo"
description = (
    "Upload a brain MRI scan to see the tumor region segmented by a U-Net model. "
    "This model was trained on the LGG Segmentation Dataset. "
    
)

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Brain MRI"),
    outputs=gr.Image(type="pil", label="Segmented Tumor (in Green)"),
    title=title,
    description=description,
    examples=[
        ["TCGA_CS_4941_19960909_16.jpg"]
    ]
)


iface.launch()