sandrawang1031 commited on
Commit
c0b640d
1 Parent(s): 4b406e7

resize image

Browse files
Files changed (3) hide show
  1. app.py +5 -7
  2. model.py +22 -5
  3. requirements.txt +2 -1
app.py CHANGED
@@ -4,9 +4,7 @@ import os
4
  from model import VirtualStagingToolV2
5
 
6
 
7
- def predict(image, style, backyard_style, color_preference, additional_info):
8
- init_image = image.convert("RGB").resize((512, 512))
9
-
10
  vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting")
11
  if backyard_style:
12
  style = backyard_style
@@ -83,10 +81,10 @@ with image_blocks as demo:
83
  )
84
 
85
  with gr.Column():
86
- mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=500)
87
- image_out_1 = gr.Image(label="Output 1", elem_id="output-img-1", type="pil").style(height=512)
88
- image_out_2 = gr.Image(label="Output 2", elem_id="output-img-2", type="pil").style(height=512)
89
- image_out_3 = gr.Image(label="Output 3", elem_id="output-img-3", type="pil").style(height=512)
90
 
91
  btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info],
92
  outputs=[image_out_1, image_out_2, image_out_3, mask_image])
 
4
  from model import VirtualStagingToolV2
5
 
6
 
7
+ def predict(init_image, style, backyard_style, color_preference, additional_info):
 
 
8
  vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting")
9
  if backyard_style:
10
  style = backyard_style
 
81
  )
82
 
83
  with gr.Column():
84
+ mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=400)
85
+ image_out_1 = gr.Image(label="Output 1", elem_id="output-img-1", type="pil").style(height=400)
86
+ image_out_2 = gr.Image(label="Output 2", elem_id="output-img-2", type="pil").style(height=400)
87
+ image_out_3 = gr.Image(label="Output 3", elem_id="output-img-3", type="pil").style(height=400)
88
 
89
  btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info],
90
  outputs=[image_out_1, image_out_2, image_out_3, mask_image])
model.py CHANGED
@@ -3,6 +3,7 @@ import matplotlib.pyplot as plt
3
  import matplotlib.patches as mpatches
4
  from matplotlib import cm
5
 
 
6
  from PIL import Image
7
  import numpy as np
8
 
@@ -22,8 +23,12 @@ class VirtualStagingToolV2():
22
  self.segmentation_version = segmentation_version
23
  self.diffusion_version = diffusion_version
24
 
25
- self.feature_extractor = AutoImageProcessor.from_pretrained(self.segmentation_version)
26
- self.segmentation_model = UperNetForSemanticSegmentation.from_pretrained(self.segmentation_version)
 
 
 
 
27
 
28
  self.diffution_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
29
  self.diffusion_version,
@@ -44,6 +49,18 @@ class VirtualStagingToolV2():
44
  mask[np.isin(prediction_array, mask_items)] = 0
45
  mask[~np.isin(prediction_array, mask_items)] = 255
46
 
 
 
 
 
 
 
 
 
 
 
 
 
47
  # # # Create a PIL Image object from the mask
48
  mask_image = Image.fromarray(mask, mode='L')
49
  # display(mask_image)
@@ -84,7 +101,7 @@ class VirtualStagingToolV2():
84
  mask_items = [1, 2, 4, 25, 32]
85
  room = 'backyard'
86
  elif 73 in label_ids or 50 in label_ids or 61 in label_ids:
87
- mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 118, 124, 129
88
  ]
89
  room = 'kitchen'
90
  elif 37 in label_ids or 65 in label_ids or (27 in label_ids and 47 in label_ids and 70 in label_ids):
@@ -128,12 +145,12 @@ class VirtualStagingToolV2():
128
 
129
 
130
  if room == 'kitchen':
131
- items = [i for i in items if i in ['kitchen island', 'cabinet', 'shelf', 'counter', 'countertop', 'stool']]
132
  elif room == 'bedroom':
133
  items = [i for i in items if i in ['bed ', 'table', 'chest of drawers', 'desk', 'armchair', 'wardrobe']]
134
  elif room == 'bathroom':
135
  items = [i for i in items if
136
- i in ['shower', 'bathtub', 'chest of drawers', 'counter', 'countertop', 'sink']]
137
  elif room == 'living room':
138
  items = [i for i in items if
139
  i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'coffee table']]
 
3
  import matplotlib.patches as mpatches
4
  from matplotlib import cm
5
 
6
+ import cv2
7
  from PIL import Image
8
  import numpy as np
9
 
 
23
  self.segmentation_version = segmentation_version
24
  self.diffusion_version = diffusion_version
25
 
26
+ if segmentation_version == "openmmlab/upernet-convnext-tiny":
27
+ self.feature_extractor = AutoImageProcessor.from_pretrained(self.segmentation_version)
28
+ self.segmentation_model = UperNetForSemanticSegmentation.from_pretrained(self.segmentation_version)
29
+ elif segmentation_version == "nvidia/segformer-b5-finetuned-ade-640-640":
30
+ self.feature_extractor = SegformerFeatureExtractor.from_pretrained(self.segmentation_version)
31
+ self.segmentation_model = SegformerForSemanticSegmentation.from_pretrained(self.segmentation_version)
32
 
33
  self.diffution_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
34
  self.diffusion_version,
 
49
  mask[np.isin(prediction_array, mask_items)] = 0
50
  mask[~np.isin(prediction_array, mask_items)] = 255
51
 
52
+ buffer_size = 10
53
+
54
+ # Dilate the binary image
55
+ kernel = np.ones((buffer_size, buffer_size), np.uint8)
56
+ dilated_image = cv2.dilate(mask, kernel, iterations=1)
57
+
58
+ # Subtract the original binary image
59
+ buffer_area = dilated_image - mask
60
+
61
+ # Apply buffer area to the original image
62
+ mask = cv2.bitwise_or(mask, buffer_area)
63
+
64
  # # # Create a PIL Image object from the mask
65
  mask_image = Image.fromarray(mask, mode='L')
66
  # display(mask_image)
 
101
  mask_items = [1, 2, 4, 25, 32]
102
  room = 'backyard'
103
  elif 73 in label_ids or 50 in label_ids or 61 in label_ids:
104
+ mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 73, 118, 124, 129
105
  ]
106
  room = 'kitchen'
107
  elif 37 in label_ids or 65 in label_ids or (27 in label_ids and 47 in label_ids and 70 in label_ids):
 
145
 
146
 
147
  if room == 'kitchen':
148
+ items = [i for i in items if i in ['cabinet', 'shelf', 'counter', 'countertop', 'stool']]
149
  elif room == 'bedroom':
150
  items = [i for i in items if i in ['bed ', 'table', 'chest of drawers', 'desk', 'armchair', 'wardrobe']]
151
  elif room == 'bathroom':
152
  items = [i for i in items if
153
+ i in ['shower', 'bathtub', 'screen door', 'cabinet']]
154
  elif room == 'living room':
155
  items = [i for i in items if
156
  i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'coffee table']]
requirements.txt CHANGED
@@ -4,4 +4,5 @@ diffusers==0.16.1
4
  accelerate==0.19.0
5
  matplotlib==3.6.2
6
  pillow==9.2.0
7
- numpy==1.23.2
 
 
4
  accelerate==0.19.0
5
  matplotlib==3.6.2
6
  pillow==9.2.0
7
+ numpy==1.23.2
8
+ opencv-python==4.7.0.72