Spaces:
Runtime error
Runtime error
| from pathlib import Path | |
| import numpy as np | |
| import random | |
| import re | |
| import textwrap | |
| from shapely.geometry.polygon import Polygon | |
| import aggdraw | |
| from PIL import Image, ImageDraw, ImageOps, ImageFilter, ImageFont, ImageColor | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM | |
| #finetuned = AutoModelForCausalLM.from_pretrained('model') | |
| finetuned = AutoModelForCausalLM.from_pretrained('.') | |
| tokenizer = AutoTokenizer.from_pretrained('gpt2') | |
| # Utility functions | |
| housegan_labels = {"living_room": 1, "kitchen": 2, "bedroom": 3, "bathroom": 4, "missing": 5, "closet": 6, | |
| "balcony": 7, "corridor": 8, "dining_room": 9, "laundry_room": 10} | |
| architext_colors = [[0, 0, 0], [249, 222, 182], [195, 209, 217], [250, 120, 128], [126, 202, 234], [190, 0, 198], [255, 255, 255], | |
| [6, 53, 17], [17, 33, 58], [132, 151, 246], [197, 203, 159], [6, 53, 17],] | |
| regex = re.compile(".*?\((.*?)\)") | |
| def draw_polygons(polygons, colors, im_size=(256, 256), b_color="white", fpath=None): | |
| image = Image.new("RGB", im_size, color="white") | |
| draw = aggdraw.Draw(image) | |
| for poly, color, in zip(polygons, colors): | |
| #get initial polygon coordinates | |
| xy = poly.exterior.xy | |
| coords = np.dstack((xy[1], xy[0])).flatten() | |
| # draw it on canvas, with the appropriate colors | |
| brush = aggdraw.Brush((0, 0, 0), opacity=255) | |
| draw.polygon(coords, brush) | |
| #get inner polygon coordinates | |
| small_poly = poly.buffer(-1, resolution=32, cap_style=2, join_style=2, mitre_limit=5.0) | |
| if small_poly.geom_type == 'MultiPolygon': | |
| mycoordslist = [list(x.exterior.coords) for x in small_poly] | |
| for coord in mycoordslist: | |
| coords = np.dstack((np.array(coord)[:,1], np.array(coord)[:, 0])).flatten() | |
| brush2 = aggdraw.Brush((0, 0, 0), opacity=255) | |
| draw.polygon(coords, brush2) | |
| elif poly.geom_type == 'Polygon': | |
| #get inner polygon coordinates | |
| xy2 = small_poly.exterior.xy | |
| coords2 = np.dstack((xy2[1], xy2[0])).flatten() | |
| # draw it on canvas, with the appropriate colors | |
| brush2 = aggdraw.Brush((color[0], color[1], color[2]), opacity=255) | |
| draw.polygon(coords2, brush2) | |
| image = Image.frombytes("RGB", (256,256), draw.tobytes()).transpose(Image.FLIP_TOP_BOTTOM) | |
| if(fpath): | |
| image.save(fpath, quality=100, subsampling=0) | |
| return draw, image | |
| def prompt_to_layout(user_prompt, top_p, top_k, fpath=None): | |
| model_prompt = '[User prompt] {} [Layout]'.format(user_prompt) | |
| input_ids = tokenizer(model_prompt, return_tensors='pt') | |
| output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, max_length=300) | |
| output = tokenizer.batch_decode(output, skip_special_tokens=True) | |
| layout = output[0].split('[User prompt]')[1].split('[Layout] ')[1].split(', ') | |
| spaces = [txt.split(':')[0] for txt in layout] | |
| coordinates = [txt.split(':')[1] for txt in layout] | |
| coordinates = [re.findall(regex, coord) for coord in coordinates] | |
| polygons = [] | |
| for coord in coordinates: | |
| polygons.append([point.split(',') for point in coord]) | |
| geom = [] | |
| for poly in polygons: | |
| geom.append(Polygon(np.array(poly, dtype=int))) | |
| colors = [architext_colors[housegan_labels[space]] for space in spaces] | |
| _, im = draw_polygons(geom, colors, fpath=fpath) | |
| legend = Image.open("legend.png") | |
| im = np.array(im) | |
| im[:40, :] = np.array(legend) | |
| im = Image.fromarray(im) | |
| return im, layout, output | |
| def mut_txt2layout(mut_output): | |
| output = mut_output[0].rstrip().split('[User prompt]')[1].split('[Layout]')[1].split(', ') | |
| spaces = [txt.split(':')[0].strip(' ') for txt in output] | |
| coordinates = [txt.split(':')[1] for txt in output] | |
| coordinates = [re.findall(regex, coord) for coord in coordinates] | |
| polygons = [] | |
| for coord in coordinates: | |
| polygons.append([point.split(',') for point in coord]) | |
| geom = [] | |
| for poly in polygons: | |
| geom.append(Polygon(np.array(poly, dtype=int))) | |
| colors = [architext_colors[housegan_labels[space]] for space in spaces] | |
| _, im = draw_polygons(geom, colors, fpath=None) | |
| legend = Image.open("legend.png") | |
| im = np.array(im) | |
| im[:40, :] = np.array(legend) | |
| im = Image.fromarray(im) | |
| return im | |
| def prompt_with_mutation(user_prompt, top_p, top_k, mut_rate, fpath=None): | |
| #Create initial layout based on prompt | |
| im, layout, output = prompt_to_layout(user_prompt, top_p=top_p, top_k=top_k) | |
| #Create mutated layout based on initial | |
| mut_len = int((1-mut_rate)*len(layout)) | |
| index1 = random.randrange(0,len(layout)-mut_len) | |
| rooms = layout[index1:index1+mut_len] | |
| rooms[-1] = rooms[-1].split(':')[0] + ':' | |
| rooms = ', '.join(rooms)# + ', ' | |
| new_prompt = '[User prompt] {} [Layout] {}'.format(user_prompt, rooms) | |
| input_ids = tokenizer(new_prompt, return_tensors='pt') | |
| mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, max_length=400) | |
| mut_output = tokenizer.batch_decode(mut_output, skip_special_tokens=True) | |
| mut_im = mut_txt2layout(mut_output) | |
| return im, mut_im | |
| # Gradio App | |
| custom_css=""" | |
| @import url("https://use.typekit.net/nid3pfr.css"); | |
| .gradio_page { | |
| display: flex; | |
| width: 100vw; | |
| min-height: 50vh; | |
| flex-direction: column; | |
| justify-content: center; | |
| align-items: center; | |
| margin: 0px; | |
| max-width: 100vw; | |
| background: #FFFFFF; | |
| } | |
| .gradio_interface { | |
| width: 100vw; | |
| max-width: 1500px; | |
| } | |
| .gradio_interface[theme=default] .panel_buttons { | |
| justify-content: flex-end; | |
| } | |
| .gradio_interface[theme=default] .panel_button { | |
| flex: 0 0 0; | |
| min-width: 150px; | |
| } | |
| .gradio_interface[theme=default] .panel_button.submit { | |
| background: #11213A; | |
| border-radius: 5px; | |
| color: #FFFFFF; | |
| text-transform: uppercase; | |
| min-width: 150px; | |
| height: 4em; | |
| letter-spacing: 0.15em; | |
| flex: 0 0 0; | |
| } | |
| .gradio_interface[theme=default] .panel_button.submit:hover { | |
| background: #000000; | |
| } | |
| .input_text { | |
| font: 200 50px garamond-premier-pro-display, serif; | |
| line-height: 115%; | |
| color: #11213A; | |
| border-radius: 0px; | |
| border: 3px solid #11213A; | |
| } | |
| .input_text:focus { | |
| border-color: #FA7880; | |
| } | |
| .gradio_interface[theme=default] .input_text input, | |
| .gradio_interface[theme=default] .input_text textarea { | |
| padding: 30px; | |
| } | |
| .input_text textarea:focus-visible { | |
| outline: none; | |
| } | |
| .panel:nth-child(1) { | |
| margin-left: 50px; | |
| margin-right: 50px; | |
| margin-top: 80px; | |
| margin-bottom: 80px; | |
| max-width: 750px; | |
| } | |
| .panel:nth-child(2) { | |
| background: #D3ECF5; | |
| } | |
| .gradio_interface[theme=default] .output_image .image_preview_holder { | |
| background: #D3ECF5; | |
| } | |
| .gradio_interface[theme=default] .component_set { | |
| background: transparent; | |
| opacity: 1 !important; | |
| }""" | |
| def gen_and_mutate(user_prompt, mutate=False, top_p=0.94, top_k=100, mut_rate=0.2): | |
| if(mutate): | |
| im, mut_im = None, None | |
| while (mut_im is None): | |
| try: | |
| im, mut_im = prompt_with_mutation(user_prompt, top_p, top_k, mut_rate) | |
| except: | |
| pass | |
| else: | |
| mut_im=Image.open("empty.png") | |
| im, _, _ = prompt_to_layout(user_prompt, top_p, top_k) | |
| return im, mut_im | |
| checkbox = gr.inputs.Checkbox(label='Mutate') | |
| topp_slider = gr.inputs.Slider(0.1, 1.0, 0.01, default=0.94, label='top_p') | |
| topk_slider = gr.inputs.Slider(0, 100, 25, default=0, label='top_k') | |
| mut_slider = gr.inputs.Slider(0.2, 0.8, 0.1, default=0.3, label='Mutation rate') | |
| textbox = gr.inputs.Textbox(placeholder='house with two bedrooms and one bathroom', lines="2", | |
| label="DESCRIBE YOUR DESIGN") | |
| generated = gr.outputs.Image(label='Generated Layout') | |
| mutated = gr.outputs.Image(label='Mutated Layout') | |
| iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox, topp_slider, topk_slider, mut_slider], outputs=[generated, mutated], | |
| css=custom_css, | |
| thumbnail="thumbnail_gradio.PNG", | |
| description='Demo of Semantic Generation of Residential Layouts \n', | |
| article='''<div> | |
| <p> This app allows users the use of natural language prompts for appartment layout generation, using a variety of semantic information:</p> | |
| <ul> | |
| <li> <strong>typology</strong>: "a house with two bedrooms and two bathrooms"</li> | |
| <li> <strong>enumeration</strong>: "a house with five rooms"</li> | |
| <li> <strong>adjacency</strong>: "the kitchen is adjacent to a bedroom", "the living room is not adjacent to the bathroom"</li> | |
| <li> <strong>location</strong>: "a house with a bedroom in the north east side"</li> | |
| </ul> | |
| <p>You can also create a mutation of the generated layout by enabling the 'Mutate' option.</p> | |
| <p> Made by: <a href='https://www.linkedin.com/in/theodorosgalanos/'>Theodoros </a> <a href='https://twitter.com/TheodoreGalanos'> Galanos</a> and <a href='https://twitter.com/tylerlastovich'>Tyler Lastovich</a> </p> | |
| </div>''') | |
| iface.launch() | |