Adds PointNet classification for edge detection

Implements a PointNet-based architecture for classifying 6D point cloud patches, designed to identify edges within 3D data. This includes a dataset class for loading and augmenting patches, a training pipeline, and functions for prediction. It also updates the 3D CNN prediction function and incorporates point merging based on overlap to refine vertex extraction.

fast_pointnet_class.py

@@ -0,0 +1,407 @@
+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 ClassificationPointNet(nn.Module):
+    """
+    PointNet implementation for binary classification from 6D point cloud patches.
+    Takes 6D point clouds (x,y,z,r,g,b) and predicts binary classification (edge/not edge).
+    """
+    def __init__(self, input_dim=6, max_points=1024):
+        super(ClassificationPointNet, self).__init__()
+        self.max_points = max_points
+        
+        # Point-wise MLPs for feature extraction (deeper network)
+        self.conv1 = nn.Conv1d(input_dim, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, 256, 1)
+        self.conv4 = nn.Conv1d(256, 512, 1)
+        self.conv5 = nn.Conv1d(512, 1024, 1)
+        self.conv6 = nn.Conv1d(1024, 2048, 1)  # Additional layer
+        
+        # Classification head (deeper with more capacity)
+        self.fc1 = nn.Linear(2048, 1024)
+        self.fc2 = nn.Linear(1024, 512)
+        self.fc3 = nn.Linear(512, 256)
+        self.fc4 = nn.Linear(256, 128)
+        self.fc5 = nn.Linear(128, 64)
+        self.fc6 = nn.Linear(64, 1)  # Single output for binary classification
+        
+        # Batch normalization layers
+        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)
+        self.bn6 = nn.BatchNorm1d(2048)
+        
+        # Dropout layers
+        self.dropout1 = nn.Dropout(0.3)
+        self.dropout2 = nn.Dropout(0.4)
+        self.dropout3 = nn.Dropout(0.5)
+        self.dropout4 = nn.Dropout(0.4)
+        self.dropout5 = nn.Dropout(0.3)
+
+    def forward(self, x):
+        """
+        Forward pass
+        Args:
+            x: (batch_size, input_dim, max_points) tensor
+        Returns:
+            classification: (batch_size, 1) tensor of logits (sigmoid for probability)
+        """
+        batch_size = x.size(0)
+        
+        # Point-wise feature extraction
+        x1 = F.relu(self.bn1(self.conv1(x)))
+        x2 = F.relu(self.bn2(self.conv2(x1)))
+        x3 = F.relu(self.bn3(self.conv3(x2)))
+        x4 = F.relu(self.bn4(self.conv4(x3)))
+        x5 = F.relu(self.bn5(self.conv5(x4)))
+        x6 = F.relu(self.bn6(self.conv6(x5)))
+        
+        # Global max pooling
+        global_features = torch.max(x6, 2)[0]  # (batch_size, 2048)
+        
+        # Classification head
+        x = F.relu(self.fc1(global_features))
+        x = self.dropout1(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout2(x)
+        x = F.relu(self.fc3(x))
+        x = self.dropout3(x)
+        x = F.relu(self.fc4(x))
+        x = self.dropout4(x)
+        x = F.relu(self.fc5(x))
+        x = self.dropout5(x)
+        classification = self.fc6(x)  # (batch_size, 1)
+        
+        return classification
+
+class PatchClassificationDataset(Dataset):
+    """
+    Dataset class for loading saved patches for PointNet classification 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: (6, max_points) tensor of point cloud data
+            label: scalar tensor for binary classification (0 or 1)
+            valid_mask: (max_points,) boolean tensor indicating valid points
+        """
+        patch_file = self.patch_files[idx]
+        
+        with open(patch_file, 'rb') as f:
+            patch_info = pickle.load(f)
+        
+        patch_6d = patch_info['patch_6d']  # (N, 6)
+        label = patch_info.get('label', 0)  # Get binary classification label (0 or 1)
+        
+        # Pad or sample points to max_points
+        num_points = patch_6d.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_6d[indices]
+            valid_mask = np.ones(self.max_points, dtype=bool)
+        else:
+            # Pad with zeros
+            patch_sampled = np.zeros((self.max_points, 6))
+            patch_sampled[:num_points] = patch_6d
+            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)
+        
+        # Convert to tensors and transpose for conv1d (channels first)
+        patch_tensor = torch.from_numpy(patch_sampled.T).float()  # (6, max_points)
+        label_tensor = torch.tensor(label, dtype=torch.float32)  # Float for BCE loss
+        valid_mask_tensor = torch.from_numpy(valid_mask)
+        
+        return patch_tensor, label_tensor, valid_mask_tensor
+
+    def _augment_patch(self, patch, valid_mask):
+        """
+        Apply data augmentation to the patch.
+        """
+        valid_points = patch[valid_mask]
+        
+        if len(valid_points) == 0:
+            return patch
+        
+        # Random rotation around z-axis
+        angle = np.random.uniform(0, 2 * np.pi)
+        cos_angle = np.cos(angle)
+        sin_angle = np.sin(angle)
+        rotation_matrix = np.array([
+            [cos_angle, -sin_angle, 0],
+            [sin_angle, cos_angle, 0],
+            [0, 0, 1]
+        ])
+        
+        # Apply rotation to xyz coordinates
+        valid_points[:, :3] = valid_points[:, :3] @ rotation_matrix.T
+        
+        # Random jittering
+        noise = np.random.normal(0, 0.01, valid_points[:, :3].shape)
+        valid_points[:, :3] += noise
+        
+        # Random scaling
+        scale = np.random.uniform(0.9, 1.1)
+        valid_points[:, :3] *= scale
+        
+        patch[valid_mask] = valid_points
+        return patch
+
+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)
+        
+        # Skip if file already exists
+        if os.path.exists(filepath):
+            continue
+        
+        # Save patch data
+        with open(filepath, 'wb') as f:
+            pickle.dump(patch, f)
+    
+    print(f"Saved {len(patches)} patches for entry {entry_id}")
+
+# Create dataloader with custom collate function to filter invalid samples
+def collate_fn(batch):
+    valid_batch = []
+    for patch_data, label, valid_mask in batch:
+        # Filter out invalid samples (no valid points)
+        if valid_mask.sum() > 0:
+            valid_batch.append((patch_data, label, valid_mask))
+    
+    if len(valid_batch) == 0:
+        return None
+    
+    # Stack valid samples
+    patch_data = torch.stack([item[0] for item in valid_batch])
+    labels = torch.stack([item[1] for item in valid_batch])
+    valid_masks = torch.stack([item[2] for item in valid_batch])
+    
+    return patch_data, labels, valid_masks
+
+# Initialize weights using Xavier/Glorot initialization
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.Linear):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.BatchNorm1d):
+        nn.init.ones_(m.weight)
+        nn.init.zeros_(m.bias)
+
+def train_pointnet(dataset_dir: str, model_save_path: str, epochs: int = 100, batch_size: int = 32, 
+                  lr: float = 0.001):
+    """
+    Train the ClassificationPointNet 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
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Training on device: {device}")
+    
+    # Create dataset and dataloader
+    dataset = PatchClassificationDataset(dataset_dir, max_points=1024, augment=True)
+    print(f"Dataset loaded with {len(dataset)} samples")
+    
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8, 
+                           collate_fn=collate_fn, drop_last=True)
+    
+    # Initialize model
+    model = ClassificationPointNet(input_dim=6, max_points=1024)
+    model.apply(init_weights)
+    model.to(device)
+    
+    # Loss function and optimizer (BCE for binary classification)
+    criterion = nn.BCEWithLogitsLoss()
+    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
+        correct = 0
+        total = 0
+        num_batches = 0
+        
+        for batch_idx, batch_data in enumerate(dataloader):
+            if batch_data is None:  # Skip invalid batches
+                continue
+                
+            patch_data, labels, valid_masks = batch_data
+            patch_data = patch_data.to(device)  # (batch_size, 6, max_points)
+            labels = labels.to(device).unsqueeze(1)  # (batch_size, 1)
+            
+            # Forward pass
+            optimizer.zero_grad()
+            outputs = model(patch_data)  # (batch_size, 1)
+            loss = criterion(outputs, labels)
+            
+            # Backward pass
+            loss.backward()
+            optimizer.step()
+            
+            # Statistics
+            total_loss += loss.item()
+            predicted = (torch.sigmoid(outputs) > 0.5).float()
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+            num_batches += 1
+            
+            if batch_idx % 50 == 0:
+                print(f"Epoch {epoch+1}/{epochs}, Batch {batch_idx}, "
+                      f"Loss: {loss.item():.6f}, "
+                      f"Accuracy: {100 * correct / total:.2f}%")
+        
+        avg_loss = total_loss / num_batches if num_batches > 0 else 0
+        accuracy = 100 * correct / total if total > 0 else 0
+        
+        print(f"Epoch {epoch+1}/{epochs} completed, "
+              f"Avg Loss: {avg_loss:.6f}, "
+              f"Accuracy: {accuracy:.2f}%")
+        
+        scheduler.step()
+
+        # Save model checkpoint every 10 epochs
+        if (epoch + 1) % 10 == 0:
+            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,
+                'accuracy': accuracy,
+            }, 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) -> ClassificationPointNet:
+    """
+    Load a trained ClassificationPointNet model.
+    
+    Args:
+        model_path: Path to the saved model
+        device: Device to load the model on
+        
+    Returns:
+        Loaded ClassificationPointNet model
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    model = ClassificationPointNet(input_dim=6, max_points=1024)
+    
+    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_class_from_patch(model: ClassificationPointNet, patch: Dict, device: torch.device = None) -> Tuple[int, float]:
+    """
+    Predict binary classification from a patch using trained PointNet.
+    
+    Args:
+        model: Trained ClassificationPointNet model
+        patch: Dictionary containing patch data with 'patch_6d' key
+        device: Device to run prediction on
+        
+    Returns:
+        tuple of (predicted_class, confidence)
+            predicted_class: int (0 for not edge, 1 for edge)
+            confidence: float representing confidence score (0-1)
+    """
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    patch_6d = patch['patch_6d']  # (N, 6)
+    
+    # Prepare input
+    max_points = 1024
+    num_points = patch_6d.shape[0]
+    
+    if num_points >= max_points:
+        # Sample points
+        indices = np.random.choice(num_points, max_points, replace=False)
+        patch_sampled = patch_6d[indices]
+    else:
+        # Pad with zeros
+        patch_sampled = np.zeros((max_points, 6))
+        patch_sampled[:num_points] = patch_6d
+    
+    # Convert to tensor
+    patch_tensor = torch.from_numpy(patch_sampled.T).float().unsqueeze(0)  # (1, 6, max_points)
+    patch_tensor = patch_tensor.to(device)
+    
+    # Predict
+    with torch.no_grad():
+        outputs = model(patch_tensor)  # (1, 1)
+        probability = torch.sigmoid(outputs).item()
+        predicted_class = int(probability > 0.5)
+        confidence = probability if predicted_class == 1 else (1 - probability)
+        
+        return predicted_class, confidence

