import gradio as gr
from matplotlib import gridspec
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import tensorflow as tf
from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
import os

feature_extractor = SegformerFeatureExtractor.from_pretrained(
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)
model = TFSegformerForSemanticSegmentation.from_pretrained(
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)

def ade_palette():
    """ADE20K 팔레트 - 각 클래스를 RGB 값에 매핑합니다."""
    return [
        [204, 87, 92],
        [112, 185, 212],
        [45, 189, 106],
        [234, 123, 67],
        [78, 56, 123],
        [210, 32, 89],
        [90, 180, 56],
        [155, 102, 200],
        [33, 147, 176],
        [255, 183, 76],
        [67, 123, 89],
        [190, 60, 45],
        [134, 112, 200],
        [56, 45, 189],
        [200, 56, 123],
        [87, 92, 204],
        [120, 56, 123],
        [45, 78, 123],
        [156, 200, 56]
    ]

labels_list = []

with open(r'labels.txt', 'r') as fp:
    for line in fp:
        labels_list.append(line[:-1])

colormap = np.asarray(ade_palette())

def label_to_color_image(label):
    if label.ndim != 2:
        raise ValueError("2-D 입력 라벨이 필요합니다.")

    if np.max(label) >= len(colormap):
        raise ValueError("라벨 값이 너무 큽니다.")
    return colormap[label]

def draw_plot(pred_img, seg):
    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')
    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.numpy().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)
    return fig

def sepia(input_text):
    # 입력 텍스트가 유효한 파일 경로인지 확인
    if not os.path.isfile(input_text):
        return "유효하지 않은 파일 경로입니다. 유효한 이미지 파일 경로를 입력해주세요."

    # 입력 텍스트를 사용하여 이미지를 로드 (이미지 파일 경로로 가정)
    input_img = Image.open(input_text)

    inputs = feature_extractor(images=input_img, return_tensors="tf")
    outputs = model(**inputs)
    logits = outputs.logits

    logits = tf.transpose(logits, [0, 2, 3, 1])
    logits = tf.image.resize(
        logits, input_img.size[::-1]
    )
    seg = tf.math.argmax(logits, axis=-1)[0]

    color_seg = np.zeros(
        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
    )
    for label, color in enumerate(colormap):
        color_seg[seg.numpy() == label, :] = color

    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
    pred_img = pred_img.astype(np.uint8)

    # Gradio에 적합한 형식으로 이미지 배열 변환
    pred_img = Image.fromarray(pred_img)

    return pred_img

# Gradio 인터페이스 정의
iface = gr.Interface(fn=sepia, inputs="image", outputs="image")

# Gradio 앱 시작
iface.launch()