#!/usr/bin/env python3 """ Extract Sample Features using miniSpinNet This script takes the output from generate_training_samples (either HDF5 file or folders) and extracts features for each sample using miniSpinNet encoder. For each sample: 1. Load all submaps (point clouds) in the sample 2. Combine all points from all submaps 3. Apply farthest point sampling to get K total points 4. Use miniSpinNet to extract features for each sampled point 5. Save sampled points + features as PLY files maintaining folder structure Usage: python ./dataset_process/extract_sample_features.py --input /path/to/training_data --output /path/to/features # For an indoor dataset (for example, NSS) python ./dataset_process/extract_sample_features.py --input ./dataset/lidar_rpf_training_data/nss_pair_v1 --output ./dataset/lidar_rpf_training_data/nss_pair_v1_processed_db_05 --des_r 0.5 --voxel_size 0.1 -r 0.5 --log_level DEBUG """ import os import sys import numpy as np import h5py import open3d as o3d import torch import argparse import logging from tqdm import tqdm from typing import List, Tuple, Optional, Dict, Union import json # Add the current directory and parent directory to the path so we can import our modules current_dir = os.path.dirname(os.path.abspath(__file__)) parent_dir = os.path.dirname(current_dir) sys.path.insert(0, current_dir) sys.path.insert(0, parent_dir) from utils.spinnet.patch_embedder import MiniSpinNet from utils.processing_utils import set_random_seeds from utils.io_utils import get_dataset_name, convert_to_hdf5, load_sample_from_folder, load_sample_from_hdf5, save_processed_sample from utils.feature_extraction_metadata_utils import save_processing_metadata, print_detailed_statistics from utils.validation_utils import _validate_and_setup_args from utils.dataset_utils import save_num_points_to_folder from utils.split_utils import copy_and_update_data_split from utils.point_sampling_utils import calculate_adaptive_sample_count_per_part, allocate_fps_points, apply_batched_fps # Setup logging logger = logging.getLogger(__name__) class FeatureExtractor: def __init__(self, model_config: Dict = None, des_r: float = 3.0, is_aligned_to_global_z: bool = True, checkpoint_path: str = None, device: str = 'auto'): if device == 'auto': self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') else: self.device = torch.device(device) # Default model configuration default_config = { 'num_points_per_patch': 512, 'rad_n': 3, 'azi_n': 20, 'ele_n': 7, 'delta': 0.8, 'voxel_sample': 10, } if model_config: default_config.update(model_config) self.model_config = default_config self.model = self._build_model() self.des_r = des_r self.is_aligned_to_global_z = is_aligned_to_global_z if checkpoint_path and os.path.exists(checkpoint_path): self._load_checkpoint(checkpoint_path) elif checkpoint_path: logger.warning(f"Checkpoint not found: {checkpoint_path}") def _build_model(self) -> MiniSpinNet: """Build miniSpinNet model.""" model = MiniSpinNet(**self.model_config) model.to(self.device) model.eval() return model def _load_checkpoint(self, checkpoint_path: str): try: state_dict = torch.load(checkpoint_path, map_location=self.device, weights_only=False) filtered = {k[5:]: v for k, v in state_dict.items() if k.startswith('Desc.')} self.model.load_state_dict(filtered, strict=False) self.model.eval() except Exception as e: logger.warning(f"Failed to load checkpoint: {e}") def extract_features(self, points: Union[np.ndarray, torch.Tensor], keypoints: Optional[Union[np.ndarray, torch.Tensor]] = None) -> Dict: def _to_tensor_with_batch(data): """Convert data to tensor with batch dimension.""" if isinstance(data, torch.Tensor): tensor = data.float().to(self.device) return tensor.unsqueeze(0) if tensor.dim() == 2 else tensor else: return torch.from_numpy(data).float().unsqueeze(0).to(self.device) # Handle empty inputs points_is_tensor = isinstance(points, torch.Tensor) is_empty = points.numel() == 0 if points_is_tensor else len(points) == 0 if is_empty: empty_result = torch.tensor([]) if points_is_tensor else np.array([]) return {'features': empty_result, 'keypoints': empty_result} # Use all points as keypoints if not specified if keypoints is None: keypoints = points.clone() if points_is_tensor else points.copy() # Convert to tensors with batch dimension points_tensor = _to_tensor_with_batch(points) keypoints_tensor = _to_tensor_with_batch(keypoints) with torch.no_grad(): try: # Extract features using miniSpinNet result = self.model( pts=points_tensor, kpts=keypoints_tensor, des_r=self.des_r, is_aligned_to_global_z=self.is_aligned_to_global_z ) # Extract features (descriptors) features = result['desc'] # Keep as tensor: (1, K, feature_dim) features = features.squeeze(0) if features.dim() == 3 else features # Remove batch dim # Return in same format as input return { 'features': features if points_is_tensor else features.cpu().numpy(), 'keypoints': keypoints, } except Exception as e: logger.warning(f"Feature extraction failed: {e}") # Return empty features on failure feature_dim = 32 # Default feature dimension for miniSpinNet keypoint_len = keypoints.shape[0] if keypoints.dim() > 1 else len(keypoints) if points_is_tensor: empty_features = torch.zeros(keypoint_len, feature_dim, device=self.device) else: empty_features = np.zeros((len(keypoints), feature_dim)) return {'features': empty_features, 'keypoints': keypoints} class SampleProcessor: def __init__(self, feature_extractor: FeatureExtractor, num_points: int = 5000, skip_point_sampling: bool = False, remove_outliers: bool = True, outlier_nb_neighbors: int = 20, outlier_std_ratio: float = 2.0, min_points_per_part: int = 100, max_points_per_part: int = 10000, global_seed: int = 42, allocation_method: str = 'point_count', voxel_size: float = 1.0, voxel_ratio: float = 0.1): self.feature_extractor = feature_extractor self.num_points = num_points self.skip_point_sampling = skip_point_sampling self.remove_outliers = remove_outliers self.outlier_nb_neighbors = outlier_nb_neighbors self.outlier_std_ratio = outlier_std_ratio self.min_points_per_part = min_points_per_part self.max_points_per_part = max_points_per_part self.global_seed = global_seed self.allocation_method = allocation_method self.voxel_size = voxel_size self.voxel_ratio = voxel_ratio if self.allocation_method not in ['point_count', 'spatial_coverage', 'voxel_adaptive']: raise ValueError(f"allocation_method must be 'point_count', 'spatial_coverage', or 'voxel_adaptive'") def _set_sample_seed(self, sample_idx: int = 0): seed = self.global_seed + sample_idx * 1000 np.random.seed(seed) torch.manual_seed(seed) def process_sample(self, parts_points: List[np.ndarray], parts_normals: List[Optional[np.ndarray]] = None) -> List[Dict]: if not parts_points or len(parts_points) == 0: return [] n_parts = len(parts_points) parts_normals = parts_normals or [None] * n_parts original_parts = [] original_normals = [] for i, (points, normals) in enumerate(zip(parts_points, parts_normals)): if len(points) == 0: continue original_parts.append(points.copy()) original_normals.append(normals.copy() if normals is not None else None) if not original_parts: return [] # Optional fast path: skip FPS and use all points as keypoints if self.skip_point_sampling: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') part_results = [] for i in range(len(original_parts)): points_tensor = torch.from_numpy(original_parts[i]).float().to(device) feat_res = self.feature_extractor.extract_features( points=points_tensor, keypoints=points_tensor, ) features = feat_res['features'] part_results.append({ 'sampled_points': original_parts[i], 'sampled_normals': original_normals[i] if original_normals[i] is not None else None, 'features': features.cpu().numpy() if isinstance(features, torch.Tensor) else features, }) return part_results fps_parts, fps_normals = original_parts, original_normals if self.remove_outliers: fps_parts = [] fps_normals = [] for i, (orig_part, orig_normals_i) in enumerate(zip(original_parts, original_normals)): if len(orig_part) < self.outlier_nb_neighbors: fps_parts.append(orig_part) fps_normals.append(orig_normals_i) continue try: pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(orig_part) # Remove statistical outliers pcd_filtered, inlier_indices = pcd.remove_statistical_outlier( nb_neighbors=self.outlier_nb_neighbors, std_ratio=self.outlier_std_ratio ) inlier_indices = np.array(inlier_indices) if len(inlier_indices) == 0: logger.warning(f"All points removed as outliers from part {i}, keeping original for FPS") fps_parts.append(orig_part) fps_normals.append(orig_normals_i) else: fps_parts.append(orig_part[inlier_indices]) fps_normals.append(orig_normals_i[inlier_indices] if orig_normals_i is not None else None) except Exception as e: logger.warning(f"Outlier removal failed for part {i}: {e}") fps_parts.append(orig_part) fps_normals.append(orig_normals_i) # Pre-FPS random downsampling for very large parts if self.max_points_per_part is not None: pre_fps_parts = [] pre_fps_normals = [] for i, (part, norms) in enumerate(zip(fps_parts, fps_normals)): pre_fps_cap = 20 * self.max_points_per_part if len(part) > pre_fps_cap: indices = np.random.choice(len(part), pre_fps_cap, replace=False) pre_fps_parts.append(part[indices].copy()) if norms is not None: pre_fps_normals.append(norms[indices].copy()) else: pre_fps_normals.append(None) else: pre_fps_parts.append(part) pre_fps_normals.append(norms) fps_parts = pre_fps_parts fps_normals = pre_fps_normals # FPS allocation if self.allocation_method == 'voxel_adaptive': # Calculate adaptive sample count based on occupied voxels after outlier removal adaptive_sample_counts_per_part = calculate_adaptive_sample_count_per_part( fps_parts, self.voxel_size, self.voxel_ratio, self.min_points_per_part, self.max_points_per_part ) total_adaptive_points = sum(adaptive_sample_counts_per_part) target_per_part = allocate_fps_points( fps_parts, self.allocation_method, self.num_points, # Not used for voxel_adaptive here, but kept for function signature consistency self.min_points_per_part, self.voxel_size, self.voxel_ratio, total_sample_points=adaptive_sample_counts_per_part ) else: if self.allocation_method == 'spatial_coverage': target_per_part = allocate_fps_points( fps_parts, self.allocation_method, self.num_points, self.min_points_per_part, self.voxel_size, self.voxel_ratio, ) else: pts_per_part = np.array([len(part) for part in fps_parts]) target_per_part = allocate_fps_points( pts_per_part, self.allocation_method, self.num_points, self.min_points_per_part, self.voxel_size, self.voxel_ratio, ) # Apply batched FPS device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Prepare batched data for PyTorch3D FPS batch_parts = [] batch_normals = [] batch_lengths = [] batch_k = [] for i, (part, part_normals, target_points) in enumerate( zip(fps_parts, fps_normals, target_per_part) ): if target_points == 0: continue batch_parts.append(torch.from_numpy(part).float()) if part_normals is not None: batch_normals.append(torch.from_numpy(part_normals).float()) else: batch_normals.append(None) batch_lengths.append(len(part)) batch_k.append(min(target_points, len(part))) # Don't exceed available points if not batch_parts: return [] max_points = max(batch_lengths) padded_parts = [] padded_normals = [] for i, (part, norms) in enumerate(zip(batch_parts, batch_normals)): if len(part) < max_points: pad_size = max_points - len(part) part_padded = torch.cat([part, torch.zeros(pad_size, 3)], dim=0) if norms is not None: norms_padded = torch.cat([norms, torch.zeros(pad_size, 3)], dim=0) else: norms_padded = None else: part_padded = part norms_padded = norms padded_parts.append(part_padded) padded_normals.append(norms_padded) # Stack into batch tensors batch_parts_tensor = torch.stack(padded_parts).to(device) # (N, max_points, 3) batch_lengths_tensor = torch.tensor(batch_lengths, dtype=torch.int64).to(device) # (N,) batch_k_tensor = torch.tensor(batch_k, dtype=torch.int64).to(device) # (N,) sampled_parts, indices_tensor = apply_batched_fps( batch_parts_tensor, batch_lengths_tensor, batch_k_tensor, self.global_seed, device ) sampled_points_list = [] sampled_normals_list = [] for i, (k_i, indices_i) in enumerate(zip(batch_k_tensor, indices_tensor)): valid_indices = indices_i[:k_i] sampled_points_list.append(batch_parts_tensor[i][valid_indices]) if padded_normals[i] is not None: normals_tensor = padded_normals[i].to(device) sampled_normals_list.append(normals_tensor[valid_indices]) else: sampled_normals_list.append(None) if not sampled_parts: return [] part_results = [] for i, (sampled_pts, sampled_norms) in enumerate(zip(sampled_points_list, sampled_normals_list)): if sampled_pts.numel() == 0: continue feature_result = self.feature_extractor.extract_features( points=torch.from_numpy(original_parts[i]).float().to(sampled_pts.device), keypoints=sampled_pts, ) features = feature_result['features'] part_result = { 'sampled_points': sampled_pts.cpu().numpy(), 'sampled_normals': sampled_norms.cpu().numpy() if sampled_norms is not None else None, 'features': features.cpu().numpy() if isinstance(features, torch.Tensor) else features, } part_results.append(part_result) return part_results def process_from_folders(input_dir: str, output_dir: str, processor: SampleProcessor, dataset_name: Optional[str] = None, hdf5_only: bool = False, dry_run: bool = False) -> Dict: dataset_name = get_dataset_name(input_dir, dataset_name) dataset_dir = os.path.join(input_dir, dataset_name) if not os.path.exists(dataset_dir): dataset_dir = input_dir # Find all sample directories (supports nested structures like ThreeDMatch) sample_dirs = [] if os.path.exists(dataset_dir): for sequence_name in os.listdir(dataset_dir): sequence_path = os.path.join(dataset_dir, sequence_name) if os.path.isdir(sequence_path): # Check for direct samples in this directory for sample_name in os.listdir(sequence_path): if sample_name.startswith('sample_'): sample_path = os.path.join(sequence_path, sample_name) if os.path.isdir(sample_path): relative_path = os.path.join(sequence_name, sample_name) sample_dirs.append((relative_path, sample_path)) for subdir_name in os.listdir(sequence_path): subdir_path = os.path.join(sequence_path, subdir_name) if os.path.isdir(subdir_path) and not subdir_name.startswith('sample_'): for sample_name in os.listdir(subdir_path): if sample_name.startswith('sample_') or sample_name.startswith('fracture_'): sample_path = os.path.join(subdir_path, sample_name) if os.path.isdir(sample_path): sample_dirs.append((os.path.join(sequence_name, subdir_name, sample_name), sample_path)) if not sample_dirs: logger.error(f"No sample directories found in {input_dir}") return {'processed_samples': 0, 'failed_samples': 0} logger.info(f"Found {len(sample_dirs)} samples to process") if hdf5_only or dry_run: processed_count = len(sample_dirs) failed_count = 0 sample_num_points = [] # Empty for dry run/hdf5_only mode else: # Process samples processed_count = 0 failed_count = 0 sample_num_points = [] # Track num_points for each sample sample_part_counts = [] # Track number of parts per sample sample_part_points = [] # Track points per part for each sample all_part_points = [] # Track all individual part point counts for relative_path, sample_path in tqdm(sample_dirs, desc="Processing samples"): try: logger.debug(f"Processing sample: {relative_path}") # Set sample-specific seed processor._set_sample_seed(processed_count) # Load sample parts_points, parts_normals, part_names = load_sample_from_folder(sample_path) if not parts_points: logger.warning(f"No parts loaded for sample: {relative_path}") failed_count += 1 sample_num_points.append(0) sample_part_counts.append(0) sample_part_points.append([]) continue # Process sample part_results = processor.process_sample(parts_points, parts_normals) # Calculate statistics for this sample part_point_counts = [len(part_result['sampled_points']) for part_result in part_results] total_sample_points = sum(part_point_counts) sample_num_points.append(total_sample_points) sample_part_counts.append(len(part_results)) sample_part_points.append(part_point_counts) all_part_points.extend(part_point_counts) # Save processed sample parts sample_output_dir = os.path.join(output_dir, dataset_name, relative_path + '_processed') save_processed_sample(part_results, part_names, sample_output_dir, input_sample_dir=sample_path) processed_count += 1 except Exception as e: logger.error(f"Failed to process sample {relative_path}: {e}") sample_num_points.append(0) # Add 0 for failed samples sample_part_counts.append(0) sample_part_points.append([]) failed_count += 1 logger.info(f"Processing complete: {processed_count} processed, {failed_count} failed") # Copy and update data_split folder if not dry_run: copy_and_update_data_split(input_dir, output_dir, dataset_name) # Save num_points data to folder structure if sample_num_points: save_num_points_to_folder(output_dir, dataset_name, sample_num_points, sample_dirs) return { 'processed_samples': processed_count, 'failed_samples': failed_count, 'total_samples': len(sample_dirs), 'sample_num_points': sample_num_points, # Include for HDF5 conversion 'sample_part_counts': sample_part_counts if not (hdf5_only or dry_run) else [], 'sample_part_points': sample_part_points if not (hdf5_only or dry_run) else [], 'all_part_points': all_part_points if not (hdf5_only or dry_run) else [] } def process_from_hdf5(hdf5_path: str, output_dir: str, processor: SampleProcessor, dataset_name: Optional[str] = None, hdf5_only: bool = False, dry_run: bool = False) -> Dict: # Auto-detect dataset name if not provided dataset_name = get_dataset_name(hdf5_path, dataset_name) processed_count = 0 failed_count = 0 try: with h5py.File(hdf5_path, 'r') as h5_file: # Find all sample paths sample_paths = [] def collect_samples(name, obj): if isinstance(obj, h5py.Group): # Check if this looks like a sample path (contains numeric submaps) if any(key.isdigit() for key in obj.keys()): sample_paths.append(name) h5_file.visititems(collect_samples) if not sample_paths: logger.error(f"No samples found in HDF5 file: {hdf5_path}") return {'processed_samples': 0, 'failed_samples': 0} logger.info(f"Found {len(sample_paths)} samples to process") if hdf5_only or dry_run: processed_count = len(sample_paths) failed_count = 0 sample_num_points = [] # Empty for dry run/hdf5_only mode else: # Process samples sample_num_points = [] # Track num_points for each sample sample_part_counts = [] # Track number of parts per sample sample_part_points = [] # Track points per part for each sample all_part_points = [] # Track all individual part point counts for sample_path in tqdm(sample_paths, desc="Processing samples"): try: logger.debug(f"Processing sample: {sample_path}") # Set sample-specific seed processor._set_sample_seed(processed_count) # Load sample parts_points, parts_normals, part_names = load_sample_from_hdf5(h5_file, sample_path) if not parts_points: logger.warning(f"No parts loaded for sample: {sample_path}") failed_count += 1 sample_num_points.append(0) sample_part_counts.append(0) sample_part_points.append([]) continue # Process sample part_results = processor.process_sample(parts_points, parts_normals) # Calculate statistics for this sample part_point_counts = [len(part_result['sampled_points']) for part_result in part_results] total_sample_points = sum(part_point_counts) sample_num_points.append(total_sample_points) sample_part_counts.append(len(part_results)) sample_part_points.append(part_point_counts) all_part_points.extend(part_point_counts) # Save processed sample parts sample_output_dir = os.path.join(output_dir, dataset_name, sample_path + '_processed') save_processed_sample(part_results, part_names, sample_output_dir) processed_count += 1 except Exception as e: logger.error(f"Failed to process sample {sample_path}: {e}") sample_num_points.append(0) # Add 0 for failed samples sample_part_counts.append(0) sample_part_points.append([]) failed_count += 1 except Exception as e: logger.error(f"Failed to open HDF5 file {hdf5_path}: {e}") return {'processed_samples': 0, 'failed_samples': 0} logger.info(f"Processing complete: {processed_count} processed, {failed_count} failed") # Copy and update data_split folder hdf5_dir = os.path.dirname(hdf5_path) if not dry_run: copy_and_update_data_split(hdf5_dir, output_dir, dataset_name) # Save num_points data to folder structure if sample_num_points: sample_dirs = [(p, '') for p in sample_paths] save_num_points_to_folder(output_dir, dataset_name, sample_num_points, sample_dirs) return { 'processed_samples': processed_count, 'failed_samples': failed_count, 'total_samples': len(sample_paths), 'sample_num_points': sample_num_points, # Include for HDF5 conversion 'sample_part_counts': sample_part_counts if not (hdf5_only or dry_run) else [], 'sample_part_points': sample_part_points if not (hdf5_only or dry_run) else [], 'all_part_points': all_part_points if not (hdf5_only or dry_run) else [] } def main(): parser = argparse.ArgumentParser(description='Extract features from training samples using miniSpinNet') # Input/Output arguments parser.add_argument('--input', '-i', type=str, required=True, help='Input path (directory with PLY files or HDF5 file)') parser.add_argument('--output', '-o', type=str, required=True, help='Output directory for processed samples with features') parser.add_argument('--dataset_name', type=str, default=None, help='Dataset name (auto-detected if not provided)') # HDF5 output arguments parser.add_argument('--save_hdf5', action='store_true', default=True, help='Convert processed PLY/NPY files to HDF5 format after processing') parser.add_argument('--hdf5_output', type=str, default=None, help='Path for output HDF5 file (auto-generated if not provided when --save_hdf5 is used)') parser.add_argument('--hdf5_only', action='store_true', default=False, help='Only convert existing PLY/NPY files to HDF5, skip feature extraction processing') # Processing arguments parser.add_argument('--num_points', '-k', type=int, default=5000, help='Number of points to sample using FPS (default: 5000), now deprecated') parser.add_argument('--global_seed', type=int, default=42, help='Global random seed for all random operations including FPS (default: 42)') parser.add_argument('--min_points_per_part', type=int, default=300, help='Minimum number of points each part should have after FPS (default: 100)') parser.add_argument('--max_points_per_part', type=int, default=10000, help='Maximum number of points each part should have after FPS (default: 10000)') parser.add_argument('--skip_point_sampling', action='store_true', default=False, help='Skip farthest point sampling and use all points as keypoints for feature extraction') # Outlier removal arguments parser.add_argument('--remove_outliers', action='store_true', default=True, help='Remove statistical outliers for FPS sampling, but keep all points for feature extraction context (default: True)') parser.add_argument('--no_remove_outliers', dest='remove_outliers', action='store_false', help='Disable statistical outlier removal') parser.add_argument('--outlier_nb_neighbors', type=int, default=20, help='Number of neighbors for outlier removal (default: 20)') parser.add_argument('--outlier_std_ratio', type=float, default=2.5, help='Standard deviation ratio for outlier removal (default: 2.5)') # Allocation method and voxel size parser.add_argument('--allocation_method', type=str, default='voxel_adaptive', choices=['point_count', 'spatial_coverage', 'voxel_adaptive'], help='Method for allocating FPS points (default: voxel_adaptive)') parser.add_argument('--voxel_size', type=float, default=1.0, help='Voxel size in meters for spatial coverage calculation (default: 1.0)') parser.add_argument('--voxel_ratio', '-r', type=float, default=0.05, help='Ratio of occupied voxels to sample points for voxel_adaptive method (default: 0.05), should also be considered together with --voxel_size, \ the smaller the ratio, the fewer the sample points, we set a larger value (for example 0.2) for scan level data') # Model arguments parser.add_argument('--checkpoint', type=str, default='./weights/weights/spinnet_3dmatch_bufferx.pth', help='Path to miniSpinNet checkpoint (default: ./weights/weights/spinnet_3dmatch_bufferx.pth), select from ./weights/weights/spinnet_3dmatch_bufferx.pth, ./weights/weights/spinnet_kitti_bufferx.pth') parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cpu', 'cuda'], help='Device to use for feature extraction (default: auto)') parser.add_argument('--is_aligned_to_global_z', action='store_true', default=True, help='Align point clouds to global Z axis before feature extraction (default: True)') parser.add_argument('--no_is_aligned_to_global_z', dest='is_aligned_to_global_z', action='store_false', help='Do not align point clouds to global Z axis before feature extraction') # miniSpinNet configuration parser.add_argument('--des_r', type=float, default=5.0, help='Description radius for miniSpinNet in meters (default: 5.0)') parser.add_argument('--num_points_per_patch', type=int, default=512, help='Number of points per patch for miniSpinNet (default: 512)') # Utility arguments parser.add_argument('--log_level', type=str, default='INFO', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], help='Logging level (default: INFO)') parser.add_argument('--dry_run', action='store_true', default=False, help='Show what would be processed without actually doing it') args = parser.parse_args() # Setup logging logging.basicConfig( level=getattr(logging, args.log_level), format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) # Set random seeds for reproducibility set_random_seeds(args.global_seed) # Validate arguments if not _validate_and_setup_args(args): return if not args.dry_run: os.makedirs(args.output, exist_ok=True) if not args.hdf5_only: if args.dry_run: processor = None else: model_config = { 'num_points_per_patch': args.num_points_per_patch, 'is_aligned_to_global_z': args.is_aligned_to_global_z, 'des_r': args.des_r, } feature_extractor = FeatureExtractor( model_config=model_config, des_r=args.des_r, is_aligned_to_global_z=args.is_aligned_to_global_z, checkpoint_path=args.checkpoint, device=args.device ) processor = SampleProcessor( feature_extractor=feature_extractor, num_points=args.num_points, skip_point_sampling=args.skip_point_sampling, remove_outliers=args.remove_outliers, outlier_nb_neighbors=args.outlier_nb_neighbors, outlier_std_ratio=args.outlier_std_ratio, min_points_per_part=args.min_points_per_part, max_points_per_part=args.max_points_per_part, global_seed=args.global_seed, allocation_method=args.allocation_method, voxel_size=args.voxel_size, voxel_ratio=args.voxel_ratio, ) else: processor = None if args.input.endswith('.hdf5') or args.input.endswith('.h5'): stats = process_from_hdf5(args.input, args.output, processor, args.dataset_name, args.hdf5_only, args.dry_run) else: stats = process_from_folders(args.input, args.output, processor, args.dataset_name, args.hdf5_only, args.dry_run) if args.save_hdf5 or args.hdf5_only: dataset_name = get_dataset_name(args.output, args.dataset_name) if not args.dry_run: convert_to_hdf5(args.output, dataset_name, args.hdf5_output, stats, args) if not args.dry_run: save_processing_metadata(args.output, stats, args) comprehensive_metadata = { 'processing_summary': { 'input_path': args.input, 'output_path': args.output, 'dataset_name': args.dataset_name, 'processed_samples': stats.get('processed_samples', 0), 'failed_samples': stats.get('failed_samples', 0), 'total_samples': stats.get('total_samples', 0), 'success_rate': f"{stats.get('processed_samples', 0) / max(stats.get('total_samples', 1), 1) * 100:.1f}%" }, 'feature_extraction_config': { 'feature_extractor': 'miniSpinNet', 'num_points_fps': args.num_points, 'min_points_per_part': args.min_points_per_part, 'des_r': args.des_r, 'num_points_per_patch': args.num_points_per_patch, 'is_aligned_to_global_z': args.is_aligned_to_global_z, 'checkpoint_path': args.checkpoint, 'device': args.device }, 'processing_config': { 'global_seed': args.global_seed, 'skip_point_sampling': args.skip_point_sampling, 'remove_outliers': args.remove_outliers, 'outlier_nb_neighbors': args.outlier_nb_neighbors, 'outlier_std_ratio': args.outlier_std_ratio, 'allocation_method': args.allocation_method, 'voxel_size': args.voxel_size, 'voxel_ratio': args.voxel_ratio, }, 'output_config': { 'save_hdf5': args.save_hdf5, 'hdf5_only': args.hdf5_only, 'hdf5_output': args.hdf5_output } } with open(os.path.join(args.output, 'comprehensive_metadata.json'), 'w') as f: json.dump(comprehensive_metadata, f, indent=2, default=str) print_detailed_statistics(stats, args) if __name__ == "__main__": main()