resize image

app.py

@@ -4,9 +4,7 @@ import os
 from model import VirtualStagingToolV2
 
 
-def predict(image, style, backyard_style, color_preference, additional_info):
-    init_image = image.convert("RGB").resize((512, 512))
-
+def predict(init_image, style, backyard_style, color_preference, additional_info):
     vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting")
     if backyard_style:
         style = backyard_style
@@ -83,10 +81,10 @@ with image_blocks as demo:
                     )
 
                 with gr.Column():
-                    mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=500)
-                    image_out_1 = gr.Image(label="Output 1", elem_id="output-img-1", type="pil").style(height=512)
-                    image_out_2 = gr.Image(label="Output 2", elem_id="output-img-2", type="pil").style(height=512)
-                    image_out_3 = gr.Image(label="Output 3", elem_id="output-img-3", type="pil").style(height=512)
+                    mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=400)
+                    image_out_1 = gr.Image(label="Output 1", elem_id="output-img-1", type="pil").style(height=400)
+                    image_out_2 = gr.Image(label="Output 2", elem_id="output-img-2", type="pil").style(height=400)
+                    image_out_3 = gr.Image(label="Output 3", elem_id="output-img-3", type="pil").style(height=400)
 
                 btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info],
                           outputs=[image_out_1, image_out_2, image_out_3, mask_image])

model.py

@@ -3,6 +3,7 @@ import matplotlib.pyplot as plt
 import matplotlib.patches as mpatches
 from matplotlib import cm
 
+import cv2
 from PIL import Image
 import numpy as np
 
@@ -22,8 +23,12 @@ class VirtualStagingToolV2():
         self.segmentation_version = segmentation_version
         self.diffusion_version = diffusion_version
 
-        self.feature_extractor = AutoImageProcessor.from_pretrained(self.segmentation_version)
-        self.segmentation_model = UperNetForSemanticSegmentation.from_pretrained(self.segmentation_version)
+        if segmentation_version == "openmmlab/upernet-convnext-tiny":
+            self.feature_extractor = AutoImageProcessor.from_pretrained(self.segmentation_version)
+            self.segmentation_model = UperNetForSemanticSegmentation.from_pretrained(self.segmentation_version)
+        elif segmentation_version == "nvidia/segformer-b5-finetuned-ade-640-640":
+            self.feature_extractor = SegformerFeatureExtractor.from_pretrained(self.segmentation_version)
+            self.segmentation_model = SegformerForSemanticSegmentation.from_pretrained(self.segmentation_version)
 
         self.diffution_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
             self.diffusion_version,
@@ -44,6 +49,18 @@ class VirtualStagingToolV2():
         mask[np.isin(prediction_array, mask_items)] = 0
         mask[~np.isin(prediction_array, mask_items)] = 255
 
+        buffer_size = 10
+
+        # Dilate the binary image
+        kernel = np.ones((buffer_size, buffer_size), np.uint8)
+        dilated_image = cv2.dilate(mask, kernel, iterations=1)
+
+        # Subtract the original binary image
+        buffer_area = dilated_image - mask
+
+        # Apply buffer area to the original image
+        mask = cv2.bitwise_or(mask, buffer_area)
+
         #     # # Create a PIL Image object from the mask
         mask_image = Image.fromarray(mask, mode='L')
         # display(mask_image)
@@ -84,7 +101,7 @@ class VirtualStagingToolV2():
             mask_items = [1, 2, 4, 25, 32]
             room = 'backyard'
         elif 73 in label_ids or 50 in label_ids or 61 in label_ids:
-            mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 118, 124, 129
+            mask_items = [0, 3, 5, 8, 14, 50, 61, 71, 73, 118, 124, 129
                           ]
             room = 'kitchen'
         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():
 
 
         if room == 'kitchen':
-            items = [i for i in items if i in ['kitchen island', 'cabinet', 'shelf', 'counter', 'countertop', 'stool']]
+            items = [i for i in items if i in ['cabinet', 'shelf', 'counter', 'countertop', 'stool']]
         elif room == 'bedroom':
             items = [i for i in items if i in ['bed ', 'table', 'chest of drawers', 'desk', 'armchair', 'wardrobe']]
         elif room == 'bathroom':
             items = [i for i in items if
-                     i in ['shower', 'bathtub', 'chest of drawers', 'counter', 'countertop', 'sink']]
+                     i in ['shower', 'bathtub', 'screen door', 'cabinet']]
         elif room == 'living room':
             items = [i for i in items if
                      i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'coffee table']]

requirements.txt

@@ -4,4 +4,5 @@ diffusers==0.16.1
 accelerate==0.19.0
 matplotlib==3.6.2
 pillow==9.2.0
-numpy==1.23.2
+numpy==1.23.2
+opencv-python==4.7.0.72