import gradio as gr
from matplotlib import gridspec
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch
from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation

MODEL_ID = "nvidia/segformer-b4-finetuned-cityscapes-1024-1024"
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForSemanticSegmentation.from_pretrained(MODEL_ID)


def city_palette():
    return [
        [128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], [190, 153, 153],
        [153, 153, 153], [250, 170, 30], [220, 220, 0], [107, 142, 35], [152, 251, 152],
        [70, 130, 180], [220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70],
        [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32],
    ]


labels_list = []
with open("labels.txt", "r", encoding="utf-8") as fp:
    for line in fp:
        labels_list.append(line.rstrip("\n"))

colormap = np.asarray(city_palette(), dtype=np.uint8)


def label_to_color_image(label):
    if label.ndim != 2:
        raise ValueError("Expect 2-D input label")
    if np.max(label) >= len(colormap):
        raise ValueError("label value too large.")
    return colormap[label]


def draw_plot(pred_img, seg_np):
    fig = plt.figure(figsize=(20, 15))
    grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])

    plt.subplot(grid_spec[0])
    plt.imshow(pred_img)
    plt.axis('off')
    plt.title('Segmentation Result', fontsize=20, pad=20)

    LABEL_NAMES = np.asarray(labels_list)
    FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
    FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)

    unique_labels = np.unique(seg_np.astype("uint8"))
    ax = plt.subplot(grid_spec[1])
    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
    ax.yaxis.tick_right()
    plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
    plt.xticks([], [])
    ax.tick_params(width=0.0, labelsize=25)
    plt.title('Detected Classes', fontsize=20, pad=20)
    return fig


def run_inference(input_img):
    # input: numpy array from gradio -> PIL
    img = Image.fromarray(input_img.astype(np.uint8)) if isinstance(input_img, np.ndarray) else input_img
    if img.mode != "RGB":
        img = img.convert("RGB")

    inputs = processor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits

    # resize to original
    upsampled = torch.nn.functional.interpolate(
        logits, size=img.size[::-1], mode="bilinear", align_corners=False
    )
    seg = upsampled.argmax(dim=1)[0].cpu().numpy().astype(np.uint8)

    # colorize & overlay
    color_seg = colormap[seg]
    pred_img = (np.array(img) * 0.5 + color_seg * 0.5).astype(np.uint8)

    fig = draw_plot(pred_img, seg)
    return fig


with gr.Blocks(theme=gr.themes.Soft(), title="도시 장면 분할") as demo:
    gr.Markdown(
        """
        # 도시 장면 영상 분할 (City Scene Segmentation)
        **Cityscapes 데이터셋으로 학습된 SegFormer 모델**을 활용한 도로 및 도시 장면 분할 데모입니다.

        도로, 건물, 차량, 보행자 등 19개 클래스를 자동으로 인식하고 분할합니다.
        """
    )

    gr.Markdown(
        """
        ---
        ###  감지 가능한 클래스 (19개)
        `도로`, `보도`, `건물`, `벽`, `울타리`, `기둥`, `신호등`, `표지판`, `식물`, 
        `지형`, `하늘`, `사람`, `자전거 탑승자`, `자동차`, `트럭`, `버스`, `기차`, `오토바이`, `자전거`
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_img = gr.Image(
                type="numpy",
                label="입력 이미지",
                height=400
            )
            submit_btn = gr.Button(
                "분할 실행",
                variant="primary",
                size="lg"
            )

            gr.Markdown("###  예시 이미지")
            gr.Examples(
                examples=[
                    "road-2.jpg",
                    "road-3.jpeg",
                ],
                inputs=input_img,
                label="도시/도로 장면 샘플"
            )

        with gr.Column(scale=1):
            output_plot = gr.Plot(label=" 분할 결과 및 범례")


    # 이벤트 핸들러
    submit_btn.click(
        fn=run_inference,
        inputs=input_img,
        outputs=output_plot
    )

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