slightly improved performance

handcrafted_solution.py

@@ -28,7 +28,7 @@ def empty_solution():
 
 def undesired_objects(image):
     image = image.astype('uint8')
-    nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(image, connectivity=8)
+    nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(image, connectivity=4)
     sizes = stats[:, -1]
     max_label = 1
     max_size = sizes[1]
@@ -73,8 +73,8 @@ def get_vertices(image_gestalt, *, color_range=4., dialations=3, erosions=1, ker
     eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask, cv2.MORPH_DILATE, kernel, iterations=dialations)
     eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask, cv2.MORPH_ERODE, kernel, iterations=erosions)
 
-    *_, apex_centroids = cv2.connectedComponentsWithStats(apex_mask, connectivity=8, stats=cv2.CV_32S)
-    *_, other_centroids = cv2.connectedComponentsWithStats(eave_end_point_mask, connectivity=8, stats=cv2.CV_32S)
+    *_, apex_centroids = cv2.connectedComponentsWithStats(apex_mask, connectivity=4, stats=cv2.CV_32S)
+    *_, other_centroids = cv2.connectedComponentsWithStats(eave_end_point_mask, connectivity=4, stats=cv2.CV_32S)
 
     return apex_centroids[1:], other_centroids[1:], apex_mask, eave_end_point_mask
 
@@ -96,7 +96,7 @@ def infer_vertices(image_gestalt, *, color_range=4.):
     intersection_mask = cv2.bitwise_and(ridge_mask, rake_mask)
     intersection_mask = cv2.morphologyEx(intersection_mask, cv2.MORPH_DILATE, np.ones((11, 11)), iterations=3)
 
-    *_, inferred_centroids = cv2.connectedComponentsWithStats(intersection_mask, connectivity=8, stats=cv2.CV_32S)
+    *_, inferred_centroids = cv2.connectedComponentsWithStats(intersection_mask, connectivity=4, stats=cv2.CV_32S)
 
     return inferred_centroids[1:], intersection_mask
 
@@ -144,15 +144,17 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., poi
     # missed_vertices = get_missed_vertices(vertices, inferred_vertices, **kwargs)
     # vertices = np.concatenate([vertices, missed_vertices])
 
-    scale = 1
-    vertex_size = np.zeros(vertices.shape[0])
-    for i, coords in enumerate(vertices):
-        # coords = np.round(coords).astype(np.uint32)
-        radius = point_radius  # np.clip(int(max_depth//2 + depth_np[coords[1], coords[0]]), 10, 30)#int(np.clip(max_depth - depth_np[coords[1], coords[0]], 10, 20))
-        vertex_size[i] = (scale * radius) ** 2  # because we are using squared distances
+    vertices = KDTree(vertices)
+
+    # scale = 1
+    # vertex_size = np.zeros(vertices.shape[0])
+    # for i, coords in enumerate(vertices):
+    #     # coords = np.round(coords).astype(np.uint32)
+    #     radius = point_radius  # np.clip(int(max_depth//2 + depth_np[coords[1], coords[0]]), 10, 30)#int(np.clip(max_depth - depth_np[coords[1], coords[0]], 10, 20))
+    #     vertex_size[i] = (scale * radius) ** 2  # because we are using squared distances
 
     for edge_class in ['eave', 'ridge', 'rake', 'valley', 'flashing', 'step_flashing']:
-        if len(vertices) < 2:
+        if len(vertices.data) < 2:
             break
         edge_color = np.array(gestalt_color_mapping[edge_class])
 
@@ -162,117 +164,100 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., poi
         mask = cv2.morphologyEx(mask,
                                 cv2.MORPH_DILATE, np.ones((3, 3)), iterations=1)
 
-        if np.any(mask):
+        if not np.any(mask):
+            continue
 
-            rho = 1  # distance resolution in pixels of the Hough grid
-            theta = np.pi / 180  # angular resolution in radians of the Hough grid
-            threshold = 20  # minimum number of votes (intersections in Hough grid cell)
-            min_line_length = 60  # minimum number of pixels making up a line
-            max_line_gap = 40  # maximum gap in pixels between connectable line segments
+        rho = 1  # distance resolution in pixels of the Hough grid
+        theta = np.pi / 180  # angular resolution in radians of the Hough grid
+        threshold = 20  # minimum number of votes (intersections in Hough grid cell)
+        min_line_length = 60  # minimum number of pixels making up a line
+        max_line_gap = 40  # maximum gap in pixels between connectable line segments
 
