add masked items option

app.py

@@ -5,13 +5,18 @@ from model import SegmentationTool
 seg_tool = SegmentationTool()
 
 
-def segment(input_img):
-    mask_image, transparent_mask_image, image, items, room = seg_tool.get_mask(image=input_img)
+def segment(input_img, masked_items):
+    mask_image, transparent_mask_image, image, items, room = (
+        seg_tool.get_mask(image=input_img, masked_items=masked_items))
     return mask_image
 
 
 demo = gr.Interface(fn=segment,
-                    inputs=gr.Image(type='pil'),
+                    inputs=[
+                        gr.Image(type='pil'),
+                        gr.CheckboxGroup([("Door", 14), ("Window", 8)],
+                                         value=[8, 14],
+                                         label="Masked Items")],
                     outputs=['image'],
                     allow_flagging='never')
 if __name__ == "__main__":

model.py

@@ -68,7 +68,7 @@ class SegmentationTool:
 
         return mask_image
 
-    def get_mask(self, image_path=None, image=None):
+    def get_mask(self, image_path=None, image=None, masked_items=None):
         if image_path:
             image = Image.open(image_path)
         else:
@@ -76,13 +76,15 @@ class SegmentationTool:
                 raise ValueError("no image provided")
 
         # display(image)
+        # print(image)
         prediction = self._predict(image)
 
         label_ids = np.unique(prediction)
 
         # mask_items = [0, 3, 5, 8, 14]
-        mask_items = [8] # windowpane
-
+        # mask_items = [8] # windowpane
+        if masked_items is None:
+            masked_items = []
         if 73 in label_ids or 50 in label_ids or 61 in label_ids:
             # mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 73, 118, 124, 129]
             room = 'kitchen'