import gradio as gr
from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
from PIL import Image
import torch

# 모델과 feature extractor 로드
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")
feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b3-finetuned-cityscapes-1024-1024")

# 이미지를 처리하는 함수
def predict(image):
    # 이미지를 모델에 맞게 변환
    processed_image = image.resize((1024, 1024))
    inputs = feature_extractor(images=processed_image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits

    # 결과를 이미지로 변환 (예: 가장 높은 확률을 가진 클래스 선택)
    result = torch.argmax(logits)
    result = result.squeeze().cpu().numpy()

    # 여기에서는 단순화를 위해 결과를 그대로 반환합니다.
    # 실제로는 결과를 적절한 형식으로 변환해야 할 수 있습니다.
    return result

# Gradio 인터페이스 생성
demo = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(type='pil'),
    outputs=gr.outputs.Image(type='pil'),
    examples=["image1.jpg", "image2.jpg", "image3.jpg"]  # 세 개의 예제 이미지 경로
)

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