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

# Segformer 모델 및 feature extractor 불러오기
feature_extractor = SegformerFeatureExtractor.from_pretrained(
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
model = TFSegformerForSemanticSegmentation.from_pretrained(
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")

# 모델 예측 함수를 정의합니다.
def classify_image(img):
    # 이미지를 전처리합니다.
    inputs = feature_extractor(images=img, return_tensors="tf")
    
    # 모델로 예측을 수행합니다.
    predictions = model(**inputs)
    
    # 예측 결과 중에서 가장 높은 확률을 가진 클래스를 선택합니다.
    predicted_label = tf.argmax(predictions.logits[0], axis=-1).numpy()

    # 라벨을 반환합니다.
    return predicted_label

# Gradio UI를 생성합니다.
iface = gr.Interface(fn=classify_image,
                     inputs="image",
                     outputs="label", live=True)

# Gradio UI를 시작합니다.
iface.launch()