Update app.py

app.py

@@ -1,53 +1,50 @@
 # URL: https://huggingface.co/spaces/gradio/image_segmentation/
 # imports
 import gradio as gr
-from transformers import DetrFeatureExtractor, DetrForSegmentation
-from PIL import Image
-import numpy as np
 import torch
-import torchvision
-import itertools
-import seaborn as sns
+import random
+import numpy as np
+from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
+
 
-# load model from hugging face
-feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50-panoptic')
-model = DetrForSegmentation.from_pretrained('facebook/detr-resnet-50-panoptic')
+# load model
+device = torch.device("cpu")
+model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade").to(device)
+model.eval()
+preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")
 
-def predict_animal_mask(im,
-                        gr_slider_confidence):
-    image = Image.fromarray(im) 
-    image = image.resize((200,200))
-    encoding = feature_extractor(images=image, return_tensors="pt") 
-    outputs = model(**encoding)
-    logits = outputs.logits
-    bboxes = outputs.pred_boxes
-    masks = outputs.pred_masks 
-    prob_per_query = outputs.logits.softmax(-1)[..., :-1].max(-1)[0] 
-    keep = prob_per_query > gr_slider_confidence/100.0
-    label_per_pixel = torch.argmax(masks[keep].squeeze(),dim=0).detach().numpy() 
-    color_mask = np.zeros(image.size+(3,))
-    palette = itertools.cycle(sns.color_palette())
-    for lbl in np.unique(label_per_pixel): 
-        color_mask[label_per_pixel==lbl,:] = np.asarray(next(palette))*255
-    pred_img = np.array(image.convert('RGB'))*0.25 + color_mask*0.75
-    pred_img = pred_img.astype(np.uint8)   
-    return pred_img
+# define core and helper fns
+def visualize_instance_seg_mask(mask):
+    image = np.zeros((mask.shape[0], mask.shape[1], 3))
+    labels = np.unique(mask)
+    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
+    for i in range(image.shape[0]):
+      for j in range(image.shape[1]):
+        image[i, j, :] = label2color[mask[i, j]]
+    image = image / 255
+    return image
 
+def query_image(img):
+    target_size = (img.shape[0], img.shape[1])
+    inputs = preprocessor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    outputs.class_queries_logits = outputs.class_queries_logits.cpu()
+    outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()
+    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()
+    results = torch.argmax(results, dim=0).numpy()
+    results = visualize_instance_seg_mask(results)
+    return results
 
-# define inputs
-gr_image_input = gr.inputs.Image()
-gr_slider_confidence = gr.inputs.Slider(0,100,5,85,
-                                        label='Set confidence threshold for masks')
-# define output
-gr_image_output = gr.outputs.Image() 
+# define interface 
 
-# define interface
-demo = gr.Interface(predict_animal_mask, 
-                inputs = [gr_image_input,gr_slider_confidence],
-                outputs = gr_image_output,
-                title = 'Image segmentation with varying confidence',
-                description = "A panoptic (semantic+instance) segmentation webapp using DETR (End-to-End Object Detection) model with ResNet-50 backbone",
-                examples=[["cheetah.jpg", 75], ["lion.jpg", 85]])
+demo = gr.Interface(
+    query_image, 
+    inputs=[gr.Image()], 
+    outputs="image",
+    title="MaskFormer Demo",
+    examples=["example_1.png", "example_2.png"]
+)
 
-# launch
-demo.launch()
+# launch 
+demo.launch()