eskayML/mask_segmentation

1import gradio as gr2import torch3import random4import numpy as np5from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation6 7device = torch.device("cpu")8model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade").to(device)9model.eval()10preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")11 12def visualize_instance_seg_mask(mask):13    image = np.zeros((mask.shape[0], mask.shape[1], 3))14    labels = np.unique(mask)15    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}16    for i in range(image.shape[0]):17      for j in range(image.shape[1]):18        image[i, j, :] = label2color[mask[i, j]]19    image = image / 25520    return image21 22def query_image(img):23    target_size = (img.shape[0], img.shape[1])24    inputs = preprocessor(images=img, return_tensors="pt")25    with torch.no_grad():26        outputs = model(**inputs)27    outputs.class_queries_logits = outputs.class_queries_logits.cpu()28    outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()29    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()30    results = torch.argmax(results, dim=0).numpy()31    results = visualize_instance_seg_mask(results)32    return results33 34demo = gr.Interface(35    query_image,36    inputs=[gr.Image()],37    outputs="image",38    title="Image Segmentation using Maskformer",39)40 41demo.launch()