simplifies query_image. uses non-deprecated function

app.py

@@ -1,6 +1,4 @@
 import gradio as gr
-import torch
-import random
 import numpy as np
 from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
 
@@ -10,23 +8,34 @@ from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmen
 preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
 model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
 
+
 def visualize_instance_seg_mask(mask):
     image = np.zeros((mask.shape[0], mask.shape[1], 3))
+    image_total_pixels = mask.shape[0] * mask.shape[1]
     labels = np.unique(mask)
-    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
+
+    label2color = {label: (np.random.randint(0, 2), np.random.randint(0, 256), np.random.randint(0, 256)) for label in labels}
+    label2count = {label: 0 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]]
+        label2count[mask[i, j]] = label2count[mask[i, j]] + 1
+
     image = image / 255
+
+    for k, v in label2count.items():
+        label2count[k] = v / image_total_pixels
+
     return image
 
+
 def query_image(img):
-    target_size = (img.shape[0], img.shape[1])
+    img_size = (img.shape[0], img.shape[1])
     inputs = preprocessor(images=img, return_tensors="pt")
     outputs = model(**inputs)
-    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0]
-    results = torch.argmax(results, dim=0).numpy()
-    results = visualize_instance_seg_mask(results)
+    results = preprocessor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
+    results = visualize_instance_seg_mask(results.numpy())
     return results
 
 
@@ -34,7 +43,7 @@ demo = gr.Interface(
     query_image,
     inputs=[gr.Image()],
     outputs="image",
-    title="maskformer-swin-tiny-ade results",
+    title="maskformer-swin-large-coco results",
     allow_flagging="never",
     analytics_enabled=None
 )