a lot of changes, I want to see how it performs (rewrote evaluation to load data in batches)

handcrafted_solution.py

@@ -14,6 +14,12 @@ from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read
 from scipy.spatial import KDTree
 from scipy.spatial.distance import cdist
 from sklearn.cluster import DBSCAN
+from scipy.spatial import cKDTree
+
+from enum import Enum
+
+
+
 
 apex_color = gestalt_color_mapping["apex"]
 eave_end_point = gestalt_color_mapping["eave_end_point"]
@@ -23,7 +29,24 @@ apex_color, eave_end_point, flashing_end_point = [np.array(i) for i in [apex_col
 unclassified = np.array([(215, 62, 138)])
 line_classes = ['eave', 'ridge', 'rake', 'valley']
 
-
+class VertexType(Enum):
+    APEX = 0
+    EAVE_END_POINT = 1
+
+class NearestNDInterpolatorWithThreshold(si.NearestNDInterpolator):
+    def __init__(self, points, values, max_distance):
+        super().__init__(points, values)
+        self.max_distance = max_distance
+        self.tree = cKDTree(points)
+
+    def __call__(self, *args):
+        # Convert the input to a 2D array of query points
+        query_points = np.array(args).T
+        distances, indices = self.tree.query(query_points)
+        values = np.full(query_points.shape[:-1], np.nan)
+        valid_mask = distances <= self.max_distance
+        values[valid_mask] = self.values[indices[valid_mask]]
+        return values.T
 def empty_solution():
     '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
     return np.zeros((2, 3)), [(0, 1)]
@@ -70,15 +93,15 @@ def remove_undesired_objects(image):
 def clean_image(image_gestalt) -> np.ndarray:
     # clears image in from of unclassified and disconected components
     image_gestalt = np.array(image_gestalt)
-    unclassified_mask = cv2.inRange(image_gestalt, unclassified + 0.0, unclassified + 0.8)
-    unclassified_mask = cv2.bitwise_not(unclassified_mask)
-    mask = remove_undesired_objects(unclassified_mask).astype(np.uint8)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((11, 11), np.uint8), iterations=11)
-    mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, np.ones((11, 11), np.uint8), iterations=2)
-
-    image_gestalt[:, :, 0] *= mask
-    image_gestalt[:, :, 1] *= mask
-    image_gestalt[:, :, 2] *= mask
+    # unclassified_mask = cv2.inRange(image_gestalt, unclassified - 1, unclassified + 1)
+    # unclassified_mask = cv2.bitwise_not(unclassified_mask)
+    # mask = remove_undesired_objects(unclassified_mask).astype(np.uint8)
+    # mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((11, 11), np.uint8), iterations=11)
+    # mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, np.ones((11, 11), np.uint8), iterations=2)
+
+    # image_gestalt[:, :, 0] *= mask
+    # image_gestalt[:, :, 1] *= mask
+    # image_gestalt[:, :, 2] *= mask
     return image_gestalt
 
 
@@ -98,10 +121,10 @@ 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=4, stats=cv2.CV_32S)
-    *_, other_centroids = cv2.connectedComponentsWithStats(eave_end_point_mask, connectivity=4, stats=cv2.CV_32S)
+    *_, apex_stats, apex_centroids = cv2.connectedComponentsWithStats(apex_mask, connectivity=4, stats=cv2.CV_32S)
+    *_, other_stats, 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
+    return apex_centroids[1:], other_centroids[1:], apex_mask, eave_end_point_mask, apex_stats[1:, cv2.CC_STAT_WIDTH]/2, other_stats[1:, cv2.CC_STAT_WIDTH]/2
 
 
 def infer_vertices(image_gestalt, *, color_range=4.):
