import torch
from torchvision import models, transforms
from PIL import Image
import gradio as gr
from typing import Union

class Preprocessor:
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def __call__(self, image: Image.Image) -> torch.Tensor:
        return self.transform(image)

class SegmentationModel:
    def __init__(self):
        self.model = models.segmentation.deeplabv3_resnet101(pretrained=True)
        self.model.eval()
        if torch.cuda.is_available():
            self.model.to('cuda')

    def predict(self, input_batch: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            if torch.cuda.is_available():
                input_batch = input_batch.to('cuda')
            output: torch.Tensor = self.model(input_batch)['out'][0]
        return output

class OutputColorizer:
    def __init__(self):
        palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
        colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
        self.colors = (colors % 255).numpy().astype("uint8")

    def colorize(self, output: torch.Tensor) -> Image.Image:
        colorized_output = Image.fromarray(output.byte().cpu().numpy(), mode='P')
        colorized_output.putpalette(self.colors.ravel())
        return colorized_output

class Segmenter:
    def __init__(self):
        self.preprocessor = Preprocessor()
        self.model = SegmentationModel()
        self.colorizer = OutputColorizer()

    def segment(self, image: Union[Image.Image, torch.Tensor]) -> Image.Image:
        input_image: Image.Image = image.convert("RGB")
        input_tensor: torch.Tensor = self.preprocessor(input_image)
        input_batch: torch.Tensor = input_tensor.unsqueeze(0)
        output: torch.Tensor = self.model.predict(input_batch)
        output_predictions: torch.Tensor = output.argmax(0)
        return self.colorizer.colorize(output_predictions)

class GradioApp:
    def __init__(self, segmenter: Segmenter):
        self.segmenter = segmenter

    def launch(self):
        with gr.Blocks() as demo:
            gr.Markdown("<h1 style='text-align: center; color: #4CAF50;'>Deeplabv3 Segmentation</h1>")
            gr.Markdown("<p style='text-align: center;'>Upload an image to perform semantic segmentation using Deeplabv3 ResNet101.</p>")
            gr.Markdown("""
                ### Model Information
                **DeepLabv3 with ResNet101** is a convolutional neural network model designed for semantic image segmentation. 
                It utilizes atrous convolution to capture multi-scale context by using different atrous rates.
                """)
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(type='pil', label="Input Image", show_label=False)
                with gr.Column():
                    image_output = gr.Image(type='pil', label="Segmented Output", show_label=False)
                    
            button = gr.Button("Segment")
            button.click(fn=self.segmenter.segment, inputs=image_input, outputs=image_output)

            gr.Markdown("### Example Images")
            gr.Examples(
                examples=[
                    ["https://www.timeforkids.com/wp-content/uploads/2024/01/Snapshot_20240126.jpg?w=1024"],
                    ["https://www.timeforkids.com/wp-content/uploads/2023/09/G3G5_230915_puffins_on_the_rise.jpg?w=1024"],
                    ["https://www.timeforkids.com/wp-content/uploads/2024/03/G3G5_240412_bug_eyed.jpg?w=1024"]
                ],
                inputs=image_input,
                outputs=image_output,
                label="Click an example to use it"
            )
            
        demo.launch()

if __name__ == "__main__":
    segmenter = Segmenter()
    app = GradioApp(segmenter)
    app.launch()