import gradio as gr
import tensorflow as tf
from PIL import Image
import requests

# 모델 로드
model = tf.saved_model.load("nvidia_segformer_b5_finetuned_cityscapes_1024")

# 레이블 및 색상 정의
label_colors = {
    "road": [204, 87, 92],
    "sidewalk": [112, 185, 212],
    "building": [196, 160, 122],
    "wall": [106, 135, 242],
    "fence": [91, 192, 222],
    "pole": [255, 192, 203],
    "traffic_light": [176, 224, 230],
    "traffic_sign": [222, 49, 99],
    "vegetation": [139, 69, 19],
    "terrain": [255, 0, 0],
    "sky": [0, 0, 255],
    "person": [255, 228, 181],
    "rider": [128, 0, 0],
    "car": [0, 128, 0],
    "truck": [255, 99, 71],
    "bus": [0, 255, 0],
    "train": [128, 0, 128],
    "motorcycle": [255, 255, 0],
    "bicycle": [128, 0, 128]
}

# Gradio 인터페이스 정의
iface = gr.Interface(
    fn=lambda image: predict_segmentation(image, model),
    inputs="image",
    outputs="image"
)
iface.launch()

# 이미지 세그멘테이션 함수 정의
def predict_segmentation(image, model):
    # 이미지 변환
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    image = image.resize((1024, 1024))  # 모델의 입력 크기에 맞게 조절
    image_array = tf.keras.preprocessing.image.img_to_array(image)
    image_array = tf.expand_dims(image_array, 0)

    # 모델 추론
    predictions = model(image_array)["output_0"]

    # 레이블별 색상 매핑
    segmented_image = tf.zeros_like(predictions)
    for label, color in label_colors.items():
        mask = tf.reduce_all(tf.equal(predictions, color), axis=-1, keepdims=True)
        for i in range(3):
            segmented_image += tf.cast(mask, tf.float32) * tf.constant(color[i], dtype=tf.float32)

    # 이미지 리턴
    segmented_image = tf.cast(segmented_image, tf.uint8)
    segmented_image = tf.image.resize(segmented_image, [image.height, image.width])
    return segmented_image.numpy()