@@ -168,8 +191,8 @@ def get_lines_and_directions(gest_seg_np, edge_class, *, color_range=4., rho, th
 
             direction = extend * direction
 
-            x1, y1 = (-direction + (x1, y1)).astype(np.int32)
-            x2, y2 = (+ direction + (x2, y2)).astype(np.int32)
+            x1, y1 = -direction + (x1, y1)
+            x2, y2 = + direction + (x2, y2)
 
             edges.append((x1, y1, x2, y2))
     return edges, line_directions
@@ -189,14 +212,19 @@ def infer_missing_vertices(ridge_edges, rake_edges):
     return ridge_ends.data[missing_candidates]
 
 
-def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, max_angle=5.,
+def get_vertices_and_edges_from_segmentation(gest_seg_np, *,
+                                             point_radius=30,
+                                             max_angle=5.,
+                                             point_radius_scale=1,
                                              **kwargs):
     '''Get the vertices and edges from the gestalt segmentation mask of the house'''
     # Apex
     connections = []
     deviation_threshold = np.cos(np.deg2rad(max_angle))
 
-    apex_centroids, eave_end_point_centroids, apex_mask, eave_end_point_mask = get_vertices(gest_seg_np)
+    (apex_centroids, eave_end_point_centroids,
+     apex_mask, eave_end_point_mask,
+     apex_radii, eave_radii) = get_vertices(gest_seg_np)
 
     vertices = np.concatenate([apex_centroids, eave_end_point_centroids])
     # inferred_vertices, inferred_mask = infer_vertices(gest_seg_np)
@@ -206,12 +234,6 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, ma
     if len(vertices) < 2:
         return [], []
 
-    # 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
 
     edges = []
     line_directions = []
@@ -251,6 +273,19 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, ma
         missed_vertices = get_missed_vertices(vertices, inferred_vertices, **kwargs)
         vertices = np.concatenate([vertices, missed_vertices])
 
+    vertex_size = np.full(len(vertices), point_radius/2)
+    apex_radii *= point_radius_scale
+    eave_radii *= point_radius_scale
+    apex_radii = np.clip(apex_radii, 10, point_radius)
+    eave_radii = np.clip(eave_radii, 10, point_radius)
+    vertex_size[:len(apex_radii)] = apex_radii
+    vertex_size[len(apex_radii):len(apex_radii) + len(eave_radii)] = eave_radii
+
+    # 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
+
     vertices = KDTree(vertices)
 
     for edge_class in ['eave',
@@ -291,12 +326,29 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, ma
     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())
+        begin_vertices, end_vertices = begin_indices[line_idx], end_indices[line_idx]
+        begin_vertices, end_vertices = np.array(begin_vertices), np.array(end_vertices)
+        begin_value = begin_edges.data[line_idx]
+        end_value = end_edges.data[line_idx]
+        begin_in_range_indices = np.where(
+            np.linalg.norm(vertices.data[begin_vertices] - begin_value, axis=1)
+            <
+            vertex_size[begin_vertices])[0]
+        end_in_range_indices = np.where(
+            np.linalg.norm(vertices.data[end_vertices] - end_value, axis=1)
+            <
+            vertex_size[end_vertices])[0]
+        begin_vertices = begin_vertices[begin_in_range_indices]
+        end_vertices = end_vertices[end_in_range_indices]
+        if len(begin_vertices) < 1 or len(end_vertices) < 1:
+            continue
+
 
-        line_idx_list.extend([line_idx] * len(begin_vertex.flatten()))
+        begin_vertices, end_vertices = np.meshgrid(begin_vertices, end_vertices)
+        begin_vertex_list.extend(begin_vertices.flatten())
+        end_vertex_list.extend(end_vertices.flatten())
+
+        line_idx_list.extend([line_idx] * len(begin_vertices.flatten()))
 
     line_idx_list = np.array(line_idx_list)
     all_connections = np.array([begin_vertex_list, end_vertex_list])
@@ -334,9 +386,9 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, ma
         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]
-    vertices += [{"xy": v, "type": "eave_end_point"} for v in eave_end_point_centroids]
+    vertices = [{"xy": v, "type": VertexType.APEX} for v in apex_centroids]
+    vertices += [{"xy": v, "type": VertexType.APEX} for v in missed_vertices]
+    vertices += [{"xy": v, "type": VertexType.EAVE_END_POINT} for v in eave_end_point_centroids]
     return vertices, connections
 
 
@@ -376,7 +428,7 @@ def merge_vertices_3d(vert_edge_per_image, merge_th=0.1, **kwargs):
         new_vertex_mapping = dict(zip(left_vertex_indices, new_indices))
 
         vertices = [v for i, v in enumerate(vertices) if i in new_vertex_mapping]
-        types += [int(v['type'] == 'apex') for v in vertices]
+        types += [int(v['type'] == VertexType.APEX) for v in vertices]
         vertices_3d = vertices_3d[left_vertex_indices]
         connections = [[new_vertex_mapping[a] + cur_start, new_vertex_mapping[b] + cur_start] for a, b in connections]
 
@@ -454,8 +506,13 @@ def prune_not_connected(all_3d_vertices, connections_3d):
     return np.array(new_verts), connected_out
 
 