-            # Run Hough on edge detected image
-            # Output "lines" is an array containing endpoints of detected line segments
-            cv2.GaussianBlur(mask, (11, 11), 0, mask)
-            lines = cv2.HoughLinesP(mask, rho, theta, threshold, np.array([]),
-                                    min_line_length, max_line_gap)
+        # Run Hough on edge detected image
+        # Output "lines" is an array containing endpoints of detected line segments
+        cv2.GaussianBlur(mask, (11, 11), 0, mask)
+        lines = cv2.HoughLinesP(mask, rho, theta, threshold, np.array([]),
+                                min_line_length, max_line_gap)
 
-            edges = []
+        edges = []
 
-            if lines is None:
-                continue
+        if lines is None:
+            continue
 
-            line_directions = np.zeros((len(lines), 2))
-            for line_idx, line in enumerate(lines):
-                for x1, y1, x2, y2 in line:
-                    if x1 < x2:
-                        x1, y1, x2, y2 = x2, y2, x1, y1
-                    direction = (np.array([x2 - x1, y2 - y1]))
-                    direction = direction / np.linalg.norm(direction)
-                    line_directions[line_idx] = direction
+        line_directions = np.zeros((len(lines), 2))
+        for line_idx, line in enumerate(lines):
+            for x1, y1, x2, y2 in line:
+                if x1 < x2:
+                    x1, y1, x2, y2 = x2, y2, x1, y1
+                direction = (np.array([x2 - x1, y2 - y1]))
+                direction = direction / np.linalg.norm(direction)
+                line_directions[line_idx] = direction
 
-                    direction = extend * direction
+                direction = extend * direction
 
-                    x1, y1 = (-direction + (x1, y1)).astype(np.int32)
-                    x2, y2 = (+ direction + (x2, y2)).astype(np.int32)
+                x1, y1 = (-direction + (x1, y1)).astype(np.int32)
+                x2, y2 = (+ direction + (x2, y2)).astype(np.int32)
 
-                    edges.append((x1, y1, x2, y2))
+                edges.append((x1, y1, x2, y2))
 
-            edges = np.array(edges)
-            if len(edges) < 1:
-                continue
-            # calculate the distances between the vertices and the edge ends
-            begin_distances = cdist(vertices, edges[:, :2], metric="sqeuclidean")
-            end_distances = cdist(vertices, edges[:, 2:], metric="sqeuclidean")
-
-            begin_in_range_mask = begin_distances < vertex_size[:, np.newaxis]
-            end_in_range_mask = end_distances < vertex_size[:, np.newaxis]
-
-            in_range_connected_mask = np.logical_and(np.any(begin_in_range_mask, axis=0),
-                                                     np.any(end_in_range_mask, axis=0))
-
-            # where both ends are in range
-            begin_in_range_mask = np.logical_and(begin_in_range_mask, in_range_connected_mask)
-            end_in_range_mask = np.logical_and(end_in_range_mask, in_range_connected_mask)
-
-            begin_candidates = np.array(np.where(begin_in_range_mask))
-            end_candidates = np.array(np.where(end_in_range_mask))
-
-            # sort the candidates by line index; required for the seamlessnes np.split
-            sorted_begin_indices = np.argsort(begin_candidates[1])
-            sorted_end_indices = np.argsort(end_candidates[1])
-            begin_candidates = begin_candidates[:, sorted_begin_indices]
-            end_candidates = end_candidates[:, sorted_end_indices]
-
-            # create all possible connections between begin and end candidates that correspond to a line
-            grouped_begins = np.split(begin_candidates[0], np.unique(begin_candidates[1], return_index=True)[1][1:])
-            grouped_ends = np.split(end_candidates[0], np.unique(end_candidates[1], return_index=True)[1][1:])
-            line_indices = np.unique(begin_candidates[1])
-
-            # create all possible connections between begin and end candidates that correspond to a line
-            begin_vertex_list = []
-            end_vertex_list = []
-            line_idx_list = []
-            for begin_vertex, end_vertex, line_idx in zip(grouped_begins, grouped_ends, line_indices):
-                begin_vertex, end_vertex = np.meshgrid(begin_vertex, end_vertex)
-                begin_vertex_list.extend(begin_vertex.flatten())
-                end_vertex_list.extend(end_vertex.flatten())
-                line_idx_list.extend([line_idx] * len(begin_vertex.flatten()))
-
-            line_idx_list = np.array(line_idx_list)
-            all_connections = np.array([begin_vertex_list, end_vertex_list])
-
-            # decrease the number of possible connections to reduce number of calculations
-            possible_connections = np.unique(all_connections, axis=1)
-            possible_connections = np.sort(possible_connections, axis=0)
-            possible_connections = np.unique(possible_connections, axis=1)
-            possible_connections = possible_connections[:, possible_connections[0, :] != possible_connections[1, :]]
-
-            if possible_connections.shape[1] < 1:
-                continue
+        edges = np.array(edges).astype(np.float64)
+        if len(edges) < 1:
+            continue
+        # calculate the distances between the vertices and the edge ends
+
+        begin_edges = KDTree(edges[:, :2])
+        end_edges = KDTree(edges[:, 2:])
+
+        begin_indices = begin_edges.query_ball_tree(vertices, point_radius)
+        end_indices = end_edges.query_ball_tree(vertices, point_radius)
+
+        line_indices = np.where(np.array([len(i) and len(j) for i, j in zip(begin_indices, end_indices)]))[0]
 
