contiguous mutation

app.py

@@ -68,10 +68,11 @@ def prompt_to_layout(user_prompt, fpath=None):
     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]
+    
+    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 output]
+    coordinates = [txt.split(':')[1] for txt in layout]
     coordinates = [re.findall(regex, coord) for coord in coordinates]
     
     polygons = []
@@ -82,20 +83,20 @@ def prompt_to_layout(user_prompt, fpath=None):
     for poly in polygons:
         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")
+    legend = Image.open("legend3.png")
     
     im = np.array(im)
     im[:40, :] = np.array(legend)
     im = Image.fromarray(im)
     
-    return im, output
+    return im, layout, output
     
 def mut_txt2layout(mut_output):
-    output = mut_output[0].split('[Layout]')[1].split(', ')
+    output = mut_output[0].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]
@@ -108,10 +109,10 @@ def mut_txt2layout(mut_output):
     for poly in polygons:
         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=None)
 
-    legend = Image.open("legend.png")
+    legend = Image.open("legend3.png")
 
     im = np.array(im)
     im[:40, :] = np.array(legend)
@@ -119,22 +120,23 @@ def mut_txt2layout(mut_output):
     
     return im
     
-def prompt_with_mutation(user_prompt, fpath=None):
+def prompt_with_mutation(user_prompt, mut_rate, fpath=None):
     
     #Create initial layout based on prompt
-    im, output = prompt_to_layout(user_prompt)
+    im, layout, 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)
+    mut_len = int((1-mut_rate)*len(layout))
+    index1 = random.randrange(0,len(layout)-mut_len)
+    rooms = layout[index1:index1+mut_len]
+    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=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
+    return im, mut_im
     
 # Gradio App
 
@@ -211,24 +213,28 @@ custom_css="""
   opacity: 1 !important;
 }"""
 
-def gen_and_mutate(user_prompt, mutate=False):    
+def gen_and_mutate(user_prompt, mutate=False, mut_rate=0.2):    
     if(mutate):
         im, mut_im = None, None
         while (mut_im is None):
-            im, mut_im, output, mut_output = prompt_with_mutation(user_prompt)
+            try:
+                im, mut_im = prompt_with_mutation(user_prompt, mut_rate)
+            except:
+                pass
     else:
         mut_im=Image.open("empty.png")
-        im, _ = prompt_to_layout(user_prompt)
-
+        im, _, _ = prompt_to_layout(user_prompt)
+        
     return im, mut_im
     
 checkbox =  gr.inputs.Checkbox(label='Mutate')
-textbox = gr.inputs.Textbox(placeholder='house with two bedrooms and one bathroom', lines="4", label="DESCRIBE YOUR DESIGN")
+slider = gr.inputs.Slider(0.2, 0.8, 0.1, label='Mutation rate')
+textbox = gr.inputs.Textbox(placeholder='house with two bedrooms and one bathroom', lines="3", 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], outputs=[generated, mutated],
+iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox, slider], outputs=[generated, mutated],
                       css=custom_css,
                       thumbnail="thumbnail_gradio.PNG",
                       description='Demo of Semantic Generation of Residential Layouts \n',