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') | |
| 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} | |
| housegan_colors = [[0, 0, 0], [197, 203, 159], [169, 89, 142], [0, 132, 66], [190, 0, 198], [255, 255, 255], | |
| [6, 53, 17], [2, 54, 192], [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, 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=0.94, top_k=100, max_length=300) | |
| output = tokenizer.batch_decode(output, skip_special_tokens=True) | |
| output = output[0].split('[Layout]')[1].split(', ') | |
| spaces = [txt.split(':')[0] 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 = [housegan_colors[housegan_labels[space]] for space in spaces] | |
| _, im = draw_polygons(geom, colors, fpath=fpath) | |
| legend = Image.open(r"C:\\Users\\user\\Desktop\\legend3.png") | |
| im = np.array(im) | |
| im[:40, :] = np.array(legend) | |
| im = Image.fromarray(im) | |
| return im, output | |
| def mut_txt2layout(mut_output): | |
| output = mut_output[0].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 = [housegan_colors[housegan_labels[space]] for space in spaces] | |
| _, im = draw_polygons(geom, colors, fpath=None) | |
| legend = Image.open(r"C:\\Users\\user\\Desktop\\legend3.png") | |
| im = np.array(im) | |
| im[:40, :] = np.array(legend) | |
| im = Image.fromarray(im) | |
| return im | |
| def prompt_with_mutation(user_prompt, fpath=None): | |
| #Create initial layout based on prompt | |
| im, output = prompt_to_layout(user_prompt) | |
| #Create mutated layout based on initial | |
| cut_off = np.random.randint(1, 3, size=1)[0] | |
| cut_off = min(cut_off, len(output)-1) | |
| to_keep = ', '.join(output[:cut_off]) + ', ' | |
| new_prompt = '[User prompt] {} [Layout] {}'.format(user_prompt, to_keep) | |
| input_ids = tokenizer(new_prompt, return_tensors='pt') | |
| mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=0.94, top_k=100, max_length=300) | |
| mut_output = tokenizer.batch_decode(mut_output, skip_special_tokens=True) | |
| mut_im = mut_txt2layout(mut_output) | |
| return im, mut_im, output, mut_output | |
| # Gradio App | |
| def gen_and_mutate(user_prompt, mutate=False): | |
| if(mutate): | |
| im, mut_im = None, None | |
| while (mut_im is None): | |
| im, mut_im, output, mut_output = prompt_with_mutation(user_prompt) | |
| else: | |
| mut_im=Image.open(r"C:\\Users\\user\\Desktop\\empty.png") | |
| im, _ = prompt_to_layout(user_prompt) | |
| return im, mut_im | |
| checkbox = gr.inputs.Checkbox(label='Mutate') | |
| textbox = gr.inputs.Textbox(placeholder='Enter a prompt describing a layout, see below for instructions') | |
| generated = gr.outputs.Image(label='Generated Layout') | |
| mutated = gr.outputs.Image(label='Mutated Layout') | |
| iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox], outputs=[generated, mutated], | |
| thumbnail=r"E:\Datasets\MyFloorplans\text2text\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() | |