-            # precalculate the possible direction vectors
-            possible_direction_vectors = vertices[possible_connections[0]] - vertices[possible_connections[1]]
-            possible_direction_vectors = possible_direction_vectors / np.linalg.norm(possible_direction_vectors,
-                                                                                     axis=1)[:, np.newaxis]
-
-            owned_lines_per_possible_connections = [list() for i in range(possible_connections.shape[1])]
-
-            # assign lines to possible connections
-            for line_idx, i, j in zip(line_idx_list, begin_vertex_list, end_vertex_list):
-                if i == j:
-                    continue
-                i, j = min(i, j), max(i, j)
-                for connection_idx, connection in enumerate(possible_connections.T):
-                    if np.all((i, j) == connection):
-                        owned_lines_per_possible_connections[connection_idx].append(line_idx)
-                        break
-
-            # check if the lines are in the same direction as the possible connection
-            for fitted_line_idx, owned_lines_per_possible_connection in enumerate(owned_lines_per_possible_connections):
-                line_deviations = np.abs(np.dot(line_directions[owned_lines_per_possible_connection],
-                                                possible_direction_vectors[fitted_line_idx]))
-                if np.any(line_deviations > deviation_threshold):
-                    connections.append(possible_connections[:, fitted_line_idx])
+        # create all possible connections between begin and end candidates that correspond to a line
+        begin_vertex_list = []
+        end_vertex_list = []
+        line_idx_list = []
+        for line_idx in line_indices:
+            begin_vertex, end_vertex = begin_indices[line_idx], end_indices[line_idx]
+            begin_vertex, end_vertex = np.meshgrid(begin_vertex, end_vertex)
+            begin_vertex_list.extend(begin_vertex.flatten())
+            end_vertex_list.extend(end_vertex.flatten())
+
+            line_idx_list.extend([line_idx] * len(begin_vertex.flatten()))
+
+        line_idx_list = np.array(line_idx_list)
+        all_connections = np.array([begin_vertex_list, end_vertex_list])
+
+        # decrease the number of possible connections to reduce number of calculations
+        possible_connections = np.unique(all_connections, axis=1)
+        possible_connections = np.sort(possible_connections, axis=0)
+        possible_connections = np.unique(possible_connections, axis=1)
+        possible_connections = possible_connections[:, possible_connections[0, :] != possible_connections[1, :]]
+
+        if possible_connections.shape[1] < 1:
+            continue
+
+        # precalculate the possible direction vectors
+        possible_direction_vectors = vertices.data[possible_connections[0]] - vertices.data[possible_connections[1]]
+        possible_direction_vectors = possible_direction_vectors / np.linalg.norm(possible_direction_vectors, axis=1)[:, np.newaxis]
+
+        owned_lines_per_possible_connections = [list() for i in range(possible_connections.shape[1])]
+
+        # assign lines to possible connections
+        for line_idx, i,j in zip(line_idx_list, begin_vertex_list, end_vertex_list):
+            if i == j:
+                continue
+            i, j = min(i, j), max(i, j)
+            for connection_idx, connection in enumerate(possible_connections.T):
+                if np.all((i, j) == connection):
+                    owned_lines_per_possible_connections[connection_idx].append(line_idx)
+                    break
+
+        # check if the lines are in the same direction as the possible connection
+        for fitted_line_idx, owned_lines_per_possible_connection in enumerate(owned_lines_per_possible_connections):
+            line_deviations = np.abs(np.dot(line_directions[owned_lines_per_possible_connection], possible_direction_vectors[fitted_line_idx]))
+            if np.any(line_deviations > deviation_threshold):
+                connections.append(possible_connections[:, fitted_line_idx])
 
     vertices = [{"xy": v, "type": "apex"} for v in apex_centroids]
     # vertices += [{"xy": v, "type": "apex"} for v in missed_vertices]
