Update app.py

app.py

@@ -27,6 +27,18 @@ def containsNumber(value):
         if character.isdigit():
             return True
     return False
+    
+def creativity(intensity):
+    if(intensity == 'Low'):
+        top_p = 0.95
+        top_k = 10
+    elif(intensity == 'Medium'):
+        top_p = 0.9
+        top_k = 50
+    if(intensity == 'High'):
+        top_p = 0.85
+        top_k = 100
+    return top_p, top_k
 
 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}
@@ -38,7 +50,7 @@ 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")
+    image = Image.new("RGBA", im_size, color="white")
     draw = aggdraw.Draw(image)
 
     for poly, color, in zip(polygons, colors):
@@ -66,7 +78,7 @@ def draw_polygons(polygons, colors, im_size=(256, 256), b_color="white", fpath=N
             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)
+    image = Image.frombytes("RGBA", (256,256), draw.tobytes()).transpose(Image.FLIP_TOP_BOTTOM)
 
     if(fpath):
         image.save(fpath, quality=100, subsampling=0)
@@ -80,14 +92,14 @@ def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
         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)
-    #print(model_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)
+    input_ids = tokenizer(model_prompt, return_tensors='pt')
+    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)
-    #print(output)
     
-    layout = output[0].lstrip().split('[User prompt]')[1].split('[Layout] ')[1].split(', ')
+    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]
@@ -99,64 +111,16 @@ def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
         
     geom = []
     for poly in polygons:
-        geom.append(Polygon(np.array(poly, dtype=int)))
+        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)))
         
-    colors = [architext_colors[housegan_labels[space]] for space in spaces]
+    colors = [architext_colors2[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').to(device)
-    mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, eos_token_id=50256, 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="""
@@ -230,33 +194,17 @@ custom_css="""
 .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')
+}""" 
+creative_slider = gr.inputs.Radio(["Low", "Medium", "High"], default="Medium", label='Creativity')
 textbox = gr.inputs.Textbox(placeholder='a 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],
+iface2 = gr.Interface(fn=prompt_to_layout, inputs=[textbox, creative_slider], 
+                      outputs=[generated],
                       css=custom_css,
+                      allow_flagging=False,
+                      allow_screenshot=False,
                       thumbnail="thumbnail_gradio.PNG",
                       description='Demo of Semantic Generation of Residential Layouts \n',
                       article='''<div>
@@ -271,38 +219,4 @@ iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox, topp_slider,
     <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>, using a finetuned <a href='https://huggingface.co/EleutherAI/gpt-neo-125M'> GPT-Neo</a> model. </p>
 </div>''')
 
-iface.launch()
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+iface.launch()