Improves wireframe extraction and prediction

Refactors wireframe extraction by merging overlapping point groups and updating connections.

Adds edge patch generation for training edge prediction models using point clouds near wireframe edges and non-edges, improving data creation and model training.

Adds visualization tools for point clouds and predicted vertices with connections to help debugging and evaluation.

Simplifies and streamlines vertex prediction by applying a PointNet model to extracted patches and filtering zero vertices.

predict.py

@@ -12,10 +12,17 @@ import open3d as o3d
 from visu import plot_reconstruction_local, plot_wireframe_local, plot_bpo_cameras_from_entry_local
 import pyvista as pv
 from fast_pointnet import save_patches_dataset, predict_vertex_from_patch
+from fast_voxel import predict_vertex_from_patch_voxel
+import time
+from fast_pointnet_class import save_patches_dataset as save_patches_dataset_class
 
-GENERATE_DATASET = True
-#DATASET_DIR = '/home/skvrnjan/personal/hohocustom/'
-DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom'
+GENERATE_DATASET = False
+DATASET_DIR = '/home/skvrnjan/personal/hohocustom/'
+#DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom'
+
+GENERATE_DATASET_EDGES = True
+#EDGES_DATASET_DIR = '/home/skvrnjan/personal/hohocustom_edges/'
+EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges'
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -544,7 +551,7 @@ def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points, all_connections)
     # and update connections accordingly
     updated_connections = []
     if extracted_points:
-        print(f"Merging groups based on point overlap... Processing {len(extracted_points)} groups")
+        #print(f"Merging groups based on point overlap... Processing {len(extracted_points)} groups")
         # Create a list to track which groups to keep
         groups_to_keep = []
         merged_groups = set()  # Track which groups have been merged
@@ -579,19 +586,36 @@ def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points, all_connections)
                         merged_points = np.vstack([merged_points, extracted_points[j]]) if len(merged_points) > 0 else extracted_points[j]
                         merged_colors = np.vstack([merged_colors, extracted_colors[j]]) if len(merged_colors) > 0 else extracted_colors[j]
                         merged_ids.update(ids_j)
+                        merged_indices.append(j)
                         merged_groups.add(j)
             
             # Add the merged group to the list of groups to keep
             if len(merged_points) > 0:
+                new_group_idx = len(groups_to_keep)
                 groups_to_keep.append((merged_points, merged_colors, np.array(list(merged_ids))))
+                
+                # Update mapping for all merged indices
+                for old_idx in merged_indices:
+                    old_to_new_mapping[old_idx] = new_group_idx
         
         # Update extracted_points, extracted_colors, and extracted_ids with filtered results
         extracted_points = [group[0] for group in groups_to_keep]
         extracted_colors = [group[1] for group in groups_to_keep]
         extracted_ids = [group[2] for group in groups_to_keep]
 
-        print(f"After merging, number of groups: {len(extracted_points)}")
+        # Update connections based on the new mapping
+        for start_idx, end_idx in all_flattened_connections:
+            if start_idx in old_to_new_mapping and end_idx in old_to_new_mapping:
+                new_start = old_to_new_mapping[start_idx]
+                new_end = old_to_new_mapping[end_idx]
+                # Only add connection if vertices are still different after merging
+                if new_start != new_end:
+                    connection = tuple(sorted((new_start, new_end)))
+                    if connection not in updated_connections:
+                        updated_connections.append(connection)
 
+        #print(f"After merging, number of groups: {len(extracted_points)}")
+        #print(f"Updated connections: {updated_connections}")
 
     # Create visualization showing extracted points for each group as balls within their mean
     if False:
@@ -629,17 +653,481 @@ def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points, all_connections)
             
             plotter.show(title=f"Extracted Points within {ball_radius}m radius - Spheres at group means")
 
