import gradio as gr
from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
from PIL import Image
import numpy as np


# 모델과 특징 추출기 불러오기
feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-512-1024")
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-512-1024")

colors = np.array([
    [255, 0, 0],      # 빨강
    [255, 228, 0],    # 노랑
    [171, 242, 0],    # 연두
    [0, 216, 255],    # 하늘
    [0, 0, 255],      # 파랑
    [255, 0, 221],    # 핑크
    [116, 116, 116],  # 회색
    [95, 0, 255],     # 보라
    [255, 94, 0],     # 주황
    [71, 200, 62],    # 초록
    [153, 0, 76],     # 마젠타
    [67, 116, 217],   # 애매한하늘+파랑
    [153, 112, 0],    # 겨자
    [87, 129, 0],     # 녹색
    [255, 169, 169],  # 분홍분홍
    [35, 30, 183],    # 어두운 파랑
    [225, 186, 133],  # 살색
    [206, 251, 201],  # 연한초록
    [165, 102, 255]   # 애매한보라
], dtype=np.uint8)

def segment_image(image):
    # 이미지를 처리하고 모델에 전달하기
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits

    # 결과 처리 및 NumPy 배열로 변환
    result = logits.argmax(dim=1)[0]
    result = result.cpu().detach().numpy()

    # 색상 팔레트 적용
    result_color = colors[result]

    # NumPy 배열을 PIL 이미지로 변환
    result_image = Image.fromarray(result_color.astype(np.uint8))

    # 원본 이미지와 추론 결과 이미지의 크기 일치시키기
    result_image = result_image.resize(image.size, Image.NEAREST)

    # 원본 이미지와 추론 결과 이미지 결합
    combined_image = Image.blend(image.convert("RGBA"), result_image.convert("RGBA"), alpha=0.5)

    return combined_image

# Gradio 인터페이스 정의
iface = gr.Interface(
    fn=segment_image,
    inputs=gr.inputs.Image(type='pil'),
    examples = ['image1.jpg', 'image2.jpg', 'image3.jpg'],
    outputs= 'image',
    title="SegFormer Image Segmentation",
    description="Upload an image to segment it using the SegFormer model trained on Cityscapes dataset."
)

# 인터페이스 실행
iface.launch()