Update app.py

app.py

@@ -4,11 +4,9 @@ import numpy as np
 import os
 import random
 import re
-import textwrap
 import torch
 from shapely.geometry.polygon import Polygon
 from shapely.affinity import scale
-import aggdraw
 from PIL import Image, ImageDraw, ImageOps, ImageFilter, ImageFont, ImageColor
 
 import gradio as gr
@@ -51,87 +49,62 @@ architext_colors = [[0, 0, 0], [249, 222, 182], [195, 209, 217], [250, 120, 128]
 regex = re.compile(".*?\((.*?)\)")
 
 def draw_polygons(polygons, colors, im_size=(512, 512), b_color="white", fpath=None):
-
     image = Image.new("RGBA", im_size, color="white")
-    #draw = aggdraw.Draw(image)
     draw = ImageDraw.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)
         draw.polygon(list(coords), fill=(0, 0, 0))
-        
         #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)
                 draw.polygon(list(coords), fill=tuple(color)) 
         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)
             draw.polygon(list(coords2), fill=tuple(color))
-
-    #image = Image.frombytes("RGBA", im_size, draw.tobytes()).transpose(Image.FLIP_TOP_BOTTOM)
     image = image.transpose(Image.FLIP_TOP_BOTTOM)
-    
     if(fpath):
         image.save(fpath, quality=100, subsampling=0)
-
     return draw, image
 
 def prompt_to_layout(user_prompt, intensity, fpath=None):
-
     if(containsNumber(user_prompt) == True):
         spaced_prompt = user_prompt.split(' ')
         new_prompt = ' '.join([word if word.isdigit() == False else num2words(int(word)).lower() for word in spaced_prompt])
         model_prompt = '[User prompt] {} [Layout]'.format(new_prompt)
-        
     top_p, top_k = creativity(intensity)
     model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
     input_ids = tokenizer(model_prompt, return_tensors='pt').to(device)
     output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, 
                                 eos_token_id=50256, max_length=400)
     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]
-    
+    coords = [txt.split(':')[1] for txt in layout]
+    coordinates = [re.findall(regex, coord) for coord in coords]
+    layout_dict = dict(zip(spaces, coords))
     polygons = []
     for coord in coordinates:
-        polygons.append([point.split(',') for point in coord])
-        
+        polygons.append([point.split(',') for point in coord])  
     geom = []
     for poly in polygons:
         scaled_poly = scale(Polygon(np.array(poly, dtype=int)), xfact=2, yfact=2, origin=(0,0))
-        geom.append(scaled_poly)
-        #geom.append(Polygon(np.array(poly, dtype=int)))
-        
+        geom.append(scaled_poly)  
     colors = [architext_colors[housegan_labels[space]] for space in spaces]
-    
     _, im = draw_polygons(geom, colors, fpath=fpath)
-    
     html = '<img class="labels" src="images/labels.png" />'
-    
     legend = Image.open("labels.png")
     imgs_comb = np.vstack([im, legend])
     imgs_comb = Image.fromarray(imgs_comb)
-    
-    return imgs_comb, layout 
+    return imgs_comb, layout_dict
      
     
 # Gradio App
@@ -257,13 +230,14 @@ creative_slider = gr.inputs.Radio(["Low", "Medium", "High"], default="Low", labe
 textbox = gr.inputs.Textbox(placeholder='An apartment with two bedrooms and one bathroom', lines="3", 
                             label="DESCRIBE YOUR IDEAL APARTMENT")
 generated = gr.outputs.Image(label='Generated Layout')
+layout = gr.outputs.Text(label='Layout Coordinates')
 
 iface = gr.Interface(fn=prompt_to_layout, inputs=[textbox, creative_slider], 
                      outputs=[generated, layout],
-                     css=custom_css,
+                     #css=custom_css,
                      theme="default",
                      allow_flagging='never',
                      allow_screenshot=False,
-                     thumbnail="thumbnail_gradio.PNG")
+                     #thumbnail="thumbnail_gradio.PNG")
 
 iface.launch(enable_queue=True)