import gradio as gr
import requests
import torch
from PIL import Image
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation


def greet(url):

    processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-large-cityscapes-semantic")
    model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-large-cityscapes-semantic")

    image = Image.open(requests.get(url, stream=True).raw)
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
    # model predicts class_queries_logits of shape `(batch_size, num_queries)`
    # and masks_queries_logits of shape `(batch_size, num_queries, height, width)`
    class_queries_logits = outputs.class_queries_logits
    masks_queries_logits = outputs.masks_queries_logits

    # you can pass them to processor for postprocessing
    predicted_semantic_map = processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]

    color_map = {
        0: (0, 0, 0),  # 클래스 0: 검은색
        1: (255, 255, 255),  # 클래스 1: 흰색
        2: (255, 0, 0),
        3: (0, 255, 0),
        4: (0, 0, 255),
        5: (255, 255, 0),
        6: (255, 0, 255),
        7: (0, 255, 255),
        # 다른 클래스에 대한 색상 지정
    }
    #semantic_image = Image.new('RGB', predicted_semantic_map.shape[1:][::-1])[0]
    semantic_image = Image.new('RGB', (predicted_semantic_map.shape[1], predicted_semantic_map.shape[0]))
    pixels = semantic_image.load()
    for y in range(predicted_semantic_map.shape[0]):
        for x in range(predicted_semantic_map.shape[1]):
            class_id = predicted_semantic_map[y, x].item()
            color = color_map.get(class_id, (0, 0, 0))
            pixels[x, y] = color

    return pixels


url = "https://www.cityscapes-dataset.com/wordpress/wp-content/uploads/2015/07/muenster00.png"
greet(url)

iface = gr.Interface(
    fn=greet,
    inputs=gr.Image(value=url),
    outputs="image"
)

iface.launch(share=True)