disable image dump

handcrafted_solution.py

@@ -12,6 +12,7 @@ from scipy.spatial.distance import cdist
 from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
 from hoho.color_mappings import gestalt_color_mapping, ade20k_color_mapping
 
+DUMP_IMG = False
 
 def empty_solution():
     '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
@@ -122,9 +123,10 @@ def get_uv_depth(vertices, depth):
     vertex_depth = depth[(uv_int[:, 1] , uv_int[:, 0])]
     return uv, vertex_depth
 
-def get_smooth_uv_depth(vertices, depth):
+def get_smooth_uv_depth(vertices, depth, gest_seg_np):
     '''Get the depth of the vertices from the depth image'''
-    #print(f'max depth = {np.max(depth)}')
+    #print(f'max depth = {np.max(depth)}')    
+
     uv = []
     for v in vertices:
         uv.append(v['xy'])
@@ -139,15 +141,17 @@ def get_smooth_uv_depth(vertices, depth):
       for i in range(max(0, x - r), min(W, x + r + 1)):
         for j in range(max(0, y - r), min(H, y + r + 1)):
             if np.sqrt((i - x)**2 + (j - y)**2) <= r:
-                local_depths.append(depth[j, i])      
+              local_depths.append(depth[j, i])
       return local_depths
     
     vertex_depth = []
     for x,y in zip(a,b):
       local_depths = get_local_depth(x,y, H, W, depth, 5)
+      #print(f'local_depths={local_depths}')
       local_mean = np.mean(local_depths)
       vertex_depth.append(local_mean)
     vertex_depth = np.array(vertex_depth)
+    
 
     return uv, vertex_depth
 
@@ -232,12 +236,13 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
     # imsave apex and eave_end
     import random
     rid = random.random()
-    filename_apex_ade = f'apex_map_on_ade_{rid}.jpg'
-    cv2.imwrite(filename_apex_ade, apex_map_on_ade)
-    filename_apex_gest = f'apex_map_on_gest_{rid}.jpg'
-    cv2.imwrite(filename_apex_gest, apex_map_on_gest)
-    filename_apex_map = f'apex_map_{rid}.jpg'
-    cv2.imwrite(filename_apex_map, apex_map)
+    if DUMP_IMG:
+      filename_apex_ade = f'apex_map_on_ade_{rid}.jpg'
+      cv2.imwrite(filename_apex_ade, apex_map_on_ade)
+      filename_apex_gest = f'apex_map_on_gest_{rid}.jpg'
+      cv2.imwrite(filename_apex_gest, apex_map_on_gest)
+      filename_apex_map = f'apex_map_{rid}.jpg'
+      cv2.imwrite(filename_apex_map, apex_map)
 
     print(f'{len(vertices)} vertices detected')
     # Connectivity
@@ -284,8 +289,9 @@ def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_
                     for a_i, a in enumerate(connected_verts):
                         for b in connected_verts[a_i+1:]:
                             connections.append((a, b))
-    filename_edges_map = f'edges_map_{rid}.jpg'
-    cv2.imwrite(filename_edges_map, line_img)
+    if DUMP_IMG: 
+      filename_edges_map = f'edges_map_{rid}.jpg'
+      cv2.imwrite(filename_edges_map, line_img)
     return vertices, connections
 
 def get_uv_dept_category(vertices, depth, ade_seg):
@@ -424,8 +430,8 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
             print (f'Not enough vertices or connections in image {i}')
             vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
             continue
-        #uv, depth_vert = get_uv_depth(vertices, depth_np) 
-        uv, depth_vert = get_smooth_uv_depth(vertices, depth_np)
+        #uv, depth_vert = get_uv_depth(vertices, depth_np)         
+        uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np)
         
         # Normalize the uv to the camera intrinsics
         xy_local = np.ones((len(uv), 3))