Implements FastPointNet for vertex prediction

Adds a FastPointNet model for predicting 3D vertex coordinates from point cloud patches.
Includes a dataset class for loading and augmenting patch data.
Also adds training and prediction functions for the model.

Patches are generated and saved for training the PointNet model,
allowing for iterative refinement of vertex predictions.
The patch generation process includes filtering COLMAP points within a
ball around identified vertices and creating a 7D point cloud
representation. GT vertex assignment and data augmentation are incorporated
to improve the training data.

fast_pointnet.py

@@ -0,0 +1,421 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import pickle
+from torch.utils.data import Dataset, DataLoader
+from typing import List, Dict, Tuple, Optional
+import json
+
+class FastPointNet(nn.Module):
+    """
+    Fast PointNet implementation for 3D vertex prediction from point cloud patches.
+    Takes 7D point clouds (x,y,z,r,g,b,filtered_flag) and predicts 3D vertex coordinates.
+    """
+    
+    def __init__(self, input_dim=7, output_dim=3, max_points=1024, predict_score=True):
+        super(FastPointNet, self).__init__()
+        self.max_points = max_points
+        self.predict_score = predict_score
+        
+        # Point-wise MLPs
+        self.conv1 = nn.Conv1d(input_dim, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, 256, 1)
+        
+        # Global feature extraction
+        self.conv4 = nn.Conv1d(256, 512, 1)
+        self.conv5 = nn.Conv1d(512, 1024, 1)
+        
+        # Shared features
+        self.shared_fc = nn.Linear(1024, 512)
+        
+        # Position prediction head
+        self.pos_fc1 = nn.Linear(512, 256)
+        self.pos_fc2 = nn.Linear(256, output_dim)
+        
+        # Score prediction head (predicts distance to GT)
+        if self.predict_score:
+            self.score_fc1 = nn.Linear(512, 256)
+            self.score_fc2 = nn.Linear(256, 128)
+            self.score_fc3 = nn.Linear(128, 1)  # Single score output
+        
+        self.dropout = nn.Dropout(0.3)
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(256)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(1024)
+        
+    def forward(self, x):
+        """
+        Forward pass
+        Args:
+            x: (batch_size, input_dim, max_points) tensor
+        Returns:
+            if predict_score=True: tuple of (position, score)
+                position: (batch_size, output_dim) tensor of predicted 3D coordinates
+                score: (batch_size, 1) tensor of predicted distance to GT
+            else: (batch_size, output_dim) tensor of predicted 3D coordinates
+        """
+        batch_size = x.size(0)
+        
+        # Point-wise feature extraction
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = F.relu(self.bn4(self.conv4(x)))
+        x = F.relu(self.bn5(self.conv5(x)))
+        
+        # Global max pooling
+        x = torch.max(x, 2)[0]  # (batch_size, 1024)
+        
+        # Shared features
+        shared_features = F.relu(self.shared_fc(x))
+        shared_features = self.dropout(shared_features)
+        
+        # Position prediction
+        pos_features = F.relu(self.pos_fc1(shared_features))
+        pos_features = self.dropout(pos_features)
+        position = self.pos_fc2(pos_features)
+        
+        if self.predict_score:
+            # Score prediction (distance to GT)
+            score_features = F.relu(self.score_fc1(shared_features))
+            score_features = self.dropout(score_features)
+            score_features = F.relu(self.score_fc2(score_features))
+            score_features = self.dropout(score_features)
+            score = F.relu(self.score_fc3(score_features))  # Ensure positive distance
+            
+            return position, score
+        else:
+            return position
+
+class PatchDataset(Dataset):
+    """
+    Dataset class for loading saved patches for PointNet training.
+    """
+    
+    def __init__(self, dataset_dir: str, max_points: int = 1024, augment: bool = True):
+        self.dataset_dir = dataset_dir
+        self.max_points = max_points
+        self.augment = augment
+        
+        # Load patch files
+        self.patch_files = []
+        for file in os.listdir(dataset_dir):
+            if file.endswith('.pkl'):
+                self.patch_files.append(os.path.join(dataset_dir, file))
+                
+        print(f"Found {len(self.patch_files)} patch files in {dataset_dir}")
+    
+    def __len__(self):
+        return len(self.patch_files)
+    
+    def __getitem__(self, idx):
+        """
+        Load and process a patch for training.
+        Returns:
+            patch_data: (7, max_points) tensor of point cloud data
+            target: (3,) tensor of target 3D coordinates
+            valid_mask: (max_points,) boolean tensor indicating valid points
+            distance_to_gt: scalar tensor of distance from initial prediction to GT
+        """
+        patch_file = self.patch_files[idx]
+        
+        with open(patch_file, 'rb') as f:
+            patch_info = pickle.load(f)
+        
+        patch_7d = patch_info['patch_7d']  # (N, 7)
+        target = patch_info['assigned_gt_vertex']  # (3,) or None
+        initial_pred = patch_info.get('initial_pred', None)  # (3,) or None
+        
+        # Skip patches without ground truth
+        if target is None:
+            # Return dummy data that will be filtered out
+            dummy_patch = np.zeros((self.max_points, 7))
+            dummy_target = np.zeros(3)
+            valid_mask = np.zeros(self.max_points, dtype=bool)
+            dummy_distance = np.array([0.0])
+            return (torch.from_numpy(dummy_patch.T).float(), 
+                   torch.from_numpy(dummy_target).float(), 
+                   torch.from_numpy(valid_mask),
+                   torch.from_numpy(dummy_distance).float())
+        
+        target = np.array(target)
+        
+        # Normalize colors from [0,1] to [-1,1]
+        patch_7d[:, 3:6] = patch_7d[:, 3:6] * 2.0 - 1.0
+        
+        # Pad or sample points to max_points
+        num_points = patch_7d.shape[0]
+        
+        if num_points >= self.max_points:
+            # Randomly sample max_points
+            indices = np.random.choice(num_points, self.max_points, replace=False)
+            patch_sampled = patch_7d[indices]
+            valid_mask = np.ones(self.max_points, dtype=bool)
+        else:
+            # Pad with zeros
+            patch_sampled = np.zeros((self.max_points, 7))
+            patch_sampled[:num_points] = patch_7d
+            valid_mask = np.zeros(self.max_points, dtype=bool)
+            valid_mask[:num_points] = True
+        
+        # Data augmentation
+        if self.augment:
+            patch_sampled = self._augment_patch(patch_sampled, valid_mask)
+            target = self._augment_target(target)
+        
+        # Convert to tensors and transpose for conv1d (channels first)
+        patch_tensor = torch.from_numpy(patch_sampled.T).float()  # (7, max_points)
+        target_tensor = torch.from_numpy(target).float()  # (3,)
+        valid_mask_tensor = torch.from_numpy(valid_mask)
+        initial_pred = torch.from_numpy(initial_pred).float()
+        
+        return patch_tensor, target_tensor, valid_mask_tensor, initial_pred
+    
+    def _augment_patch(self, patch: np.ndarray, valid_mask: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Apply data augmentation to patch and target with consistent transformations."""
+        # Only augment valid points
+        valid_points = patch[valid_mask]
+        
+        if len(valid_points) == 0:
+            return patch, target
+        
+        # Random rotation around Z-axis
+        angle = np.random.uniform(0, 2 * np.pi)
+        cos_angle, sin_angle = np.cos(angle), np.sin(angle)
+        rotation_matrix = np.array([
+            [cos_angle, -sin_angle, 0],
+            [sin_angle, cos_angle, 0],
+            [0, 0, 1]
+        ])
+        
+        # Apply rotation to patch coordinates
+        valid_points[:, :3] = valid_points[:, :3] @ rotation_matrix.T
+        
+        # Apply same rotation to target
+        target_augmented = target @ rotation_matrix.T
+        
+        # Add small random noise to coordinates
+        noise = np.random.normal(0, 0.01, valid_points[:, :3].shape)
+        valid_points[:, :3] += noise
+        
+        # Color jittering
+        color_noise = np.random.normal(0, 0.02, valid_points[:, 3:6].shape)
+        valid_points[:, 3:6] = np.clip(valid_points[:, 3:6] + color_noise, 0, 1)
+        
+        patch[valid_mask] = valid_points
+        return patch, target_augmented
+
+def save_patches_dataset(patches: List[Dict], dataset_dir: str, entry_id: str):
+    """
+    Save patches from prediction pipeline to create a training dataset.
+    
+    Args:
+        patches: List of patch dictionaries from generate_patches()
+        dataset_dir: Directory to save the dataset
+        entry_id: Unique identifier for this entry/image
+    """
+    os.makedirs(dataset_dir, exist_ok=True)
+    
+    for i, patch in enumerate(patches):
+        # Create unique filename
+        filename = f"{entry_id}_patch_{i}.pkl"
+        filepath = os.path.join(dataset_dir, filename)
+        
+        # Save patch data
+        with open(filepath, 'wb') as f:
+            pickle.dump(patch, f)
+    
+    print(f"Saved {len(patches)} patches for entry {entry_id}")
+
+def train_pointnet(dataset_dir: str, model_save_path: str, epochs: int = 100, batch_size: int = 32, lr: float = 0.001, 
+                  score_weight: float = 0.1):
+    """
+    Train the FastPointNet model on saved patches.
+    
+    Args:
+        dataset_dir: Directory containing saved patch files
+        model_save_path: Path to save the trained model
+        epochs: Number of training epochs
+        batch_size: Training batch size
+        lr: Learning rate
+        score_weight: Weight for the distance prediction loss
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Training on device: {device}")
+    
+    # Create dataset and dataloader
+    dataset = PatchDataset(dataset_dir, max_points=1024, augment=True)
+    
+    # Filter out invalid samples
+    valid_indices = []
+    for i in range(len(dataset)):
+        _, target, valid_mask, _ = dataset[i]
+        if valid_mask.sum() > 0 and not torch.all(target == 0):
+            valid_indices.append(i)
+    
+    print(f"Found {len(valid_indices)} valid patches out of {len(dataset)}")
+    
+    # Create subset with valid samples
+    valid_dataset = torch.utils.data.Subset(dataset, valid_indices)
+    dataloader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
+    
+    # Initialize model with score prediction
+    model = FastPointNet(input_dim=7, output_dim=3, max_points=1024, predict_score=True)
+    model.to(device)
+    
+    # Loss functions
+    position_criterion = nn.MSELoss()
+    score_criterion = nn.MSELoss()
+    
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5)
+    
+    # Training loop
+    model.train()
+    for epoch in range(epochs):
+        total_loss = 0.0
+        total_pos_loss = 0.0
+        total_score_loss = 0.0
+        num_batches = 0
+        
+        for batch_idx, (patch_data, targets, valid_masks, distances) in enumerate(dataloader):
+            patch_data = patch_data.to(device)  # (batch_size, 7, max_points)
+            targets = targets.to(device)  # (batch_size, 3)
+            distances = distances.to(device)  # (batch_size, 1)
+            
+            # Forward pass
+            optimizer.zero_grad()
+            predictions, predicted_scores = model(patch_data)
+            
+            # Compute actual distance from predictions to targets
+            actual_distances = torch.norm(predictions - targets, dim=1, keepdim=True)
+            
+            # Compute losses
+            pos_loss = position_criterion(predictions, targets)
+            score_loss = score_criterion(predicted_scores, actual_distances)
+            
+            # Combined loss
+            total_batch_loss = pos_loss + score_weight * score_loss
+            
+            # Backward pass
+            total_batch_loss.backward()
+            optimizer.step()
+            
+            total_loss += total_batch_loss.item()
+            total_pos_loss += pos_loss.item()
+            total_score_loss += score_loss.item()
+            num_batches += 1
+            
+            if batch_idx % 50 == 0:
+                print(f"Epoch {epoch+1}/{epochs}, Batch {batch_idx}, "
+                      f"Total Loss: {total_batch_loss.item():.6f}, "
+                      f"Pos Loss: {pos_loss.item():.6f}, "
+                      f"Score Loss: {score_loss.item():.6f}")
+        
+        avg_loss = total_loss / num_batches if num_batches > 0 else 0
+        avg_pos_loss = total_pos_loss / num_batches if num_batches > 0 else 0
+        avg_score_loss = total_score_loss / num_batches if num_batches > 0 else 0
+        
+        print(f"Epoch {epoch+1}/{epochs} completed, "
+              f"Avg Total Loss: {avg_loss:.6f}, "
+              f"Avg Pos Loss: {avg_pos_loss:.6f}, "
+              f"Avg Score Loss: {avg_score_loss:.6f}")
+        
+        scheduler.step()
+
+        # Save model checkpoint every epoch
+        checkpoint_path = model_save_path.replace('.pth', f'_epoch_{epoch+1}.pth')
+        torch.save({
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'epoch': epoch + 1,
+            'loss': avg_loss,
+        }, checkpoint_path)
+    
+    # Save the trained model
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'epoch': epochs,
+    }, model_save_path)
+    
+    print(f"Model saved to {model_save_path}")
+    return model
+
+def load_pointnet_model(model_path: str, device: torch.device = None, predict_score: bool = True) -> FastPointNet:
+    """
+    Load a trained FastPointNet model.
+    
+    Args:
+        model_path: Path to the saved model
+        device: Device to load the model on
+        predict_score: Whether the model predicts scores
+        
+    Returns:
+        Loaded FastPointNet model
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    model = FastPointNet(input_dim=7, output_dim=3, max_points=1024, predict_score=predict_score)
+    
+    checkpoint = torch.load(model_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    
+    model.to(device)
+    model.eval()
+    
+    return model
+
+def predict_vertex_from_patch(model: FastPointNet, patch_7d: np.ndarray, device: torch.device = None) -> Tuple[np.ndarray, float]:
+    """
+    Predict 3D vertex coordinates and confidence score from a patch using trained PointNet.
+    
+    Args:
+        model: Trained FastPointNet model
+        patch_7d: (N, 7) numpy array of point cloud data
+        device: Device to run prediction on
+        
+    Returns:
+        tuple of (predicted_coordinates, confidence_score)
+            predicted_coordinates: (3,) numpy array of predicted 3D coordinates
+            confidence_score: float representing predicted distance to GT (lower is better)
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    model.eval()
+    
+    # Prepare input
+    max_points = 1024
+    num_points = patch_7d.shape[0]
+    
+    if num_points >= max_points:
+        # Sample points
+        indices = np.random.choice(num_points, max_points, replace=False)
+        patch_sampled = patch_7d[indices]
+    else:
+        # Pad with zeros
+        patch_sampled = np.zeros((max_points, 7))
+        patch_sampled[:num_points] = patch_7d
+    
+    # Convert to tensor
+    patch_tensor = torch.from_numpy(patch_sampled.T).float().unsqueeze(0)  # (1, 7, max_points)
+    patch_tensor = patch_tensor.to(device)
+    
+    # Predict
+    with torch.no_grad():
+        if model.predict_score:
+            position, score = model(patch_tensor)
+            position = position.cpu().numpy().squeeze()
+            score = score.cpu().numpy().squeeze()
+            return position, score
+        else:
+            position = model(patch_tensor)
+            position = position.cpu().numpy().squeeze()
+            return position, None

