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

# 모델과 feature extractor 로드
model_name = "nvidia/segformer-b0-finetuned-ade-512-512"
model = SegformerForSemanticSegmentation.from_pretrained(model_name)
feature_extractor = SegformerFeatureExtractor.from_pretrained(model_name)

def segment_image(image):
    # 이미지 처리
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    # 마스크 생성
    upsampled_logits = torch.nn.functional.interpolate(
        outputs.logits, size=image.size[::-1], mode="bilinear", align_corners=False
    )
    upsampled_predictions = upsampled_logits.argmax(dim=1)
    mask = upsampled_predictions.squeeze().numpy()

    # 결과 반환
    return Image.fromarray(np.uint8(mask * 255))

# 예시 이미지 경로
example_images = ["image1.jpg", "image2.jpg", "image3.jpg"]

# Gradio 인터페이스 설정
demo = gr.Interface(
    fn=segment_image,
    inputs=gr.inputs.Image(type="pil"),
    outputs="image",
    title="머신러닝 7주차 과제_3",
    examples=example_images
)

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