-    return extracted_points, extracted_colors, extracted_ids, whole_pcloud
+    return extracted_points, extracted_colors, extracted_ids, whole_pcloud, updated_connections
+
+def visu_pcloud_and_preds(colmap_rec, extracted_ids, extracted_points, extracted_colors, predicted_vertices, connections):
+    if extracted_ids:
+        plotter = pv.Plotter()
+        
+        # Add all COLMAP points in gray
+        all_points = []
+        all_colors = []
+        for pid, p3D in colmap_rec.points3D.items():
+            all_points.append(p3D.xyz)
+            all_colors.append([0.8, 0.8, 0.8])  # Gray color
+
+        if all_points:
+            all_points = np.array(all_points)
+            all_colors = np.array(all_colors)
+            point_cloud = pv.PolyData(all_points)
+            point_cloud["colors"] = np.array(all_colors)
+            plotter.add_mesh(point_cloud, scalars="colors", rgb=True, point_size=3, render_points_as_spheres=True)
+        
+        for group_idx, (group_points, group_colors) in enumerate(zip(extracted_points, extracted_colors)):
+            if len(group_points) > 0:
+                # Calculate mean position for this group
+                group_mean = np.mean(group_points, axis=0)
+                
+                # Create a sphere at the mean position
+                sphere = pv.Sphere(radius=0.2, center=group_mean)
+                # Generate a random color for each group
+                group_color = np.random.rand(3)
+                plotter.add_mesh(sphere, color=group_color, opacity=0.5)
+                
+                # Add the extracted points for this group in the same color
+                group_cloud = pv.PolyData(group_points)
+                plotter.add_mesh(group_cloud, color=group_color, point_size=6, render_points_as_spheres=True)
+
+                # Add predicted vertex as sphere if it exists and is valid
+                if group_idx < len(predicted_vertices):
+                    pred_vertex = predicted_vertices[group_idx]
+                    if not np.allclose(pred_vertex, [0.0, 0.0, 0.0]):  # Check if it's not a zero vertex
+                        pred_sphere = pv.Sphere(radius=0.15, center=pred_vertex)
+                        plotter.add_mesh(pred_sphere, color="black", opacity=1.)
+
+        # Add connections between predicted vertices
+        if len(predicted_vertices) > 0 and len(connections) > 0:
+            valid_pred_vertices = []
+            valid_indices = []
+            for i, pred_vertex in enumerate(predicted_vertices):
+                if not np.allclose(pred_vertex, [0.0, 0.0, 0.0]):
+                    valid_pred_vertices.append(pred_vertex)
+                    valid_indices.append(i)
+            
+            if len(valid_pred_vertices) > 1:
+                valid_pred_vertices = np.array(valid_pred_vertices)
+                
+                # Create lines for connections
+                for start_idx, end_idx in connections:
+                    if start_idx in valid_indices and end_idx in valid_indices:
+                        # Map to valid vertex indices
+                        valid_start = valid_indices.index(start_idx)
+                        valid_end = valid_indices.index(end_idx)
+                        
+                        # Create line between vertices
+                        line_points = np.array([valid_pred_vertices[valid_start], valid_pred_vertices[valid_end]])
+                        line = pv.Line(line_points[0], line_points[1])
+                        plotter.add_mesh(line, color="red", line_width=3)
+
+        ball_radius = 1.0  # meters
+        plotter.show(title=f"Extracted Points within {ball_radius}m radius - Spheres at group means")
+
+def generate_edge_patches(frame):
+    vertices = frame['wf_vertices']
+    connections = frame['wf_edges']
+
+    vertices = np.array(vertices) if vertices else np.empty((0, 3))
+
+    positive_patches = []
+    negative_patches = []
+
+    cylinder_radius = 0.5
+
+    colmap = frame['colmap_binary']
+
+    # Create 6D point cloud from COLMAP data
+    colmap_points_6d = []
+    for pid, p3D in colmap.points3D.items():
+        # Combine xyz coordinates and RGB color
+        point_6d = np.concatenate([p3D.xyz, p3D.color / 255.0])  # Normalize color to [0,1]
+        colmap_points_6d.append(point_6d)
+
+    colmap_points_6d = np.array(colmap_points_6d) if colmap_points_6d else np.empty((0, 6))
+    
+    colmap_points_6d[:, 3:] = colmap_points_6d[:, 3:] * 2 - 1
+
+    # Extract 3D coordinates for faster vectorized operations
+    colmap_points_3d = colmap_points_6d[:, :3]
+
+    # For each connection, create a positive edge patch
+    for connection in connections:
+        start_idx, end_idx = connection
+        
+        # Get start and end vertices from the connections
+        start_vertex = vertices[start_idx]
+        end_vertex = vertices[end_idx]
+        
+        # Create line vector from start to end
+        line_vector = end_vertex - start_vertex
+        line_length = np.linalg.norm(line_vector)
+        
+        # Normalize line vector
+        line_direction = line_vector / line_length
+        
+        # Extend the line by 25 cm (0.25 meters) on both ends for more context
+        extension_length = 0.25  # 25 cm in meters
+        extended_start = start_vertex - extension_length * line_direction
+        extended_end = end_vertex + extension_length * line_direction
+        extended_line_length = line_length + 2 * extension_length
+        
+        # Vectorized distance calculation
+        # Vector from extended start to all points
+        start_to_points = colmap_points_3d - extended_start[np.newaxis, :]
+        
+        # Project onto line direction to get distance along extended line
+        projection_lengths = np.dot(start_to_points, line_direction)
+        
+        # Filter points within extended line segment bounds
+        within_bounds = (projection_lengths >= 0) & (projection_lengths <= extended_line_length)
+        
+        # Find closest points on extended line segment for all points
+        closest_points_on_line = extended_start[np.newaxis, :] + projection_lengths[:, np.newaxis] * line_direction[np.newaxis, :]
+        
+        # Calculate perpendicular distances from points to line
+        perpendicular_distances = np.linalg.norm(colmap_points_3d - closest_points_on_line, axis=1)
+        
+        # Find points within cylinder
+        within_cylinder = within_bounds & (perpendicular_distances <= cylinder_radius)
+        
+        if np.sum(within_cylinder) <= 10:
+            continue
+
+        points_in_cylinder = colmap_points_6d[within_cylinder]
+        point_indices_in_cylinder = np.where(within_cylinder)[0]
+        
+        # Center the patch at the midpoint of the original line (not extended)
+        line_midpoint = (start_vertex + end_vertex) / 2
+        
+        # Shift points to center around origin
+        points_centered = points_in_cylinder.copy()
+        points_centered[:, :3] -= line_midpoint
+        
+        # Create positive edge patch
+        positive_patch = {
+            'patch_6d': points_centered,
+            'connection': connection,
+            'line_start': start_vertex - line_midpoint,
+            'line_end': end_vertex - line_midpoint,
+            'cylinder_radius': cylinder_radius,
+            'point_indices': point_indices_in_cylinder,
+            'label': 1,  # Positive label for edge
+            'center': line_midpoint
+        }
+        
+        
+        positive_patches.append(positive_patch)
+
+    # Generate negative edge patches by sampling random unconnected vertex pairs
+    num_negative_patches = len(positive_patches)
+    
+    if num_negative_patches > 0 and len(vertices) >= 2:
+        # Create set of connected pairs for fast lookup
+        connected_pairs = set(tuple(sorted(conn)) for conn in connections)
+        
+        # Generate all possible vertex pairs
+        vertex_indices = np.arange(len(vertices))
+        all_pairs = np.array(np.meshgrid(vertex_indices, vertex_indices)).T.reshape(-1, 2)
+        
+        # Filter out pairs where both indices are the same
+        all_pairs = all_pairs[all_pairs[:, 0] != all_pairs[:, 1]]
+        
+        # Sort pairs to match connected_pairs format
+        all_pairs_sorted = np.sort(all_pairs, axis=1)
+        
+        # Find unconnected pairs
+        unconnected_mask = np.array([tuple(pair) not in connected_pairs for pair in all_pairs_sorted])
+        unconnected_pairs = all_pairs[unconnected_mask]
+        
+        if len(unconnected_pairs) > 0:
+            # Pre-compute positive patch cylinder info for overlap checks
+            positive_cylinders = []
+            for pos_patch in positive_patches:
+                start_world = pos_patch['line_start'] + pos_patch['center']
+                end_world = pos_patch['line_end'] + pos_patch['center']
+                positive_cylinders.append({
+                    'start': start_world,
+                    'end': end_world,
+                    'radius': pos_patch['cylinder_radius']
+                })
+            
+            # Randomly sample negative pairs without replacement
+            num_to_sample = min(num_negative_patches * 3, len(unconnected_pairs))  # Sample more to account for rejections
+            sampled_indices = np.random.choice(len(unconnected_pairs), size=num_to_sample, replace=False)
+            sampled_pairs = unconnected_pairs[sampled_indices]
+            
+            for idx1, idx2 in sampled_pairs:
+                if len(negative_patches) >= num_negative_patches:
+                    break
+                    
+                start_vertex = vertices[idx1]
+                end_vertex = vertices[idx2]
+                
+                # Create line vector from start to end
+                line_vector = end_vertex - start_vertex
+                line_length = np.linalg.norm(line_vector)
+                
+                # Normalize line vector
+                line_direction = line_vector / line_length
+                
+                # Extend the line by 25 cm (0.25 meters) on both ends for more context
+                extension_length = 0.25  # 25 cm in meters
+                extended_start = start_vertex - extension_length * line_direction
+                extended_end = end_vertex + extension_length * line_direction
+                extended_line_length = line_length + 2 * extension_length
+                
+                # Check cylinder overlap with positive patches
+                current_cylinder = {
+                    'start': extended_start,
+                    'end': extended_end,
+                    'radius': cylinder_radius
+                }
+                
+                has_overlap = False
+                for pos_cylinder in positive_cylinders:
+                    # Calculate cylinder-cylinder intersection volume
+                    overlap_volume = calculate_cylinder_overlap_volume(current_cylinder, pos_cylinder)
+                    
+                    # Calculate volumes of both cylinders
+                    current_volume = np.pi * cylinder_radius**2 * extended_line_length
+                    pos_height = np.linalg.norm(pos_cylinder['end'] - pos_cylinder['start'])
+                    pos_volume = np.pi * pos_cylinder['radius']**2 * pos_height
+                    
+                    # Calculate IoU
+                    union_volume = current_volume + pos_volume - overlap_volume
+                    if union_volume > 0:
+                        iou = overlap_volume / union_volume
+                        if iou > 0.5:  # 0.2 IoU threshold
+                            has_overlap = True
+                            break
+                
+                if has_overlap:
+                    continue  # Skip this negative patch due to cylinder overlap
+                
+                # Vectorized distance calculation
+                # Vector from extended start to all points
+                start_to_points = colmap_points_3d - extended_start[np.newaxis, :]
+                
+                # Project onto line direction to get distance along extended line
+                projection_lengths = np.dot(start_to_points, line_direction)
+                
+                # Filter points within extended line segment bounds
+                within_bounds = (projection_lengths >= 0) & (projection_lengths <= extended_line_length)
+                
+                # Find closest points on extended line segment for all points
+                closest_points_on_line = extended_start[np.newaxis, :] + projection_lengths[:, np.newaxis] * line_direction[np.newaxis, :]
+                
+                # Calculate perpendicular distances from points to line
+                perpendicular_distances = np.linalg.norm(colmap_points_3d - closest_points_on_line, axis=1)
+                
+                # Find points within cylinder
+                within_cylinder = within_bounds & (perpendicular_distances <= cylinder_radius)
+                
+                if np.sum(within_cylinder) <= 10:
+                    continue
+
+                points_in_cylinder = colmap_points_6d[within_cylinder]
+                point_indices_in_cylinder = np.where(within_cylinder)[0]
+                
+                # Center the patch at the midpoint of the original line (not extended)
+                line_midpoint = (start_vertex + end_vertex) / 2
+                
+                # Shift points to center around origin
+                points_centered = points_in_cylinder.copy()
+                points_centered[:, :3] -= line_midpoint
+                
+                # Create negative edge patch
+                negative_patch = {
+                    'patch_6d': points_centered,
+                    'connection': (idx1, idx2),
+                    'line_start': start_vertex - line_midpoint,
+                    'line_end': end_vertex - line_midpoint,
+                    'cylinder_radius': cylinder_radius,
+                    'point_indices': point_indices_in_cylinder,
+                    'label': 0,  # Negative label for non-edge
+                    'center': line_midpoint  # Center of the patch
+                }
+                
+                negative_patches.append(negative_patch)
+
+    print(f"Generated {len(positive_patches)} positive patches and {len(negative_patches)} negative patches")
+    all_patches = positive_patches + negative_patches
+
+    # Visualize edge patches
+    if False:  # Set to True to enable visualization
+        # Create plotter
+        plotter = pv.Plotter()
+        
+        # Add whole point cloud in gray
+        if len(colmap_points_6d) > 0:
+            whole_cloud = pv.PolyData(colmap_points_3d)
+            gray_colors = np.full((len(colmap_points_3d), 3), [0.5, 0.5, 0.5])
+            whole_cloud["colors"] = gray_colors
+            plotter.add_mesh(whole_cloud, scalars="colors", rgb=True, point_size=3, render_points_as_spheres=True)
+        
+        # Visualize each patch
+        for patch_idx, patch in enumerate(all_patches):
+            # Use green for positive (edge), red for negative (non-edge)
+            patch_color = 'green' if patch['label'] == 1 else 'red'
+            
+            # Get patch data
+            points_in_cylinder = patch['patch_6d'][:, :3]  # xyz coordinates
+            line_start = patch['line_start']
+            line_end = patch['line_end']
+            center = patch['center']  # Use center instead of calculating midpoint
+            
+            # Shift points back to world coordinates for visualization
+            points_world = points_in_cylinder + center
+            
+            # Add points inside cylinder with patch-specific color
+            if len(points_world) > 0:
+                cylinder_cloud = pv.PolyData(points_world)
+                plotter.add_mesh(cylinder_cloud, color=patch_color, point_size=8, render_points_as_spheres=True)
+            
+            # Add start and end points as larger spheres
+            start_sphere = pv.Sphere(radius=0.1, center=line_start + center)
+            end_sphere = pv.Sphere(radius=0.1, center=line_end + center)
+            plotter.add_mesh(start_sphere, color='black', opacity=0.8)
+            plotter.add_mesh(end_sphere, color='white', opacity=0.8)
+            
+            # Add line between start and end
+            line_points = np.array([line_start + center, line_end + center])
+            line = pv.Line(line_points[0], line_points[1])
+            plotter.add_mesh(line, color=patch_color, line_width=5)
+            
+            # Add cylinder wireframe to show extraction bounds
+            cylinder_center = center
+            cylinder_direction = (line_end - line_start) / np.linalg.norm(line_end - line_start)
+            cylinder_height = np.linalg.norm(line_end - line_start) + 2 * 0.25  # Including extensions
+            
+            # Create cylinder mesh for visualization
+            cylinder_mesh = pv.Cylinder(center=cylinder_center, direction=cylinder_direction, 
+                                      radius=patch['cylinder_radius'], height=cylinder_height)
+            plotter.add_mesh(cylinder_mesh, color=patch_color, opacity=0.2, style='wireframe')
+        
+        # Set title based on label distribution
+        positive_count = sum(1 for patch in all_patches if patch['label'] == 1)
+        negative_count = sum(1 for patch in all_patches if patch['label'] == 0)
+        title = f"Edge Patches - Positive (Green): {positive_count}, Negative (Red): {negative_count}"
+        
+        plotter.show(title=title)
+
+    return all_patches
 