-def predict(entry, visualize=False, scale_estimation_coefficient=2.5, clustering_eps=100, dist_coeff=0, pointcloud_depth_coeff = 1, **kwargs) -> Tuple[
-    np.ndarray, List[int]]:
+def predict(entry, visualize=False,
+            scale_estimation_coefficient=2.5,
+            clustering_eps=100,
+            dist_coeff=0,
+            pointcloud_depth_coeff = 1,
+            interpolation_radius=200,
+            **kwargs) -> Tuple[np.ndarray, List[int]]:
     if 'gestalt' not in entry or 'depthcm' not in entry or 'K' not in entry or 'R' not in entry or 't' not in entry:
         print('Missing required fields in the entry')
         return (entry['__key__'], *empty_solution())
@@ -486,10 +543,13 @@ def predict(entry, visualize=False, scale_estimation_coefficient=2.5, clustering
     clustered_keys = np.split(point_keys, cluster_indices[1:])
 
     biggest_cluster_index = np.argmax([len(i) for i in clustered_points])
-    # biggest_cluster = clustered_points[biggest_cluster_index]
+    biggest_cluster = clustered_points[biggest_cluster_index]
     biggest_cluster_keys = clustered_keys[biggest_cluster_index]
     biggest_cluster_keys = set(biggest_cluster_keys)
 
+    points3d_kdtree = KDTree(biggest_cluster)
+
+
     for i, (gest, depthcm, K, R, t, imagekey) in enumerate(zip(entry['gestalt'],
                                                                entry['depthcm'],
                                                                entry['K'],
@@ -498,89 +558,127 @@ def predict(entry, visualize=False, scale_estimation_coefficient=2.5, clustering
                                                                entry['__imagekey__']
                                                                )):
 
+        gest_seg = gest.resize(depthcm.size)
+        gest_seg_np = np.array(gest_seg).astype(np.uint8)
+        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}')
+            vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
+            continue
+        depth_np = np.array(depthcm) / scale_estimation_coefficient
+        uv, depth_vert_from_depth_map = get_uv_depth(vertices, depth_np)
         try:
-            # gest_seg = gest.resize(depthcm.size)
-            gest_seg_np = np.array(gest).astype(np.uint8)
-            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}')
-                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
-                continue
-            belonging_points = []
-            for point_id in image_dict[imagekey].point3D_ids[np.where(image_dict[imagekey].point3D_ids != -1)]:
+            belonging_points3d = []
+            belonging_points2d = []
+            point_indices = np.where(image_dict[imagekey].point3D_ids != -1)[0]
+            for idx, point_id in zip(point_indices, image_dict[imagekey].point3D_ids[point_indices]):
                 if point_id in biggest_cluster_keys:
-                    belonging_points.append(entry["points3d"][point_id])
+                    belonging_points3d.append(entry["points3d"][point_id].xyz)
+                    belonging_points2d.append(image_dict[imagekey].xys[idx])
 
-            if len(belonging_points) < 1:
+            if len(belonging_points3d) < 1:
                 print(f'No 3D points in image {i}')
                 vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
                 raise KeyError
-            projected2d, _ = cv2.projectPoints(np.array([i.xyz for i in belonging_points]), R, t, K, dist_coeff)
-            important = np.where(np.all(projected2d >= 0, axis=2))
+            belonging_points3d = np.array(belonging_points3d)
+            belonging_points2d = np.array(belonging_points2d)
+            # projected2d, _ = cv2.projectPoints(belonging_points3d, R, t, K, dist_coeff)
+            important = np.where(np.all(belonging_points2d >= 0, axis=1))
             # Normalize the uv to the camera intrinsics
             world_to_cam = np.eye(4)
             world_to_cam[:3, :3] = R
             world_to_cam[:3, 3] = t
 
-            homo_belonging_points = cv2.convertPointsToHomogeneous(np.array([i.xyz for i in belonging_points]))
+            homo_belonging_points = cv2.convertPointsToHomogeneous(belonging_points3d)
             depth = cv2.convertPointsFromHomogeneous(cv2.transform(homo_belonging_points, world_to_cam))
             depth = depth[:, 0, 2]
-            depth = depth[important[0]] * pointcloud_depth_coeff
-            projected2d = projected2d[important]
+            # projected2d = projected2d[:, 0, :]
+            depth = depth[important[0]]
+            # projected2d = projected2d[important[0]]
+            projected2d = belonging_points2d[important[0]]
+            # print(projected2d.shape)
+            # print(depth.shape)
+            depth *= pointcloud_depth_coeff
             if len(depth) < 1:
                 print(f'No 3D points in image {i}')
                 vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
                 raise KeyError
             # print(projected2d.shape, depth.shape)
 
