Handle the case where "bedroom1" is not an exact match but is still a bedroom.

app.py

@@ -45,7 +45,6 @@ def creativity(intensity):
 
 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}
-default_color = [255, 255, 255]
 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],]
 
@@ -118,7 +117,12 @@ def prompt_to_layout(user_prompt, intensity, fpath=None):
     for poly in polygons:
         scaled_poly = scale(Polygon(np.array(poly, dtype=int)), xfact=2, yfact=2, origin=(0,0))
         geom.append(scaled_poly)
-    colors = [architext_colors[housegan_labels[space]] if space in housegan_labels else default_color for space in spaces]
+    colors = []
+    for space in spaces:
+        for key in housegan_labels.keys():
+            if key in space:
+                colors.append(architext_colors[housegan_labels[key]])
+                break
     _, im = draw_polygons(geom, colors, fpath=fpath)
     html = '<img class="labels" src="images/labels.png" />'
     legend = Image.open("labels.png")