import gradio as gr
from transformers import AutoImageProcessor, BeitForSemanticSegmentation
from PIL import Image
import torch
import numpy as np

# 모델과 이미지 프로세서를 로드합니다.
processor = AutoImageProcessor.from_pretrained("microsoft/beit-base-finetuned-ade-640-640")
model = BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640")

def predict(image):
    # 이미지를 전처리합니다.
    inputs = processor(images=image, return_tensors="pt")
    
    # 모델을 사용하여 예측을 수행합니다.
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    # 결과 처리: 가장 높은 점수를 갖는 클래스 찾기
    seg = logits.argmax(dim=1).squeeze().cpu().numpy()
    
    # 결과를 시각화하기 위해 색상을 할당합니다.
    unique_classes = np.unique(seg)
    seg_colored = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    
    for cls in unique_classes:
        mask = seg == cls
        seg_colored[mask] = np.random.randint(0, 255, size=3)
    
    return Image.fromarray(seg_colored)

# Gradio 인터페이스 설정
gr_interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Image Segmentation with BEiT",
    description="Upload an image to perform segmentation using the microsoft/beit-base-finetuned-ade-640-640 model."
)

if __name__ == "__main__":
    gr_interface.launch()