import gradio as gr
import numpy as np
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation


feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-coco")
model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-coco")
# feature_extractor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
# model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")

def visualize_instance_seg_mask(mask, id2label):
    image = np.zeros((mask.shape[0], mask.shape[1], 3))
    image_total_pixels = mask.shape[0] * mask.shape[1]
    label_ids = np.unique(mask)

    id2color = {id: (np.random.randint(0, 2), np.random.randint(0, 256), np.random.randint(0, 256)) for id in label_ids}
    id2count = {id: 0 for id in label_ids}

    for i in range(image.shape[0]):
      for j in range(image.shape[1]):
        image[i, j, :] = id2color[mask[i, j]]
        id2count[mask[i, j]] = id2count[mask[i, j]] + 1

    image = image / 255

    label2count = {id2label[id]: id2count[id] / image_total_pixels for id in label_ids}
    print(label2count)

    return image


def query_image(img):
    img_size = (img.shape[0], img.shape[1])
    inputs = feature_extractor(images=img, return_tensors="pt")
    outputs = model(**inputs)
    results = feature_extractor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
    results = visualize_instance_seg_mask(results.numpy(), model.config.id2label)
    return results


demo = gr.Interface(
    query_image,
    inputs=[gr.Image()],
    outputs="image",
    title="maskformer-swin-large-coco results",
    allow_flagging="never",
    analytics_enabled=None
)

demo.launch(show_api=False)