predict.py

@@ -10,6 +10,11 @@ from PIL import Image as PImage
 import cv2
 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
+
+GENERATE_DATASET = True
+DATASET_DIR = '/home/skvrnjan/personal/hohocustom/'
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -389,11 +394,6 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
     good_entry = convert_entry_to_human_readable(entry)
     colmap_rec = good_entry['colmap_binary']
 
-    colmap_pcloud = []
-    for i, p3D in colmap_rec.points3D.items():
-        p3D.color = np.array([0, 0, 0])
-        colmap_pcloud.append(p3D)
-
     vert_edge_per_image = {}
     for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(zip(good_entry['gestalt'],
                                                 good_entry['depth'], 
@@ -413,7 +413,13 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
         gest_seg = gest.resize(depth_size)
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
         
-        vertices_ours, connections_ours, vertices_3d_ours = our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id, ade_seg, depth, K=K, R=R, t=t)
+        vertices_ours, connections_ours, vertices_3d_ours, patches = our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id, ade_seg, depth, K=K, R=R, t=t, frame=good_entry)
+        
+        if GENERATE_DATASET:
+            save_patches_dataset(patches, DATASET_DIR, img_id)
+
+            continue
+
         vertices, connections, vertices_3d = vertices_ours, connections_ours, vertices_3d_ours
         # Get 2D vertices and edges first
         #vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=25.)
