with text/semantic labels

app.py

@@ -3,39 +3,38 @@ import numpy as np
 from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
 
 
-# preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")
-# model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade")
-preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-coco")
-model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-large-coco")
+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):
+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]
-    labels = np.unique(mask)
+    label_ids = np.unique(mask)
 
-    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}
+    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, :] = label2color[mask[i, j]]
-        label2count[mask[i, j]] = label2count[mask[i, j]] + 1
+        image[i, j, :] = id2color[mask[i, j]]
+        id2count[mask[i, j]] = id2count[mask[i, j]] + 1
 
     image = image / 255
 
-    for k, v in label2count.items():
-        label2count[k] = v / image_total_pixels
+    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 = preprocessor(images=img, return_tensors="pt")
+    inputs = feature_extractor(images=img, return_tensors="pt")
     outputs = model(**inputs)
-    results = preprocessor.post_process_semantic_segmentation(outputs=outputs, target_sizes=[img_size])[0]
-    results = visualize_instance_seg_mask(results.numpy())
+    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