-def predict_wireframe(entry, pnet_model) -> Tuple[np.ndarray, List[int]]:
+def calculate_cylinder_overlap_volume(cyl1, cyl2):
+    """
+    Calculate the intersection volume between two cylinders using numpy vectorization.
+    Returns approximate overlap volume.
+    """
+    # Get cylinder parameters
+    p1_start, p1_end = cyl1['start'], cyl1['end']
+    p2_start, p2_end = cyl2['start'], cyl2['end']
+    r1, r2 = cyl1['radius'], cyl2['radius']
+    
+    # Calculate cylinder axes
+    axis1 = p1_end - p1_start
+    axis2 = p2_end - p2_start
+    len1 = np.linalg.norm(axis1)
+    len2 = np.linalg.norm(axis2)
+    
+    if len1 == 0 or len2 == 0:
+        return 0.0
+    
+    axis1_norm = axis1 / len1
+    axis2_norm = axis2 / len2
+    
+    # Calculate distance between cylinder axes using line-line distance formula
+    w = p1_start - p2_start
+    a = np.dot(axis1_norm, axis1_norm)
+    b = np.dot(axis1_norm, axis2_norm)
+    c = np.dot(axis2_norm, axis2_norm)
+    d = np.dot(axis1_norm, w)
+    e = np.dot(axis2_norm, w)
+    
+    denom = a * c - b * b
+    if abs(denom) < 1e-10:  # Lines are parallel
+        # Calculate perpendicular distance between parallel lines
+        cross_product = np.cross(axis1_norm, w)
+        if axis1_norm.shape[0] == 3:  # 3D case
+            dist = np.linalg.norm(cross_product)
+        else:  # 2D case
+            dist = abs(cross_product)
+    else:
+        # Calculate closest points on both lines
+        t1 = (b * e - c * d) / denom
+        t2 = (a * e - b * d) / denom
+        
+        # Clamp to cylinder bounds
+        t1 = np.clip(t1, 0, len1)
+        t2 = np.clip(t2, 0, len2)
+        
+        # Calculate distance between closest points
+        point1 = p1_start + t1 * axis1_norm
+        point2 = p2_start + t2 * axis2_norm
+        dist = np.linalg.norm(point1 - point2)
+    
+    # If cylinders don't intersect radially, return 0
+    if dist >= (r1 + r2):
+        return 0.0
+    
+    # Calculate overlapping length along both axes
+    # Project cylinder 2 endpoints onto cylinder 1 axis
+    proj_start = np.dot(p2_start - p1_start, axis1_norm)
+    proj_end = np.dot(p2_end - p1_start, axis1_norm)
+    
+    # Find overlap interval
+    overlap_start = max(0, min(proj_start, proj_end))
+    overlap_end = min(len1, max(proj_start, proj_end))
+    overlap_length = max(0, overlap_end - overlap_start)
+    
+    if overlap_length <= 0:
+        return 0.0
+    
+    # Approximate volume calculation
+    # For simplicity, assume uniform overlap along the length
+    if dist < abs(r1 - r2):
+        # One cylinder is inside the other
+        smaller_radius = min(r1, r2)
+        overlap_volume = np.pi * smaller_radius**2 * overlap_length
+    else:
+        # Partial overlap - use geometric approximation
+        # This is a simplified calculation for the intersection area of two circles
+        r_smaller = min(r1, r2)
+        r_larger = max(r1, r2)
+        
+        if dist < (r1 + r2):
+            # Calculate intersection area of two circles (approximate)
+            # Using lens area formula
+            d1 = (r1**2 - r2**2 + dist**2) / (2 * dist) if dist > 0 else 0
+            d2 = dist - d1
+            
+            if d1 >= 0 and d1 <= r1 and d2 >= 0 and d2 <= r2:
+                area1 = r1**2 * np.arccos(d1/r1) - d1 * np.sqrt(r1**2 - d1**2)
+                area2 = r2**2 * np.arccos(d2/r2) - d2 * np.sqrt(r2**2 - d2**2)
+                intersection_area = area1 + area2
+            else:
+                intersection_area = np.pi * r_smaller**2
+            
+            overlap_volume = intersection_area * overlap_length
+        else:
+            overlap_volume = 0.0
+    
+    return max(0.0, overlap_volume)
+
+def predict_wireframe(entry, pnet_model, voxel_model) -> Tuple[np.ndarray, List[int]]:
     """
     Predict 3D wireframe from a dataset entry.
     """
     good_entry = convert_entry_to_human_readable(entry)
     colmap_rec = good_entry['colmap_binary']
 