@@ -444,7 +450,7 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
             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)
+        vertices_3d = create_3d_wireframe_single_image(vertices, connections, depth, colmap_rec, img_id, ade_seg, K, R, t)
         #vertices_3d = gt_verts3d
         # Store original 2D vertices, connections, and computed 3D points
 
@@ -480,6 +486,9 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
         # Visualize the point cloud
         o3d.visualization.draw_geometries([pcd], window_name="COLMAP Point Cloud")
     '''
+    if GENERATE_DATASET:
+        return empty_solution()
+
     # 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)
@@ -627,9 +636,9 @@ def get_apex_or_eave_points(apex, uv, gest_seg_np, house_mask, valid_indices, po
                     final_valid_indices = valid_point_indices[depth_filter]
                     
                     # Add corresponding points to filtered lists
-                    filtered_points_xyz.extend(points_xyz_world[final_valid_indices])
-                    filtered_point_idxs.extend(points_idxs[final_valid_indices])
-                    filtered_points_color.extend([color] * np.sum(depth_filter))
+                    filtered_points_xyz.append(points_xyz_world[final_valid_indices])
+                    filtered_point_idxs.append(points_idxs[final_valid_indices])
+                    filtered_points_color.append([color] * np.sum(depth_filter))
 
                     # Find the point with lowest depth in the filtered points
                     if len(final_valid_indices) > 0:
@@ -637,9 +646,6 @@ def get_apex_or_eave_points(apex, uv, gest_seg_np, house_mask, valid_indices, po
                         lowest_depth_point = final_valid_indices[lowest_depth_idx]
 
                         filtered_vertices_apex.append(points_xyz_world[lowest_depth_point])
-                        filtered_points_xyz.append(points_xyz_world[lowest_depth_point])
-                        filtered_point_idxs.append(points_idxs[lowest_depth_point])
-                        filtered_points_color.append(np.array([1., 1., 0.]))
                         filtered_vertices_apex_uv.append(centroids[i])
 
     return filtered_points_xyz, filtered_point_idxs, filtered_points_color, filtered_vertices_apex, filtered_vertices_apex_uv
@@ -653,9 +659,9 @@ def get_vertexes(uv, gest_seg_np, house_mask, valid_indices, points_xyz_world, p
     filtered_point_idxs = filtered_point_idxs_apex + filtered_point_idxs_eave
     filtered_points_color = filtered_points_color_apex + filtered_points_color_eave
 
-    filtered_points_xyz = np.array(filtered_points_xyz[::-1]) if filtered_points_xyz else np.empty((0, 3))
-    filtered_point_idxs = np.array(filtered_point_idxs[::-1]) if filtered_point_idxs else np.empty((0,))
-    filtered_points_color = np.array(filtered_points_color[::-1]) if filtered_points_color else np.empty((0, 3))
+    #filtered_points_xyz = np.array(filtered_points_xyz[::-1]) if filtered_points_xyz else np.empty((0, 3))
+    #filtered_point_idxs = np.array(filtered_point_idxs[::-1]) if filtered_point_idxs else np.empty((0,))
+    #filtered_points_color = np.array(filtered_points_color[::-1]) if filtered_points_color else np.empty((0, 3))
     filtered_vertices_apex = np.array(filtered_vertices_apex) if filtered_vertices_apex else np.empty((0, 3))
     filtered_vertices_apex_uv = np.array(filtered_vertices_apex_uv) if filtered_vertices_apex_uv else np.empty((0, 2))
     filtered_vertices_eave = np.array(filtered_vertices_eave) if filtered_vertices_eave else np.empty((0, 3))
@@ -803,7 +809,156 @@ def visualize_3d_wireframe(colmap_rec, filtered_points_xyz, filtered_points_colo
 
     #o3d.visualization.draw_geometries(geometries, window_name=f"Combined Point Cloud - {img_id_substring}")
 
-def our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id_substring, ade_seg, depth, K=None, R=None, t=None, ):
+def generate_patches(colmap_rec, filtered_points_idxs, frame, filtered_vertices):
+    patches = []
+
+    gt_vertices = frame['wf_vertices']
+    
+    # Process each group of filtered points
+    for group_idx, point_idxs in enumerate(filtered_points_idxs):
+        if len(point_idxs) == 0:
+            continue
+            
+        # Get 3D coordinates and colors for this group
+        group_points_3d = []
+        group_colors = []
+        assigned_gt_vertex = None
+        
+        for pid in point_idxs:
+            p3d = colmap_rec.points3D[pid]
+            group_points_3d.append(p3d.xyz)
+            group_colors.append(p3d.color)
+        
+        if len(group_points_3d) == 0:
+            continue
+            
+        group_points_3d = np.array(group_points_3d)
+        group_colors = np.array(group_colors)
+        
+        # Calculate centroid of filtered points
+        # Find the closest GT vertex to the centroid of filtered points
+        centroid = np.mean(group_points_3d, axis=0)
+        
+        if len(gt_vertices) > 0:
+            # Calculate distances from centroid to all GT vertices
+            distances_to_gt = []
+            for gt_vertex in gt_vertices:
+                distance = np.linalg.norm(gt_vertex - centroid)
+                distances_to_gt.append(distance)
+            
+            # Find the closest GT vertex
+            min_distance_idx = np.argmin(distances_to_gt)
+            closest_gt_vertex = gt_vertices[min_distance_idx]
+            min_distance = distances_to_gt[min_distance_idx]
+            
+            # Define ball radius (you can adjust this value)
+            ball_radius = 2.0  # meters
+            
+            # Use closest GT vertex as centroid if it's within the ball radius
+            if min_distance <= ball_radius:
+                assigned_gt_vertex = closest_gt_vertex
+            # If no GT vertex is close enough, skip this group
+            else:
+                continue
+        else:
+            # No GT vertices available, use original centroid
+            centroid = np.mean(group_points_3d, axis=0)
+        
+        # Define ball radius (you can adjust this value)
+        ball_radius = 2.0  # meters
+        
+        # Find all COLMAP points within the ball around centroid
+        patch_points_3d = []
+        patch_colors = []
+        patch_point_ids = []
+        
+        for pid, p3d in colmap_rec.points3D.items():
+            distance = np.linalg.norm(p3d.xyz - centroid)
+            if distance <= ball_radius:
+                patch_points_3d.append(p3d.xyz)
+                patch_colors.append(p3d.color)
+                patch_point_ids.append(pid)
+        
+        if len(patch_points_3d) == 0:
+            continue
+            
+        patch_points_3d = np.array(patch_points_3d)
+        
+        # Calculate offset to center the patch
+        patch_centroid = np.mean(patch_points_3d, axis=0)
+        offset = -patch_centroid
+        
+        # Shift points to center them around origin
+        patch_points_3d += offset
+        
+        # Also shift the assigned GT vertex by the same offset if it exists
+        if assigned_gt_vertex is not None:
+            assigned_gt_vertex = assigned_gt_vertex + offset
+        patch_colors = np.array(patch_colors)
+        
+        # Create 7D point cloud for this patch
+        # [x, y, z, r, g, b, in_filtered_flag]
+        patch_7d = np.zeros((len(patch_points_3d), 7))
+        patch_7d[:, :3] = patch_points_3d  # xyz coordinates
+        patch_7d[:, 3:6] = patch_colors / 255.0  # rgb colors normalized to [0,1]
+        
+        # Set in_filtered_flag: 1 if point was in original filtered set, 0 otherwise
+        for i, pid in enumerate(patch_point_ids):
+            if pid in point_idxs:
+                patch_7d[i, 6] = 1.0
+            else:
+                patch_7d[i, 6] = 0.0
+        
+        if filtered_vertices[group_idx] is not None:
+            initial_pred = filtered_vertices[group_idx] + offset
+        else:
+            initial_pred = None
+
+        patches.append({
+            'patch_7d': patch_7d,
+            'centroid': centroid,
+            'radius': ball_radius,
+            'point_ids': patch_point_ids,
+            'filtered_point_ids': point_idxs,
+            'group_idx': group_idx,
+            'assigned_gt_vertex': assigned_gt_vertex,
+            'offset': offset,
+            'initial_pred': initial_pred
+        })
+
+        if False:
+            # Create plotter
+            plotter = pv.Plotter()
+            
+            # Create point cloud for this patch
+            patch_cloud = pv.PolyData(patch_points_3d)
+            
+            # Color points: red for filtered points, blue for other points
+            patch_point_colors = []
+            for i, pid in enumerate(patch_point_ids):
+                if pid in point_idxs:
+                    patch_point_colors.append([255, 0, 0])  # Red for filtered points
+                else:
+                    patch_point_colors.append([0, 0, 255])  # Blue for other points
+            
+            patch_cloud["colors"] = np.array(patch_point_colors)
+            plotter.add_mesh(patch_cloud, scalars="colors", rgb=True, point_size=8, render_points_as_spheres=True)
+            
+            # Create sphere to visualize GT vertex if available
+            if assigned_gt_vertex is not None:
+                gt_sphere = pv.Sphere(radius=0.1, center=assigned_gt_vertex)
+                plotter.add_mesh(gt_sphere, color="green", opacity=0.5)
+            
+            if initial_pred is not None:
+                # Create sphere to visualize initial prediction
+                pred_sphere = pv.Sphere(radius=0.1, center=initial_pred)
+                plotter.add_mesh(pred_sphere, color="orange", opacity=0.5)
+            
+            plotter.show(title=f"Patch {group_idx}")
+    
+    return patches
+
+def our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id_substring, ade_seg, depth, K=None, R=None, t=None, frame=None):
     """
     Identify apex and eave-end vertices, then detect lines for eave/ridge/rake/valley.
     Also find all COLMAP points that project into apex or eave_end masks.
@@ -838,6 +993,8 @@ def our_get_vertices_and_edges(gest_seg_np, colmap_rec, img_id_substring, ade_se
 
     vertices_formatted, connections, all_vertices_3d = get_connections(gest_seg_np, filtered_vertices_apex, filtered_vertices_eave, filtered_vertices_apex_uv, filtered_vertices_eave_uv)
 
+    patches = generate_patches(colmap_rec, filtered_point_idxs, frame, all_vertices_3d)
+
     #visualize_3d_wireframe(colmap_rec, filtered_points_xyz, filtered_points_color, all_vertices_3d, connections) 
 
-    return vertices_formatted, connections, all_vertices_3d
+    return vertices_formatted, connections, all_vertices_3d, patches

train.py

@@ -22,7 +22,7 @@ scores_iou = []
 show_visu = False
 
 idx = 0
-for a in ds['validation']:
+for a in ds['train']:
     #plot_all_modalities(a)
     #pred_vertices, pred_edges = predict_wireframe(a)
     try: