bien.nguyen1 commited on
Commit
e04da3e
1 Parent(s): 307dfdb
Files changed (2) hide show
  1. app.py +3 -3
  2. model.py +3 -3
app.py CHANGED
@@ -5,9 +5,9 @@ from model import SegmentationTool
5
  seg_tool = SegmentationTool()
6
 
7
 
8
- def segment(input_img, masked_items):
9
  mask_image, transparent_mask_image, image, items, room = (
10
- seg_tool.get_mask(image=input_img, masked_items=masked_items))
11
  return mask_image
12
 
13
 
@@ -16,7 +16,7 @@ demo = gr.Interface(fn=segment,
16
  gr.Image(type='pil'),
17
  gr.CheckboxGroup([("Door", 14), ("Window", 8)],
18
  value=[8, 14],
19
- label="Masked Items")],
20
  outputs=['image'],
21
  allow_flagging='never')
22
  if __name__ == "__main__":
 
5
  seg_tool = SegmentationTool()
6
 
7
 
8
+ def segment(input_img, mask_items):
9
  mask_image, transparent_mask_image, image, items, room = (
10
+ seg_tool.get_mask(image=input_img, mask_items=mask_items))
11
  return mask_image
12
 
13
 
 
16
  gr.Image(type='pil'),
17
  gr.CheckboxGroup([("Door", 14), ("Window", 8)],
18
  value=[8, 14],
19
+ label="Mask Items")],
20
  outputs=['image'],
21
  allow_flagging='never')
22
  if __name__ == "__main__":
model.py CHANGED
@@ -68,7 +68,7 @@ class SegmentationTool:
68
 
69
  return mask_image
70
 
71
- def get_mask(self, image_path=None, image=None, masked_items=None):
72
  if image_path:
73
  image = Image.open(image_path)
74
  else:
@@ -83,8 +83,8 @@ class SegmentationTool:
83
 
84
  # mask_items = [0, 3, 5, 8, 14]
85
  # mask_items = [8] # windowpane
86
- if masked_items is None:
87
- masked_items = []
88
  if 73 in label_ids or 50 in label_ids or 61 in label_ids:
89
  # mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 73, 118, 124, 129]
90
  room = 'kitchen'
 
68
 
69
  return mask_image
70
 
71
+ def get_mask(self, image_path=None, image=None, mask_items=None):
72
  if image_path:
73
  image = Image.open(image_path)
74
  else:
 
83
 
84
  # mask_items = [0, 3, 5, 8, 14]
85
  # mask_items = [8] # windowpane
86
+ if mask_items is None:
87
+ mask_items = []
88
  if 73 in label_ids or 50 in label_ids or 61 in label_ids:
89
  # mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 73, 118, 124, 129]
90
  room = 'kitchen'