+    if GENERATE_DATASET_EDGES:
+        patches = generate_edge_patches(good_entry)
+        save_patches_dataset_class(patches, EDGES_DATASET_DIR, good_entry['order_id'])
+        return empty_solution()
+
     vert_edge_per_image = {}
     idxs_points = []
+    all_connections = []
     for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(zip(good_entry['gestalt'],
                                                 good_entry['depth'], 
                                                 good_entry['K'],
@@ -660,6 +1148,7 @@ def predict_wireframe(entry, pnet_model) -> Tuple[np.ndarray, List[int]]:
         
         vertices_ours, connections_ours, vertices_3d_ours, patches, filtered_point_idxs = our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id, ade_seg, depth, K=K, R=R, t=t, frame=good_entry)
         idxs_points.append(filtered_point_idxs)
+        all_connections.append(connections_ours)
 
         '''
         if GENERATE_DATASET:
@@ -722,21 +1211,119 @@ def predict_wireframe(entry, pnet_model) -> Tuple[np.ndarray, List[int]]:
 
         vert_edge_per_image[i] = vertices, connections, vertices_3d
     
-    extracted_points, extracted_colors, extracted_ids, whole_pcloud = extract_vertices_from_whole_pcloud(colmap_rec, idxs_points)
+    extracted_points, extracted_colors, extracted_ids, whole_pcloud, connections = extract_vertices_from_whole_pcloud(colmap_rec, idxs_points, all_connections)
 
     patches = generate_patches_v2(extracted_points, extracted_colors, extracted_ids, whole_pcloud, good_entry['wf_vertices'])
 
+    # Predict vertices from patches using the neural network
+    predicted_vertices = []
+    for patch in patches:
+        pred_vertex, pred_dist, pred_class = predict_vertex_from_patch(pnet_model, patch, device='cuda')
+
+        #visu_patch_and_pred(patch, pred_vertex, pred_dist, pred_class)
+
+        if pred_class > 0.5:
+            predicted_vertices.append(pred_vertex)
+        else:
+            predicted_vertices.append(np.array([0.0, 0.0, 0.0]))  # Append a zero vertex if not predicted
+    
+        #pred_vertex_voxel, pred_dist_voxel, pred_class_voxel = predict_vertex_from_patch_voxel(voxel_model, patch, device='cuda')
+        #visu_patch_and_pred(patch, pred_vertex_voxel, pred_dist_voxel, pred_class_voxel)
+
+    predicted_vertices = np.array(predicted_vertices) if predicted_vertices else np.empty((0, 3))
+
+    #visu_pcloud_and_preds(colmap_rec, extracted_ids, extracted_points, extracted_colors, predicted_vertices, connections)
+
     if GENERATE_DATASET:
         save_patches_dataset(patches, DATASET_DIR, img_id)
         return empty_solution()
 
     # Merge vertices from all images
-    all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 0.1)
-    all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
+    #all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 0.1)
+    #all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
     #all_3d_vertices_clean, connections_3d_clean  = prune_not_connected(all_3d_vertices, connections_3d, keep_largest=False)
     #all_3d_vertices_clean, connections_3d_clean  = prune_too_far(all_3d_vertices_clean, connections_3d_clean, colmap_rec, th = 1.5)
 
-    if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1 and False:
+    #if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1 and False:
+    #    print (f'Not enough vertices or connections in the 3D vertices')
+    #    return empty_solution()
+
+    # Filter out zero vertices and update connections accordingly
+    non_zero_mask = ~np.all(np.isclose(predicted_vertices, [0.0, 0.0, 0.0]), axis=1)
+    valid_indices = np.where(non_zero_mask)[0]
+
+    if len(valid_indices) == 0:
+        print("No valid predicted vertices found")
+        return empty_solution()
+
+    # Filter vertices to only include non-zero ones
+    filtered_vertices = predicted_vertices[valid_indices]
+
+    # Create mapping from old indices to new indices
+    old_to_new_mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(valid_indices)}
+
+    # Filter and update connections
+    filtered_connections = []
+    for start_idx, end_idx in connections:
+        if start_idx in old_to_new_mapping and end_idx in old_to_new_mapping:
+            new_start = old_to_new_mapping[start_idx]
+            new_end = old_to_new_mapping[end_idx]
+            if new_start != new_end:  # Ensure we don't connect a vertex to itself
+                filtered_connections.append((new_start, new_end))
+
+    #print(f"Filtered vertices: {len(filtered_vertices)} from {len(predicted_vertices)}")
+    #print(f"Filtered connections: {len(filtered_connections)} from {len(connections)}")
+
+    predicted_vertices = np.array(filtered_vertices)
+    connections = filtered_connections
+
+    return predicted_vertices, connections 
+
+def predict_wireframe_old(entry) -> Tuple[np.ndarray, List[int]]:
+    """
+    Predict 3D wireframe from a dataset entry.
+    """
+    good_entry = convert_entry_to_human_readable(entry)
+    vert_edge_per_image = {}
+    for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(zip(good_entry['gestalt'],
+                                                good_entry['depth'], 
+                                                good_entry['K'],
+                                                good_entry['R'],
+                                                good_entry['t'],
+                                                good_entry['image_ids'],
+                                                good_entry['ade'] # Added ade20k segmentation
+                                                )):
+        colmap_rec = good_entry['colmap_binary']
+        K = np.array(K)
+        R = np.array(R)
+        t = np.array(t)
+        # Resize gestalt segmentation to match depth map size
+        depth_size = (np.array(depth).shape[1], np.array(depth).shape[0]) # W, H
+        gest_seg = gest.resize(depth_size)
+        gest_seg_np = np.array(gest_seg).astype(np.uint8)
+        
+        # Get 2D vertices and edges first
+        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=25.)
+        
+        # Check if we have enough to proceed
+        if (len(vertices) < 2) or (len(connections) < 1):
+            print(f'Not enough vertices or connections found in image {i}, skipping.')
+            vert_edge_per_image[i] = [], [], np.empty((0, 3))
+            continue
+            
+        # Call the refactored function to get 3D points
+        vertices_3d = create_3d_wireframe_single_image(
+            vertices, connections, depth, colmap_rec, img_id, ade_seg, K, R, t
+        )
+        # Store original 2D vertices, connections, and computed 3D points
+        vert_edge_per_image[i] = vertices, connections, vertices_3d
+    
+    # Merge vertices from all images
+    all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 0.5)
+    all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d, keep_largest=False)
+    all_3d_vertices_clean, connections_3d_clean = prune_too_far(all_3d_vertices_clean, connections_3d_clean, colmap_rec, th = 1.5)
+    
+    if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
         print (f'Not enough vertices or connections in the 3D vertices')
         return empty_solution()
 
@@ -1352,6 +1939,27 @@ def our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id_substring, ade_se
     
     H, W = gest_seg_np.shape[:2]
 
+    # Get camera parameters from COLMAP reconstruction if not provided
+    if False:
+        # Find the matching COLMAP image
+        found_img = None
+        for img_id_c, col_img_obj in colmap_rec.images.items():
+            if img_id_substring in col_img_obj.name:
+                found_img = col_img_obj
+                break
+        
+        if found_img is not None:
+            # Get camera intrinsic matrix
+            K = found_img.camera.calibration_matrix()
+            
+            # Get world-to-camera transformation matrix
+            world_to_cam = found_img.cam_from_world.matrix()
+            R = world_to_cam[:3, :3]
+            t = world_to_cam[:3, 3]
+        else:
+            print(f"Image substring {img_id_substring} not found in COLMAP.")
+            return [], [], [], [], []
+
     points_cam, points_xyz_world, points_idxs = get_visible_points(colmap_rec, img_id_substring, R=R, t=t)
 
     uv, valid_indices = project_points_to_2d(points_cam, K, H, W)