x

fast_pointnet_class_v2.py

@@ -0,0 +1,464 @@
+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 10D point cloud patches.
+    Takes 10D point clouds and predicts binary classification (edge/not edge).
+    """
+    def __init__(self, input_dim=10, max_points=1024): # Changed input_dim default to 10
+        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) # Changed input_dim here
+        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, input_dim: int = 10): # Added input_dim
+        self.dataset_dir = dataset_dir
+        self.max_points = max_points
+        self.augment = augment
+        self.input_dim = input_dim # Store input_dim
+        
+        # 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: (input_dim, 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)
+        
+        # Assuming the key in patch_info is now 'patch_10d' or similar, or that patch_info['patch_data'] is (N, 10)
+        # For this example, let's assume the key is 'patch_data' and it holds the 10D data.
+        # If your key is 'patch_10d', change 'patch_data' to 'patch_10d' below.
+        patch_data_nd = patch_info.get('patch_data', patch_info.get('patch_10d', patch_info.get('patch_6d'))) # Try to get 10d, fallback to 6d for now
+        if patch_data_nd.shape[1] != self.input_dim:
+            # This is a fallback or error handling if the loaded data isn't 10D.
+            # You might want to raise an error or handle this case specifically.
+            # For now, if it's 6D, we'll pad it to 10D with zeros as a placeholder.
+            # This part needs to be adjusted based on how your 10D data is actually stored.
+            print(f"Warning: Patch {patch_file} has {patch_data_nd.shape[1]} dimensions, expected {self.input_dim}. Padding with zeros if necessary.")
+            if patch_data_nd.shape[1] < self.input_dim:
+                padding = np.zeros((patch_data_nd.shape[0], self.input_dim - patch_data_nd.shape[1]))
+                patch_data_nd = np.concatenate((patch_data_nd, padding), axis=1)
+            elif patch_data_nd.shape[1] > self.input_dim:
+                patch_data_nd = patch_data_nd[:, :self.input_dim]
+
+
+        label = patch_info.get('label', 0)  # Get binary classification label (0 or 1)
+        
+        # Pad or sample points to max_points
+        num_points = patch_data_nd.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_data_nd[indices]
+            valid_mask = np.ones(self.max_points, dtype=bool)
+        else:
+            # Pad with zeros
+            patch_sampled = np.zeros((self.max_points, self.input_dim)) # Changed to self.input_dim
+            patch_sampled[:num_points] = patch_data_nd
+            valid_mask = np.zeros(self.max_points, dtype=bool)
+            valid_mask[:num_points] = True
+        
+        # Data augmentation
+        if self.augment:
+            # Note: _augment_patch currently only augments xyz (first 3 dims).
+            # If other dimensions are geometric and need augmentation, this function needs an update.
+            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()  # (input_dim, 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.
+        Note: This implementation only augments the first 3 dimensions (assumed to be XYZ).
+        If your 10D representation has other geometric features that need augmentation,
+        this function should be updated accordingly.
+        """
+        valid_points_data = patch[valid_mask]
+        
+        if len(valid_points_data) == 0:
+            return patch
+        
+        # Extract XYZ for augmentation (first 3 columns)
+        valid_points_xyz = valid_points_data[:, :3].copy() # Operate on a copy
+
+        # 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_xyz = valid_points_xyz @ rotation_matrix.T
+        
+        # Random jittering
+        noise = np.random.normal(0, 0.01, valid_points_xyz.shape)
+        valid_points_xyz += noise
+        
+        # Random scaling
+        scale = np.random.uniform(0.9, 1.1)
+        valid_points_xyz *= scale
+        
+        # Update the original patch data
+        augmented_patch = patch.copy()
+        augmented_patch[valid_mask, :3] = valid_points_xyz
+        
+        return augmented_patch
+
+def save_patches_dataset(patches: List[Dict], dataset_dir: str, entry_id: str):
+    """
+    Save patches from prediction pipeline to create a training dataset.
+    Ensure 'patch_data' (or 'patch_10d') in the patch dictionary contains the 10D data.
+    
+    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
+        
+        # Ensure the patch data being saved is 10D.
+        # Example: patch_data_key = 'patch_10d' or 'patch_data'
+        # if 'patch_data' not in patch or patch['patch_data'].shape[1] != 10:
+        #     print(f"Warning: Patch {i} for entry {entry_id} does not seem to be 10D. Skipping or error handling needed.")
+        #     continue
+
+        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, input_dim: int = 10): # Added input_dim
+    """
+    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
+        input_dim: Dimensionality of the input points (e.g., 10 for 10D)
+    """
+    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, input_dim=input_dim) # Pass input_dim
+    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=input_dim, max_points=1024) # Pass input_dim
+    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, input_dim, 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 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,
+            'accuracy': accuracy,
+            'input_dim': input_dim, # Save input_dim with checkpoint
+        }, checkpoint_path)
+    
+    # Save the trained model
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'epoch': epochs,
+        'input_dim': input_dim, # Save input_dim with final model
+    }, 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')
+    
+    checkpoint = torch.load(model_path, map_location=device)
+    
+    # Load input_dim from checkpoint if available, otherwise default to 10
+    # For older models saved without input_dim, you might need to specify it or assume a default.
+    input_dim = checkpoint.get('input_dim', 10) 
+    
+    model = ClassificationPointNet(input_dim=input_dim, max_points=1024) # Use loaded or default input_dim
+    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.
+    Assumes the model's input_dim matches the data.
+    
+    Args:
+        model: Trained ClassificationPointNet model
+        patch: Dictionary containing patch data. Expects a key like 'patch_data' or 'patch_10d' with (N, 10) shape.
+        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')
+
+    # Determine input_dim from the model
+    input_dim = model.conv1.in_channels
+
+    # Assuming the key in patch_info is now 'patch_10d' or similar, or that patch_info['patch_data'] is (N, 10)
+    # For this example, let's assume the key is 'patch_data' and it holds the 10D data.
+    # If your key is 'patch_10d', change 'patch_data' to 'patch_10d' below.
+    patch_data_nd = patch.get('patch_data', patch.get('patch_10d', patch.get('patch_6d'))) # Try to get 10d, fallback to 6d
+    
+    if patch_data_nd.shape[1] != input_dim:
+        # Handle dimension mismatch, e.g., by padding or raising an error
+        print(f"Warning: Input patch has {patch_data_nd.shape[1]} dimensions, but model expects {input_dim}. Adjusting...")
+        if patch_data_nd.shape[1] < input_dim:
+            padding = np.zeros((patch_data_nd.shape[0], input_dim - patch_data_nd.shape[1]))
+            patch_data_nd = np.concatenate((patch_data_nd, padding), axis=1)
+        elif patch_data_nd.shape[1] > input_dim:
+            patch_data_nd = patch_data_nd[:, :input_dim]
+            
+    # Prepare input
+    max_points = model.max_points # Use max_points from the model instance
+    num_points = patch_data_nd.shape[0]
+    
+    if num_points >= max_points:
+        # Sample points
+        indices = np.random.choice(num_points, max_points, replace=False)
+        patch_sampled = patch_data_nd[indices]
+    else:
+        # Pad with zeros
+        patch_sampled = np.zeros((max_points, input_dim)) # Use model's input_dim
+        patch_sampled[:num_points] = patch_data_nd
+    
+    # Convert to tensor
+    patch_tensor = torch.from_numpy(patch_sampled.T).float().unsqueeze(0)  # (1, input_dim, max_points)
+    patch_tensor = patch_tensor.to(device)
+    
+    # Predict
+    model.eval() # Ensure model is in eval mode
+    with torch.no_grad():
+        outputs = model(patch_tensor)  # (1, 1)
+        probability = torch.sigmoid(outputs).item()
+        predicted_class = int(probability > 0.5)
+        
+        return predicted_class, probability
+

predict.py

@@ -15,8 +15,8 @@ import cv2
 from fast_pointnet_v2 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
-from fast_pointnet_class import predict_class_from_patch
+from fast_pointnet_class_v2 import save_patches_dataset as save_patches_dataset_class
+from fast_pointnet_class_v2 import predict_class_from_patch
 #from fast_pointnet_class_10d import predict_class_from_patch as predict_class_from_patch_10d
 from scipy.spatial.distance import cdist
 from scipy.optimize import linear_sum_assignment
@@ -28,9 +28,9 @@ GENERATE_DATASET = False
 #DATASET_DIR = '/home/skvrnjan/personal/hohocustom/'
 DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_v4/'
 
-GENERATE_DATASET_EDGES = False
+GENERATE_DATASET_EDGES = True
 #EDGES_DATASET_DIR = '/home/skvrnjan/personal/hohocustom_edges/'
-EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges_10d_1m/'
+EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges_10d_v4/'
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -1010,18 +1010,13 @@ def generate_edge_patches(frame, pred_vertices, colmap_pcloud):
         elif len(point_gestalt_list) == 1:
             fused_gestalt.append(point_gestalt_list[0])
         else:
-            # Convert to numpy array for easier manipulation
-            gestalt_values = np.array(point_gestalt_list)
-            
-            # Method 1: Average the RGB values
-            fused_value = np.mean(gestalt_values, axis=0).astype(np.uint8)
+            # Convert to tuples for hashable voting
+            gestalt_tuples = [tuple(gestalt_val) for gestalt_val in point_gestalt_list]
             
-            # Method 2: Majority voting per channel (commented out alternative)
-            # fused_value = np.array([
-            #     np.bincount(gestalt_values[:, 0]).argmax(),
-            #     np.bincount(gestalt_values[:, 1]).argmax(), 
-            #     np.bincount(gestalt_values[:, 2]).argmax()
-            # ])
+            # Use Counter for majority voting
+            counts = Counter(gestalt_tuples)
+            most_common_tuple = counts.most_common(1)[0][0]
+            fused_value = np.array(most_common_tuple, dtype=np.uint8)
             
             fused_gestalt.append(fused_value)
     
@@ -1078,7 +1073,7 @@ def generate_edge_patches(frame, pred_vertices, colmap_pcloud):
         # Find points within cylinder
         within_cylinder = within_bounds & (perpendicular_distances <= cylinder_radius)
         
-        if np.sum(within_cylinder) <= 10:
+        if np.sum(within_cylinder) <= 5:
             continue
 
         points_in_cylinder = colmap_points_10d[within_cylinder]

train.py

@@ -26,8 +26,8 @@ import time
 
 # --- Argument Parsing ---
 parser = argparse.ArgumentParser(description="Train and evaluate HoHo model with custom config.")
-parser.add_argument('--vertex_threshold', type=float, default=0.6, help='Vertex threshold for prediction.')
-parser.add_argument('--edge_threshold', type=float, default=0.65, help='Edge threshold for prediction.')
+parser.add_argument('--vertex_threshold', type=float, default=0.72, help='Vertex threshold for prediction.')
+parser.add_argument('--edge_threshold', type=float, default=0.72, help='Edge threshold for prediction.')
 parser.add_argument('--only_predicted_connections', type=lambda x: (str(x).lower() == 'true'), default=True, help='Use only predicted connections (True/False).')
 parser.add_argument('--max_samples', type=int, default=50000, help='Maximum number of samples to process.')
 parser.add_argument('--results_dir', type=str, default="results", help='Directory to save result files.')
@@ -75,7 +75,7 @@ voxel_model = None
 
 idx = 0
 prediction_times = []
-for a in tqdm(ds['validation'], desc="Processing dataset"):
+for a in tqdm(ds['train'], desc="Processing dataset"):
     #plot_all_modalities(a)
     #pred_vertices, pred_edges = predict_wireframe_old(a)
     #pred_vertices, pred_edges = predict_wireframe(a.copy(), pnet_model, voxel_model, pnet_class_model, config)