fast_voxel.py

@@ -531,7 +531,7 @@ def load_3dcnn_model(model_path: str, device: torch.device = None, voxel_size: i
     
     return model
 
-def predict_vertex_from_patch(model: Fast3DCNN, patch: np.ndarray, device: torch.device = None, voxel_size: int = 32) -> Tuple[np.ndarray, float, float]:
+def predict_vertex_from_patch_voxel(model: Fast3DCNN, patch: np.ndarray, device: torch.device = None, voxel_size: int = 32) -> Tuple[np.ndarray, float, float]:
     """
     Predict 3D vertex coordinates, confidence score, and classification from a patch using trained 3D CNN.
     

predict.py

@@ -393,18 +393,18 @@ def visu_patch_and_pred(patch, pred, pred_dist, pred_class):
     plotter = pv.Plotter()
     
     # Create point cloud for this patch
-    offset = patch.get('offset', None)  # Offset if available
+    offset = patch.get('cluster_center', None)  # Offset if available
     patch_points_3d = np.array(patch['patch_7d'][:, :3])
-    patch_points_3d = patch_points_3d - offset
+    patch_points_3d = patch_points_3d + offset
     patch_cloud = pv.PolyData(patch_points_3d)
     
-    point_idxs = patch['filtered_point_ids']  # List of point indices that are filtered
-    patch_point_ids = patch['point_ids']  # Assuming the 7th column contains point IDs
-    assigned_gt_vertex = patch.get('assigned_gt_vertex', None)  # GT vertex if available
-    initial_pred = patch.get('initial_pred', None)  # Initial prediction if available
-    initial_pred = initial_pred - offset
+    point_idxs = patch['cluster_point_ids']  # List of point indices that are filtered
+    patch_point_ids = patch['cube_point_ids']  # Assuming the 7th column contains point IDs
+    assigned_gt_vertex = patch.get('assigned_wf_vertex', None)  # GT vertex if available
+    initial_pred = None
     
-    assigned_gt_vertex = assigned_gt_vertex - offset
+    if assigned_gt_vertex is not None:
+        assigned_gt_vertex = assigned_gt_vertex + offset
 
     # Color points: red for filtered points, blue for other points
     patch_point_colors = []
@@ -436,17 +436,34 @@ def visu_patch_and_pred(patch, pred, pred_dist, pred_class):
     title_text = f"Patch x\nPred dist: {pred_dist:.4f}\nPred class: {pred_class}"
     plotter.show(title=title_text)
 
-def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points):
+def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points, all_connections):
     # Filter COLMAP points and colors based on idxs_points
     filtered_colmap_points = []
     filtered_colmap_colors = []
     filtered_colmap_ids = []
     all_filtered_ids_list = []
     all_extracted_groups = []
+    all_flattened_connections = []
+    group_to_flattened_mapping = {}  # Maps (group_idx, local_vertex_idx) to flattened_idx
 
+    # Flatten all groups and create mapping for connections
+    flattened_idx = 0
     for group_idx, point_ids_group in enumerate(idxs_points):
-        for point_ids in point_ids_group:
+        cur_connections = all_connections[group_idx]
+        group_to_flattened_mapping[group_idx] = {}
+        
+        for local_idx, point_ids in enumerate(point_ids_group):
             all_extracted_groups.append(point_ids)
+            group_to_flattened_mapping[group_idx][local_idx] = flattened_idx
+            flattened_idx += 1
+        
+        # Convert connections to flattened indices
+        for conn in cur_connections:
+            start_idx, end_idx = conn
+            if start_idx in group_to_flattened_mapping[group_idx] and end_idx in group_to_flattened_mapping[group_idx]:
+                flattened_start = group_to_flattened_mapping[group_idx][start_idx]
+                flattened_end = group_to_flattened_mapping[group_idx][end_idx]
+                all_flattened_connections.append((flattened_start, flattened_end))
 
     # Collect all filtered point IDs from all images
     for group_idxs in idxs_points:
@@ -523,13 +540,15 @@ def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points):
         extracted_colors.append(np.array(group_extracted_colors) if group_extracted_colors else np.empty((0, 3)))
         extracted_ids.append(np.array(group_extracted_ids) if group_extracted_ids else np.empty((0,)))
 
-
     # Filter extracted_points to merge groups that share more than 50% of their points
+    # and update connections accordingly
+    updated_connections = []
     if extracted_points:
         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
+        old_to_new_mapping = {}  # Maps old flattened index to new index
         
         for i, (points_i, colors_i, ids_i) in enumerate(zip(extracted_points, extracted_colors, extracted_ids)):
             if i in merged_groups or len(ids_i) == 0:
@@ -539,6 +558,7 @@ def extract_vertices_from_whole_pcloud(colmap_rec, idxs_points):
             merged_points = points_i.copy()
             merged_colors = colors_i.copy()
             merged_ids = set(ids_i)
+            merged_indices = [i]  # Track which original indices are merged
             
             # Check all subsequent groups for overlap
             for j in range(i + 1, len(extracted_points)):