-            interpolator = si.NearestNDInterpolator(projected2d, depth, rescale=True)
-            # interpolator = si.CloughTocher2DInterpolator(projected2d, depth, np.nan)
+            # interpolator = si.NearestNDInterpolator(projected2d, depth, rescale=True)
+            interpolator = NearestNDInterpolatorWithThreshold(projected2d, depth, interpolation_radius)
+            # interpolator = si.LinearNDInterpolator(projected2d, depth, np.nan)
 
-            vertex_coordinates = np.array([v['xy'] for v in vertices])
-            xi, yi = vertex_coordinates[:, 0], vertex_coordinates[:, 1]
-            depth_vert = interpolator(xi, yi)
-            xy_local = np.ones((len(vertex_coordinates), 3))
-            xy_local[:, 0] = (vertex_coordinates[:, 0] - K[0, 2]) / K[0, 0]
-            xy_local[:, 1] = (vertex_coordinates[:, 1] - K[1, 2]) / K[1, 1]
-            # Get the 3D vertices
-            vertices_3d_local = depth_vert[..., None] * (xy_local / np.linalg.norm(xy_local, axis=1)[..., None])
-            world_to_cam = np.eye(4)
-            world_to_cam[:3, :3] = R
-            world_to_cam[:3, 3] = t.reshape(-1)
-            cam_to_world = np.linalg.inv(world_to_cam)
-            vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
-            vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
+            uv = np.array([v['xy'] for v in vertices])
+            xi, yi = uv[:, 0], uv[:, 1]
+            depth_vert_from_pointcloud = interpolator(xi, yi)
+            depthmap_used = False
 
+            # Get the 3D vertices
         except KeyError:
-            gest_seg = gest.resize(depthcm.size)
-            gest_seg_np = np.array(gest_seg).astype(np.uint8)
+            #Revert to the depthmap
             # Metric3D
-            depth_np = np.array(depthcm) / scale_estimation_coefficient
-            cv2.GaussianBlur(depth_np, (21, 21), 1, depth_np)
-            # cv2.medianBlur(depth_np, 5)
-            # depth_np = np.zeros_like(depth_np)
-            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}')
-                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
-                continue
-            uv, depth_vert = get_uv_depth(vertices, depth_np)
+            depthmap_used = True
+
             # Normalize the uv to the camera intrinsics
-            xy_local = np.ones((len(uv), 3))
-            xy_local[:, 0] = (uv[:, 0] - K[0, 2]) / K[0, 0]
-            xy_local[:, 1] = (uv[:, 1] - K[1, 2]) / K[1, 1]
-            # Get the 3D vertices
-            vertices_3d_local = depth_vert[..., None] * (xy_local / np.linalg.norm(xy_local, axis=1)[..., None])
-            world_to_cam = np.eye(4)
-            world_to_cam[:3, :3] = R
-            world_to_cam[:3, 3] = t.reshape(-1)
-            cam_to_world = np.linalg.inv(world_to_cam)
-            vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
-            vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
+
+        xy_local = np.ones((len(uv), 3))
+        xy_local[:, 0] = (uv[:, 0] - K[0, 2]) / K[0, 0]
+        xy_local[:, 1] = (uv[:, 1] - K[1, 2]) / K[1, 1]
+        # Get the 3D vertices
+
+        depth_vert_nan_idxs = None
+        if depthmap_used:
+            # norm_factor = np.max(np.linalg.norm(xy_local, axis=1)[..., None])
+            depth_vert = depth_vert_from_depth_map
+        else:
+            # 1. query detected vertices in projected2d
+            # if the vertex is beyond some radius, use the depthmap
+            # isnt uv
+            depth_vert_nan_idxs = np.where(np.isnan(depth_vert_from_pointcloud))[0]
+            depth_vert_from_pointcloud[depth_vert_nan_idxs] = depth_vert_from_depth_map[depth_vert_nan_idxs]
+            depth_vert = depth_vert_from_pointcloud
+
+        norm_factor = np.linalg.norm(xy_local, axis=1)[..., None]
+        if depth_vert_nan_idxs is not None and len(depth_vert_nan_idxs) > 0:
+            norm_factor_min = np.min(norm_factor[depth_vert_nan_idxs])
+            if len(depth_vert_nan_idxs) != len(norm_factor):
+                norm_factor_max = np.max(norm_factor[~np.isin(np.arange(len(norm_factor)), depth_vert_nan_idxs)])
+            else:
+                norm_factor_max = np.max(norm_factor)
+        else:
+            norm_factor_min = np.min(norm_factor)
+            norm_factor_max = np.max(norm_factor)
+
+        vertices_3d_local = depth_vert[..., None] * xy_local
+        if depthmap_used:
+            vertices_3d_local /= norm_factor_max
+        else:
+            vertices_3d_local[depth_vert_nan_idxs] /= norm_factor_max
+            vertices_3d_local[~np.isin(np.arange(len(vertices_3d_local)), depth_vert_nan_idxs)] /= norm_factor_min
+
+        # vertices_3d_local = depth_vert[..., None] * (xy_local / norm_factor)
+        world_to_cam = np.eye(4)
+        world_to_cam[:3, :3] = R
+        world_to_cam[:3, 3] = t
+
+        cam_to_world = np.linalg.inv(world_to_cam)
+        vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
+        vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
+
+        # not_nan_items = np.all(~np.isnan(vertices_3d), axis=1)
+        # _, closest_fitted = points3d_kdtree.query(vertices_3d[not_nan_items])
+
+        # vertices_3d[not_nan_items] = points3d_kdtree.data[closest_fitted]
 
         vert_edge_per_image[i] = vertices, connections, vertices_3d
     all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, **kwargs)
     all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
