Adds color visualization script

Adds a script to generate a color legend image, mapping semantic labels to specific colors.

Also, converts some numpy arrays to lists to prevent issues with serialization.

color_visu.py

@@ -0,0 +1,144 @@
+import cv2
+import numpy as np
+
+# Color mappings (as provided in the context file)
+gestalt_color_mapping = {
+    "unclassified": (215, 62, 138),
+    "apex": (235, 88, 48),
+    "eave_end_point": (248, 130, 228),
+    "flashing_end_point": (71, 11, 161),
+    "ridge": (214, 251, 248),
+    "rake": (13, 94, 47),
+    "eave": (54, 243, 63),
+    "post": (187, 123, 236),
+    "ground_line": (136, 206, 14),
+    "flashing": (162, 162, 32),
+    "step_flashing": (169, 255, 219),
+    "hip": (8, 89, 52),
+    "valley": (85, 27, 65),
+    "roof": (215, 232, 179),
+    "door": (110, 52, 23),
+    "garage": (50, 233, 171),
+    "window": (230, 249, 40),
+    "shutter": (122, 4, 233),
+    "fascia": (95, 230, 240),
+    "soffit": (2, 102, 197),
+    "horizontal_siding": (131, 88, 59),
+    "vertical_siding": (110, 187, 198),
+    "brick": (171, 252, 7),
+    "concrete": (32, 47, 246),
+    "other_wall": (112, 61, 240),
+    "trim": (151, 206, 58),
+    "unknown": (127, 127, 127),
+    "transition_line": (0,0,0),
+}
+
+edge_color_mapping = {
+    'cornice_return': (215, 62, 138),
+    'cornice_strip': (235, 88, 48),
+    'eave':  (54, 243, 63),
+    "flashing": (162, 162, 32),
+    'hip': (8, 89, 52),
+    'rake': (13, 94, 47),
+    'ridge': (214, 251, 248),
+    "step_flashing": (169, 255, 219),
+    'transition_line': (200,0,50),
+    'valley': (85, 27, 65),
+}
+
+# Parameters for the legend image
+swatch_width = 150
+swatch_height = 25
+text_color_bgr = (0, 0, 0)  # Black
+font = cv2.FONT_HERSHEY_SIMPLEX
+font_scale = 0.4
+font_thickness = 1
+horizontal_padding = 20
+vertical_padding = 15 # Padding at top/bottom of sections and between items
+text_offset_x = 10 # Space between swatch and text
+item_row_height = swatch_height + 5 # Total height for one swatch + text row, including small spacing
+section_title_font_scale = 0.6
+section_title_font_thickness = 1
+section_title_area_height = 40 # Space allocated for section title
+
+all_mappings_with_titles = [
+    ("Gestalt Colors", gestalt_color_mapping),
+    ("Edge Colors", edge_color_mapping),
+]
+
+# Calculate maximum width needed for labels to determine overall image width
+max_label_width = 0
+for _, color_map in all_mappings_with_titles:
+    for label in color_map.keys():
+        (text_w, _), _ = cv2.getTextSize(label, font, font_scale, font_thickness)
+        if text_w > max_label_width:
+            max_label_width = text_w
+
+image_width = horizontal_padding + swatch_width + text_offset_x + max_label_width + horizontal_padding
+
+# Calculate total image height
+image_height = vertical_padding # Initial top padding
+for i, (_, color_map) in enumerate(all_mappings_with_titles):
+    image_height += section_title_area_height
+    image_height += len(color_map) * item_row_height
+    if i < len(all_mappings_with_titles) - 1:
+        image_height += vertical_padding # Padding between sections
+image_height += vertical_padding # Final bottom padding
+
+# Create a white canvas
+legend_image = np.full((image_height, image_width, 3), 255, dtype=np.uint8) # White background
+
+current_y = vertical_padding
+
+for title_text, color_map in all_mappings_with_titles:
+    # Draw section title
+    (title_w, title_h), title_baseline = cv2.getTextSize(title_text, font, section_title_font_scale, section_title_font_thickness)
+    title_x = horizontal_padding # Align title to the left
+    # Center title vertically within its allocated space
+    title_y_baseline = current_y + (section_title_area_height - title_h) // 2 + title_h - title_baseline // 2
+    cv2.putText(legend_image, title_text, (title_x, title_y_baseline), font, section_title_font_scale, text_color_bgr, section_title_font_thickness)
+    current_y += section_title_area_height
+
+    # Draw color swatches and labels for the current map
+    for label, rgb_color in color_map.items():
+        # OpenCV uses BGR, so convert RGB to BGR
+        bgr_color = (rgb_color[2], rgb_color[1], rgb_color[0])
+
+        swatch_x1 = horizontal_padding
+        swatch_y1 = current_y
+        swatch_x2 = swatch_x1 + swatch_width
+        swatch_y2 = swatch_y1 + swatch_height
+
+        # Draw the color swatch (filled rectangle)
+        cv2.rectangle(legend_image, (swatch_x1, swatch_y1), (swatch_x2, swatch_y2), bgr_color, -1)
+        # Draw a thin black border around the swatch for clarity
+        cv2.rectangle(legend_image, (swatch_x1, swatch_y1), (swatch_x2, swatch_y2), (0,0,0), 1)
+
+
+        # Get text size for precise vertical centering
+        (text_w, text_h), baseline = cv2.getTextSize(label, font, font_scale, font_thickness)
+        
+        # Calculate position for the text (vertically centered with the swatch)
+        text_x = swatch_x2 + text_offset_x
+        # The 'org' for putText is the bottom-left corner of the text string.
+        # To center text_h within swatch_height:
+        text_y_baseline = swatch_y1 + (swatch_height - text_h) // 2 + text_h
+        
+        cv2.putText(legend_image, label, (text_x, text_y_baseline), font, font_scale, text_color_bgr, font_thickness)
+        
+        current_y += item_row_height # Move to the next item row
+    
+    current_y += vertical_padding # Add padding after the items in this section
+
+# Save the image
+output_filename = "gestalt_color_legend.png"
+try:
+    cv2.imwrite(output_filename, legend_image)
+    print(f"Legend image saved as {output_filename}")
+except Exception as e:
+    print(f"Error saving image: {e}")
+
+# To display the image (optional, uncomment if you have a display environment)
+# cv2.imshow("Gestalt Color Legend", legend_image)
+# cv2.waitKey(0)
+# cv2.destroyAllWindows()

predict.py

@@ -9,7 +9,7 @@ def convert_entry_to_human_readable(entry):
         if 'colmap' in k:
             out[k] = read_colmap_rec(v)
         elif k in ['wf_vertices', 'wf_edges', 'K', 'R', 't', 'depth']:
-            out[k] = np.array(v)
+            out[k] = v
         else:
             out[k]=v
     out['__key__'] = entry['order_id']

train.py

@@ -25,6 +25,8 @@ idx = 0
 for a in ds['train']:
     colmap = read_colmap_rec(a['colmap_binary'])
 
+    #plot_all_modalities(a)
+
     try:
         pred_vertices, pred_edges = predict_wireframe(a)
     except: