add backyard

app.py

@@ -4,13 +4,17 @@ import os
 from model import VirtualStagingToolV2
 
 
-def predict(image, style, color_preference):
+def predict(image, style, backyard_style, color_preference):
     init_image = image.convert("RGB").resize((512, 512))
 
     vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting")
+    if backyard_style:
+        style = backyard_style
+
     output_images, transparent_mask_image = vs_tool.virtual_stage(
         image=init_image, style=style, color_preference=color_preference, number_images=1)
-    return output_images[0], transparent_mask_image, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+
+    return output_images[0], transparent_mask_image
 
 
 image_blocks = gr.Blocks()
@@ -31,21 +35,34 @@ with image_blocks as demo:
                             ["Bohemian", "Coastal", "Contemporary", "Farmhouse", "French country",
                              "Glam", "Industrial", "Japandi", "Mid-century modern", "Minimal", "Modern",
                              "Rustic", "Scandinavian", "Traditional", "Transitional", ],
-                            label="Design theme", elem_id="input-style"
+                            label="Design theme", elem_id="input-style",
+                            info="only select if the image is NOT a backyard"
                         )
 
-                        color_preference = gr.Textbox(placeholder='Enter color preference',
-                                                      label="Color preference", elem_id="input-color")
-                        btn = gr.Button("Inpaint!").style(
-                            margin=False,
-                            rounded=(False, True, True, False),
-                            full_width=False,
+                        backyard_style = gr.Dropdown(
+                            ["Beautiful garden", "Charming playhouse garden landscaping",
+                             "Cottage garden's colorful planting palette", "Cozy corner with fire pit and seating",
+                             "Garden landscaping with gravel landscaping", "Hip california garden landscaping",
+                             "Lush green lawn", "Mediterranean garden landscaping",
+                             "Moss garden", "Outdoor dining and sitting area",
+                             "Party-ready outdoor space with pool, spa, and fire feature",
+                             "Resort-style landscaping and pool", "Round swimming pool with lawn and pool house"],
+                            label="Backyard design theme", elem_id="input-backyard-style",
+                            info="only select if the image is a backyard"
                         )
+
+                    color_preference = gr.Textbox(placeholder='Enter color preference',
+                                                  label="Color preference", elem_id="input-color")
+                    btn = gr.Button("Inpaint!").style(
+                        margin=False,
+                        rounded=(False, True, True, False),
+                        full_width=False,
+                    )
                 with gr.Column():
-                    mask_image = gr.Image(label="Mask image", elem_id="mask-img").style(height=400)
-                    image_out = gr.Image(label="Output", elem_id="output-img").style(height=400)
+                    mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=400)
+                    image_out = gr.Image(label="Output", elem_id="output-img", type="pil").style(height=400)
 
-            btn.click(fn=predict, inputs=[image, style, color_preference], outputs=[image_out, mask_image])
+            btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference], outputs=[image_out, mask_image])
 
             gr.Markdown("## Image Examples")
             gr.Examples(
@@ -56,7 +73,9 @@ with image_blocks as demo:
                           os.path.join(os.path.dirname(__file__),
                                        "examples/oct-2019-idh-bathroom-reno-ideas-new-gallery-2.jpg"),
                           os.path.join(os.path.dirname(__file__),
-                                       "examples/tips-for-decorating-a-beautiful-bedroom-1976169-hero-e960fbb8311c4b9b875a1813962d34eb.jpg")
+                                       "examples/tips-for-decorating-a-beautiful-bedroom-1976169-hero-e960fbb8311c4b9b875a1813962d34eb.jpg"),
+                          os.path.join(os.path.dirname(__file__),
+                                       "examples/backyard_additions.jpg")
                           ],
                 inputs=image
             )

model.py

@@ -120,8 +120,8 @@ class VirtualStagingToolV2():
         mask_image, transparent_mask_image, init_image, items, room = self.get_mask(image_path, image)
         if not style:
             raise ValueError('style not provided.')
-        if not color_preference:
-            raise ValueError('color_preference not provided.')
+        if color_preference:
+            color_preference = f"in {color_preference}"
 
         if room == 'kitchen':
             items = [i for i in items if i in ['kitchen island', 'cabinet', 'shelf', 'counter', 'countertop', 'stool']]
@@ -130,9 +130,16 @@ class VirtualStagingToolV2():
         elif room == 'bathroom':
             items = [i for i in items if
                      i in ['shower', 'bathtub', 'chest of drawers', 'counter', 'countertop', 'sink']]
+        elif room == 'livingroom':
+            items = [i for i in items if
+                     i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'fireplace', 'coffee table']]
 
         items = ', '.join(items)
-        prompt = f'{items}, high resolution, in the {style} style {room} in {color_preference}'
+
+        if room == 'backyard':
+            prompt = f'realistic, high resolution, {room} with {style}'
+        else:
+            prompt = f'realistic {items}, high resolution, in the {style} style {room} {color_preference}'
         print(prompt)
 
         output_images = self._edit_image(init_image, mask_image, prompt, number_images)