@@ -282,10 +267,8 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, color_range=4., poi
 
 def get_uv_depth(vertices, depth):
     '''Get the depth of the vertices from the depth image'''
-    uv = []
-    for v in vertices:
-        uv.append(v['xy'])
-    uv = np.array(uv)
+
+    uv = np.array([v['xy'] for v in vertices])
     uv_int = uv.astype(np.int32)
     H, W = depth.shape[:2]
     uv_int[:, 0] = np.clip(uv_int[:, 0], 0, W - 1)
@@ -387,7 +370,7 @@ def predict(entry, visualize=False, scale_estimation_coefficient=2.5, **kwargs)
         gest_seg = gest.resize(depth.size)
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
         # Metric3D
-        depth_np = np.array(depth) / scale_estimation_coefficient  # 2.5 is the scale estimation coefficient
+        depth_np = np.array(depth) / scale_estimation_coefficient
         vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs)
         if (len(vertices) < 2) or (len(connections) < 1):
             print(f'Not enough vertices or connections in image {i}')

test_solution.ipynb

@@ -6,8 +6,8 @@
    "metadata": {
     "collapsed": true,
     "ExecuteTime": {
-     "end_time": "2024-05-30T16:17:13.293274Z",
-     "start_time": "2024-05-30T16:17:10.165200Z"
+     "end_time": "2024-05-30T19:40:29.917386Z",
+     "start_time": "2024-05-30T19:40:26.483885Z"
     }
    },
    "source": [
@@ -44,8 +44,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-30T16:17:13.299692Z",
-     "start_time": "2024-05-30T16:17:13.294282Z"
+     "end_time": "2024-05-30T19:40:29.922957Z",
+     "start_time": "2024-05-30T19:40:29.918391Z"
     }
    },
    "cell_type": "code",
@@ -64,8 +64,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-30T16:17:13.303805Z",
-     "start_time": "2024-05-30T16:17:13.300698Z"
+     "end_time": "2024-05-30T19:40:29.927010Z",
+     "start_time": "2024-05-30T19:40:29.923961Z"
     }
    },
    "cell_type": "code",
@@ -83,8 +83,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-30T16:58:50.530994Z",
-     "start_time": "2024-05-30T16:57:19.315889Z"
+     "end_time": "2024-05-30T19:46:53.539149Z",
+     "start_time": "2024-05-30T19:45:17.379935Z"
     }
    },
    "cell_type": "code",
@@ -93,10 +93,16 @@
     "\n",
     "solution = []\n",
     "from concurrent.futures import ProcessPoolExecutor\n",
-    "with ProcessPoolExecutor(max_workers=12) as pool:\n",
+    "with ProcessPoolExecutor(max_workers=10) as pool:\n",
     "    results = []\n",
     "    for i, sample in enumerate(tqdm(dataset)):\n",
-    "        results.append(pool.submit(predict, sample,  point_radius=25, max_angle=15, extend=30, merge_th=3.0, min_missing_distance=10000.0, scale_estimation_coefficient=4))\n",
+    "        results.append(pool.submit(predict, sample,\n",
+    "                                   point_radius=25, \n",
+    "                                   max_angle=15, \n",
+    "                                   extend=30, \n",
+    "                                   merge_th=3.0, \n",
+    "                                   min_missing_distance=10000.0, \n",
+    "                                   scale_estimation_coefficient=4))\n",
     "\n",
     "    for i, result in enumerate(tqdm(results)):\n",
     "        key, pred_vertices, pred_edges = result.result()\n",
@@ -116,18 +122,18 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "346it [00:11, 30.58it/s] \n",
-      "100%|██████████| 346/346 [01:18<00:00,  4.40it/s]\n"
+      "346it [00:11, 29.44it/s] \n",
+      "100%|██████████| 346/346 [01:23<00:00,  4.15it/s]\n"
      ]
     }
    ],
-   "execution_count": 37
+   "execution_count": 7
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-05-30T16:58:53.239999Z",
-     "start_time": "2024-05-30T16:58:53.049258Z"
+     "end_time": "2024-05-30T19:42:15.830560Z",
+     "start_time": "2024-05-30T19:42:15.185713Z"
     }
    },
    "cell_type": "code",
@@ -160,12 +166,18 @@
        "DescribeResult(nobs=173, minmax=(1.3699674819647794, 3.507189015362208), mean=2.159528576281002, variance=0.19517651860816773, skewness=0.5493231777924736, kurtosis=-0.14096706768318912)"
       ]
      },
-     "execution_count": 38,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
-   "execution_count": 38
+   "execution_count": 5
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "best mean=2.159528576281002",
+   "id": "1d3cde94dcfc4c56"
   },
   {
    "metadata": {