add more vis

hoho/color_mappings.py

@@ -1,3 +1,5 @@
+import numpy as np 
+
 gestalt_color_mapping = {
     "unclassified": (215, 62, 138),
     "apex": (235, 88, 48),
@@ -26,6 +28,7 @@ gestalt_color_mapping = {
     "other_wall": (112, 61, 240),
     "trim": (151, 206, 58),
     "unknown": (127, 127, 127),
+    "transition_line": (0,0,0),
 }
 
 ade20k_color_mapping = {
@@ -180,3 +183,24 @@ ade20k_color_mapping = {
     'clock': (102, 255, 0),
     'flag': (92, 0, 255),
 }
+
+
+# edge_colors = np.asarray([(214, 251, 248),
+#                           (13, 94, 47),
+#                           (54, 243, 63),
+#                           (187, 123, 236),
+#                           (162, 162, 32),
+#                           (169, 255, 219),
+#                           (8, 89, 52),
+#                           (85, 27, 65),
+#                           (0, 0, 0)]
+
+
+# edge_colors = np.array([[ 54, 243,  63], 
+#                         [214, 251, 248], 
+#                         [169, 255, 219], 
+#                         [ 13,  94,  47], 
+#                         [162, 162,  32], 
+#                         [187, 123, 236], 
+#                         [ 85,  27,  65], 
+#                         [  0,  0,    0]])

hoho/hoho.py

@@ -176,6 +176,66 @@ def proc(row, split='train'):
     return Sample(out)
 
 
+
+
+
+from . import read_write_colmap
+def decode_colmap(s):
+    # with open('colmap_solve/points3D.bin', 'wb') as stream:
+    
+    with temp_working_directory():
+
+        with open('points3D.bin', 'wb') as stream:
+            stream.write(s['points3d'])
+
+
+        with open('cameras.bin', 'wb') as stream:
+            stream.write(s['cameras'])
+
+
+        with open('images.bin', 'wb') as stream:
+            stream.write(s['images'])
+    
+        
+        cameras, images, points3D = read_write_colmap.read_model(
+            path='.', ext='.bin'
+        )
+    return cameras, images, points3D
+
+from PIL import Image
+import io
+def decode(row):
+    cameras, images, points3D = decode_colmap(row)
+    
+    out = {}
+    
+    for k, v in row.items():
+        # colname = k.split('.')[0]
+        if k in {'ade20k', 'depthcm', 'gestalt'}:
+            # print(k, len(v), type(v))
+            v = [Image.open(io.BytesIO(im)) for im in v]
+            if k in out:
+                out[k].extend(v)
+            else:
+                out[k] = v
+        elif k in {'wireframe', 'mesh'}:
+            # out.update({a: b.tolist() for a,b in v.items()})
+            v = dict(np.load(io.BytesIO(v)))
+            out.update({a: b for a,b in v.items()})
+        elif k in 'kr':
+            out[k.upper()] = v
+        elif k == 'cameras':
+            out[k] = cameras
+        elif k == 'images':
+            out[k] = images
+        elif k =='points3d':
+            out[k] = points3D
+        else:
+            out[k] = v
+            
+    return Sample(out)
+
+
 class Sample(Dict):
     def __repr__(self):
         return str({k: v.shape if hasattr(v, 'shape') else [type(v[0])] if isinstance(v, list) else type(v) for k,v in self.items()})

hoho/vis.py

@@ -1,6 +1,9 @@
 import trimesh
 import numpy as np
 from copy import deepcopy
+from PIL import Image
+
+from . import color_mappings
 
 def line(p1, p2, c=(255,0,0), resolution=10, radius=0.05):
     '''draws a 3d cylinder along the line (p1, p2)'''
@@ -38,6 +41,19 @@ def line(p1, p2, c=(255,0,0), resolution=10, radius=0.05):
          
     return mesh
 
+def show_wf(row, radius=10):
+    EDGE_CLASSES = ['eave',
+                    'ridge',
+                    'step_flashing',
+                    'rake',
+                    'flashing',
+                    'post',
+                    'valley',
+                    'hip',
+                    'transition_line']
+    return [line(a,b, radius=radius, c=color_mappings.gestalt_color_mapping[EDGE_CLASSES[cls_id]]) for (a,b), cls_id in zip(np.stack([*row['wf_vertices']])[np.stack(row['wf_edges'])], row['edge_semantics'])]
+    # return [line(a,b, radius=radius, c=color_mappings.edge_colors[cls_id]) for (a,b), cls_id in zip(np.stack([*row['wf_vertices']])[np.stack(row['wf_edges'])], row['edge_semantics'])]
+
 
 def show_grid(edges, meshes=None, row_length=5):
     '''
@@ -85,3 +101,57 @@ def show_grid(edges, meshes=None, row_length=5):
     return trimesh.Scene(out)
 
 
+
+
+def visualize_order_images(row_order):
+    return create_image_grid(row_order['ade20k'] + row_order['gestalt'] + [visualize_depth(dm) for dm in row_order['depthcm']], num_per_row=len(row_order['ade20k']))
+
+def create_image_grid(images, target_length=312, num_per_row=2):
+    # Calculate the target size for the first image
+    first_img = images[0]
+    aspect_ratio = first_img.width / first_img.height
+    new_width = int((target_length ** 2 * aspect_ratio) ** 0.5)
+    new_height = int((target_length ** 2 / aspect_ratio) ** 0.5)
+    
+    # Resize the first image
+    resized_images = [img.resize((new_width, new_height), Image.Resampling.LANCZOS) for img in images]
+    
+    # Calculate the grid size
+    num_rows = (len(resized_images) + num_per_row - 1) // num_per_row
+    grid_width = new_width * num_per_row
+    grid_height = new_height * num_rows
+    
+    # Create a new image for the grid
+    grid_img = Image.new('RGB', (grid_width, grid_height))
+    
+    # Paste the images into the grid
+    for i, img in enumerate(resized_images):
+        x_offset = (i % num_per_row) * new_width
+        y_offset = (i // num_per_row) * new_height
+        grid_img.paste(img, (x_offset, y_offset))
+    
+    return grid_img
+
+
+import matplotlib
+def visualize_depth(depth, min_depth=None, max_depth=None, cmap='rainbow'):
+    depth = np.array(depth)
+    
+    if min_depth is None:
+        min_depth = np.min(depth)
+    if max_depth is None:
+        max_depth = np.max(depth)
+    
+    
+    # Normalize the depth to be between 0 and 1
+    depth = (depth - min_depth) / (max_depth - min_depth)
+    depth = np.clip(depth, 0, 1)
+    
+    # Use the matplotlib colormap to convert the depth to an RGB image
+    cmap = matplotlib.cm.get_cmap(cmap)
+    depth_image = (cmap(depth) * 255).astype(np.uint8)
+    
+    # Convert the depth image to a PIL image
+    depth_image = Image.fromarray(depth_image)
+    
+    return depth_image