Update app.py

app.py

@@ -11,14 +11,30 @@ preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-s
 
 def visualize_instance_seg_mask(mask):
     print(mask)
+    print(mask.shape)
     image = np.zeros((mask.shape[0], mask.shape[1], 3))
     labels = np.unique(mask)
     print("================unquie labels")
-    print(labels)
+    wall=[]
+    floor=[]
+    window=[]
+    other=[]
     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]]
+            if mask[i, j]==0:
+                wall.append([i,j])
+            
+            elif mask[i, j]==3
+                 floor.append([i,j])
+            
+            elif mask[i, j]==8
+                window.append([i,j])
+            else:
+                other.append([i,j])
+
+            image[i, j, :] = label2color[mask[i, j]]
+    print({"floor":floor,"window":window,"wall":wall,"other":other})        
     image = image / 255
     return image