generate 3 outputs

app.py

@@ -13,9 +13,9 @@ def predict(image, style, backyard_style, color_preference, additional_info):
 
     output_images, transparent_mask_image = vs_tool.virtual_stage(
         image=init_image, style=style, color_preference=color_preference,
-        additional_info=additional_info, number_images=1)
+        additional_info=additional_info, number_images=3)
 
-    return output_images[0], transparent_mask_image
+    return output_images[0], output_images[1], output_images[2], transparent_mask_image
 
 
 image_blocks = gr.Blocks()
@@ -65,9 +65,12 @@ 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 = gr.Image(label="Output", elem_id="output-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)
 
-            btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info], outputs=[image_out, mask_image])
+                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])
 
             gr.Markdown("## Image Examples")
             gr.Examples(

model.py

@@ -93,6 +93,7 @@ class VirtualStagingToolV2():
         elif 7 in label_ids:
             room = 'bedroom'
         elif 23 in label_ids or 49 in label_ids:
+            mask_items = [0, 3, 5, 8, 14, 49]
             room = 'living room'
         elif 15 in label_ids and 19 in label_ids:
             room = 'dining room'
@@ -135,7 +136,7 @@ class VirtualStagingToolV2():
                      i in ['shower', 'bathtub', 'chest of drawers', 'counter', 'countertop', 'sink']]
         elif room == 'living room':
             items = [i for i in items if
-                     i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'fireplace', 'coffee table']]
+                     i in ['table', 'sofa', 'chest of drawers', 'armchair', 'cabinet', 'coffee table']]
         elif room == 'dining room':
             items = [i for i in items if i in ['table', 'chair', 'cabinet']]