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

os.environ["OPENCV_LOG_LEVEL"] = "SILENT"

def load_model():
    model = smp.Unet(
        encoder_name="resnet34",
        encoder_weights=None,
        in_channels=3,
        classes=1
    )
    model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))
    model.eval()
    return model

model = load_model()

def segment_image(input_img, threshold):
    if input_img is None:
        return None
        
    transform = A.Compose([
        A.Resize(256, 256),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        ToTensorV2()
    ])
    
    img_np = np.array(input_img)
    augmented = transform(image=img_np)
    img_tensor = augmented["image"].unsqueeze(0)
    
    with torch.no_grad():
        output = model(img_tensor)
        mask = torch.sigmoid(output).squeeze().numpy()
        
    binary_mask = (mask > threshold).astype(np.uint8)
    num_labels, labels = cv2.connectedComponents(binary_mask)
    
    if num_labels <= 1:
        return Image.fromarray(np.zeros((256, 256, 3), dtype=np.uint8))

    label_hue = np.uint8(179 * labels / np.max(labels))
    blank_ch = np.ones_like(label_hue) * 255
    labeled_img = cv2.merge([label_hue, blank_ch, blank_ch])
    labeled_img = cv2.cvtColor(labeled_img, cv2.COLOR_HSV2RGB)
    labeled_img[labels == 0] = 0
    
    return Image.fromarray(labeled_img)

interface = gr.Interface(
    fn=segment_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Aerial Image"),
        gr.Slider(minimum=0.01, maximum=0.99, value=0.25, step=0.01, label="Confidence Threshold")
    ],
    outputs=gr.Image(type="pil", label="Instance Segmentation"),
    title="Landvisor Building Segmentation",
    description="Building segmentation via transfer learning and pfCE."
)

if __name__ == "__main__":
    interface.launch()