+    # highest_edges = np.argpartition(all_3d_vertices_clean[:, 1], 4)[:4].tolist()
+    #
+    # connections_3d_clean.append(highest_edges[:2])
+    # connections_3d_clean.append(highest_edges[2:])
     # all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d)
     if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
         print(f'Not enough vertices or connections in the 3D vertices')

script.py

@@ -116,7 +116,18 @@ def save_submission(submission, path):
     sub = pd.DataFrame(submission, columns=["__key__", "wf_vertices", "wf_edges"])
     sub.to_parquet(path)
     print(f"Submission saved to {path}")
-
+batch_size = 48  # You can adjust this according to your needs
+
+# Define a generator function to yield batches of samples
+def batch_generator(dataset, batch_size):
+    batch = []
+    for i, sample in enumerate(dataset):
+        batch.append(sample)
+        if len(batch) == batch_size:
+            yield batch
+            batch = []
+    if batch:  # Yield the remaining samples
+        yield batch
 if __name__ == "__main__":
     from handcrafted_solution import predict
     print ("------------ Loading dataset------------ ")
@@ -127,28 +138,30 @@ if __name__ == "__main__":
     solution = []
     from concurrent.futures import ProcessPoolExecutor
     with ProcessPoolExecutor(max_workers=8) as pool:
-        results = []
-        for i, sample in enumerate(tqdm(dataset)):
-            results.append(pool.submit(predict, sample,
-                                       visualize=False,
-                                       point_radius=15,
-                                       max_angle=5,
-                                       extend=25,
-                                       merge_th=100.0,
-                                       min_missing_distance=350.0,
-                                       scale_estimation_coefficient=2.54,
-                                       clustering_eps=100,
-                                       dist_coeff=0.1,
-                                       pointcloud_depth_coeff=1.05,
-                                       ))
-        
-        for i, result in enumerate(tqdm(results)):
-            key, pred_vertices, pred_edges = result.result()
-            solution.append({
-                            '__key__': key,
-                            'wf_vertices': pred_vertices.tolist(),
-                            'wf_edges': pred_edges
-                        })
+        for batch in tqdm(batch_generator(dataset, batch_size), desc='Batches'):
+            results = []
+            for i, sample in enumerate(batch):
+                results.append(pool.submit(predict, sample,
+                                           point_radius=40,
+                                           max_angle=15,
+                                           extend=25,
+                                           merge_th=80.0,
+                                           min_missing_distance=350.0,
+                                           scale_estimation_coefficient=2.54,
+                                           clustering_eps=120,
+                                           interpolation_radius=200,
+                                           point_radius_scale=0.5,
+                                           # dist_coeff=0,
+                                           # pointcloud_depth_coeff=1,
+                                           ))
+
+            for result in tqdm(results, desc='Results', total=len(results), position=0):
+                key, pred_vertices, pred_edges = result.result()
+                solution.append({
+                    '__key__': key,
+                    'wf_vertices': pred_vertices.tolist(),
+                    'wf_edges': pred_edges
+                })
             if i % 100 == 0:
                 # incrementally save the results in case we run out of time
                 print(f"Processed {i} samples")