virtual-staging / app.py
sandrawang1031's picture
resize image
c0b640d
import gradio as gr
import os
from model import VirtualStagingToolV2
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
output_images, transparent_mask_image = vs_tool.virtual_stage(
image=init_image, style=style, color_preference=color_preference,
additional_info=additional_info, number_images=3)
return output_images[0], output_images[1], output_images[2], transparent_mask_image
image_blocks = gr.Blocks()
with image_blocks as demo:
gr.Markdown(
"""
# Virtual Home Staging
""")
with gr.Group():
with gr.Box():
with gr.Row():
with gr.Column():
image = gr.Image(source='upload', elem_id="image_upload",
type="pil", label="Upload",
).style(height=400)
with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
style = gr.Dropdown(
["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",
info="only select if the image is NOT a backyard"
)
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")
additional_info = gr.Textbox(placeholder='Enter additional information',
label="Additional information", elem_id="input-add-info")
btn = gr.Button("Inpaint!").style(
margin=False,
rounded=(False, True, True, False),
full_width=False,
)
gr.Markdown("## Image Examples")
gr.Examples(
examples=[os.path.join(os.path.dirname(__file__),
"examples/exciting-small-kitchen-ideas-1821197-hero-d00f516e2fbb4dcabb076ee9685e877a.jpg"),
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"),
os.path.join(os.path.dirname(__file__),
"examples/backyard_additions.jpg"),
os.path.join(os.path.dirname(__file__),
"examples/living-room-gallery-shelves-l-shaped-couch-ELeyNpyyqpZ8hosOG3EG1X-b5a39646574544e8a75f2961332cd89a.jpg"),
os.path.join(os.path.dirname(__file__),
"examples/modern-dining-room-ideas-4147451-hero-d6333998f8b34620adfd4d99ac732586.jpg"),
],
inputs=image
)
with gr.Column():
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])
image_blocks.queue(concurrency_count=3)
image_blocks.launch()