dataset_process/visualize_registered_pointclouds.py

#!/usr/bin/env python3
"""
Visualize Registered Point Clouds using Estimated Poses

This script loads point clouds from an input directory and transforms them using
estimated poses from a result directory. Visualizes one sample at a time with
keyboard navigation (spacebar for next sample).

Usage:
    python visualize_registered_pointclouds.py \
        --input ./dataset/lidar_rpf_training_data/scannet_test_v1 \
        --results ./dataset/lidar_rpf_training_data/scannet_test_v1_result \
        --sequence scene0025_01
    
    # Visualize all sequences
    python visualize_registered_pointclouds.py \
        --input ./dataset/lidar_rpf_training_data/scannet_test_v1 \
        --results ./dataset/lidar_rpf_training_data/scannet_test_v1_result
"""

import argparse
import os
import sys

# Add project root so "dataset_process" can be imported when run from any directory
_script_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.dirname(_script_dir)
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)

import glob
import re
import logging
import json
import numpy as np
import open3d as o3d
import natsort
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Union
from dataset_process.utils.io_utils import CMAP_DEFAULT

# Logger will be configured in main() after parsing arguments
logger = logging.getLogger(__name__)

class RegisteredPointCloudVisualizer:
    """Visualizer for registered point clouds using estimated poses."""
    
    def __init__(self,
                 input_dir: str,
                 poses_dir: str,
                 point_size: float = 3.0,
                 background_color: List[float] = [1.0, 1.0, 1.0],
                 no_coordinate_frame: bool = True,
                 estimate_normals: bool = False,
                 normal_estimation_radius: float = 0.1,
                 generation: str = "generation00",
                 max_points_per_fragment: int = 10000,
                 remove_outliers: bool = False,
                 outlier_nb_neighbors: int = 20,
                 outlier_std_ratio: float = 2.0):
        """
        Initialize the registered point cloud visualizer.
        
        Args:
            input_dir: Root directory containing point cloud files
            poses_dir: Root directory containing pose files
            point_size: Size of points in the visualization
            background_color: Background color as [R, G, B] in range [0, 1]
            no_coordinate_frame: Whether to hide coordinate frame
            estimate_normals: If True, estimate normals for point clouds that don't have them
            normal_estimation_radius: Radius for normal estimation
            generation: Generation subdirectory name (default: "generation00")
            max_points_per_fragment: Maximum points to load per fragment
            remove_outliers: If True, apply stochastic outlier removal to input point clouds
            outlier_nb_neighbors: Number of neighbors for outlier removal (default: 20)
            outlier_std_ratio: Standard deviation ratio for outlier removal (default: 2.0)
        """
        self.input_dir = input_dir
        self.poses_dir = poses_dir
        self.point_size = point_size
        self.background_color = background_color
        self.no_coordinate_frame = no_coordinate_frame
        self.estimate_normals = estimate_normals
        self.normal_estimation_radius = normal_estimation_radius
        self.generation = generation
        self.max_points_per_fragment = max_points_per_fragment
        self.remove_outliers = remove_outliers
        self.outlier_nb_neighbors = outlier_nb_neighbors
        self.outlier_std_ratio = outlier_std_ratio
        
        # Navigation state
        self.current_sequence = None
        self.current_sample_idx = None
        self.sequence_samples_cache = {}  # Cache for sequence sample lists
        self.vis = None
        self.colored_pcds = []
        self.original_pcds = []  # Store original (untransformed) point clouds
        self.registered_pcds = []  # Store registered (transformed) point clouds
        self.generated_pcds = []  # Store generated point clouds from processed directory
        self.generated_gt_pcd = None  # Store merged GT point cloud
        self.random_yaw_pcds = []  # Store input point clouds with random yaw rotations
        self.view_mode = "registered"  # "original", "registered", "generated", "gt", or "random_yaw"
        
        # Generation cycling state
        self.cycle_generations = False  # Whether to cycle through generations instead of samples
        self.current_generation_idx = 0  # Current generation index when cycling
        self.available_generations = []  # List of available generation names for current sample
        
        # Background toggle state
        self.use_white_background = True
        self.white_background = [1.0, 1.0, 1.0]
        self.black_background = [0.0, 0.0, 0.0]
        
        if np.mean(self.background_color) > 0.5:
            self.use_white_background = True
        else:
            self.use_white_background = False
    
    def compute_adaptive_normal_radius(self, points: np.ndarray) -> float:
        """
        Compute adaptive normal estimation radius based on point cloud scale.
        Uses 2% of average extent, with a minimum of 0.05.
        
        Args:
            points: Point cloud points as numpy array (N, 3)
        
        Returns:
            Adaptive radius value
        """
        if len(points) == 0:
            return self.normal_estimation_radius
        
        # Compute extent (bounding box size)
        min_coords = np.min(points, axis=0)
        max_coords = np.max(points, axis=0)
        extent = max_coords - min_coords
        
        # Use adaptive radius based on point cloud scale (2% of average extent)
        adaptive_radius = max(np.mean(extent) * 0.02, 0.05)
        
        return adaptive_radius
    
    def apply_outlier_removal(self, pcd: o3d.geometry.PointCloud) -> o3d.geometry.PointCloud:
        """
        Apply stochastic outlier removal to a point cloud.
        
        Args:
            pcd: Open3D point cloud
        
        Returns:
            Point cloud with outliers removed
        """
        if not self.remove_outliers:
            return pcd
        
        points = np.asarray(pcd.points)
        if len(points) < self.outlier_nb_neighbors:
            # Not enough points for outlier removal
            logger.debug(f"Skipping outlier removal: only {len(points)} points (need at least {self.outlier_nb_neighbors})")
            return pcd
        
        try:
            # Apply statistical outlier removal
            # This returns (cleaned_pcd, indices) where indices are the inlier indices
            pcd_cleaned, ind = pcd.remove_statistical_outlier(
                nb_neighbors=self.outlier_nb_neighbors,
                std_ratio=self.outlier_std_ratio
            )
            
            num_removed = len(points) - len(np.asarray(pcd_cleaned.points))
            if num_removed > 0:
                logger.debug(f"Removed {num_removed} outliers ({num_removed/len(points)*100:.1f}%)")
            
            # Colors and normals are automatically preserved by remove_statistical_outlier
            # as it filters by indices, so we don't need to manually copy them
            
            return pcd_cleaned
        except Exception as e:
            logger.warning(f"Failed to apply outlier removal: {e}, returning original point cloud")
            return pcd
    
    def center_pcd(self, pcd: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Center point cloud at origin."""
        center = np.mean(pcd, axis=0)
        pcd = pcd - center
        return pcd, center
    
    def translate_point_cloud(self, pcd: o3d.geometry.PointCloud, translation: np.ndarray) -> o3d.geometry.PointCloud:
        """
        Translate a point cloud by a given translation vector.
        
        Args:
            pcd: Open3D point cloud
            translation: Translation vector (3D)
        
        Returns:
            Translated point cloud
        """
        points = np.asarray(pcd.points)
        translated_points = points - translation  # Subtract to move towards origin
        
        translated_pcd = o3d.geometry.PointCloud()
        translated_pcd.points = o3d.utility.Vector3dVector(translated_points)
        
        # Copy normals if available
        if pcd.has_normals():
            translated_pcd.normals = pcd.normals
        
        # Copy colors if available
        if pcd.has_colors():
            translated_pcd.colors = pcd.colors
        
        return translated_pcd
    
    def apply_random_yaw_rotation(self, pcd: o3d.geometry.PointCloud, skip_rotation: bool = False) -> o3d.geometry.PointCloud:
        """
        Apply random yaw rotation (around z-axis) with small roll and pitch perturbations to a point cloud.
        First centers the point cloud at origin, then applies random rotations:
        - Yaw: random rotation around z-axis (0 to 2π)
        - Roll: random rotation around x-axis (±15 degrees)
        - Pitch: random rotation around y-axis (±15 degrees)
        
        Args:
            pcd: Open3D point cloud
            skip_rotation: If True, only center the point cloud without applying rotation
        
        Returns:
            Point cloud with random rotations applied (or just centered if skip_rotation=True)
        """
        points = np.asarray(pcd.points)
        
        if len(points) == 0:
            return pcd
        
        # Center the point cloud at origin
        center = np.mean(points, axis=0)
        centered_points = points - center
        
        # If skip_rotation is True, just return centered point cloud
        if skip_rotation:
            centered_pcd = o3d.geometry.PointCloud()
            centered_pcd.points = o3d.utility.Vector3dVector(centered_points)
            
            # Copy normals if available
            if pcd.has_normals():
                centered_pcd.normals = pcd.normals
            
            # Copy colors if available
            if pcd.has_colors():
                centered_pcd.colors = pcd.colors
            
            return centered_pcd
        
        # Generate random rotation angles
        yaw_angle = np.random.uniform(0, 2 * np.pi)  # Full yaw rotation
        roll_angle = np.random.uniform(-15, 15) * np.pi / 180  # ±15 degrees roll
        pitch_angle = np.random.uniform(-15, 15) * np.pi / 180  # ±15 degrees pitch
        
        # Create rotation matrices
        # Roll rotation around x-axis
        cos_roll = np.cos(roll_angle)
        sin_roll = np.sin(roll_angle)
        R_roll = np.array([
            [1, 0, 0],
            [0, cos_roll, -sin_roll],
            [0, sin_roll, cos_roll]
        ])
        
        # Pitch rotation around y-axis
        cos_pitch = np.cos(pitch_angle)
        sin_pitch = np.sin(pitch_angle)
        R_pitch = np.array([
            [cos_pitch, 0, sin_pitch],
            [0, 1, 0],
            [-sin_pitch, 0, cos_pitch]
        ])
        
        # Yaw rotation around z-axis
        cos_yaw = np.cos(yaw_angle)
        sin_yaw = np.sin(yaw_angle)
        R_yaw = np.array([
            [cos_yaw, -sin_yaw, 0],
            [sin_yaw,  cos_yaw, 0],
            [0,        0,       1]
        ])
        
        # Combine rotations: R = R_yaw * R_pitch * R_roll
        # This applies roll first, then pitch, then yaw
        rotation_matrix = R_yaw @ R_pitch @ R_roll
        
        # Apply rotation to centered points
        rotated_points = (rotation_matrix @ centered_points.T).T
        
        # Create new point cloud
        rotated_pcd = o3d.geometry.PointCloud()
        rotated_pcd.points = o3d.utility.Vector3dVector(rotated_points)
        
        # Transform normals if available
        if pcd.has_normals():
            normals = np.asarray(pcd.normals)
            if len(normals) == len(points):
                # Only rotate normals (no translation for normals)
                rotated_normals = (rotation_matrix @ normals.T).T
                rotated_pcd.normals = o3d.utility.Vector3dVector(rotated_normals)
        
        # Copy colors if available
        if pcd.has_colors():
            rotated_pcd.colors = pcd.colors
        
        return rotated_pcd
    
    def find_max_points_generated_part_index(self, generated_pcds: List[o3d.geometry.PointCloud]) -> Optional[int]:
        """
        Find the index of the generated part with the most points.
        
        Args:
            generated_pcds: List of generated point clouds
        
        Returns:
            Index of the part with the most points, or None if list is empty
        """
        if not generated_pcds:
            return None
        
        max_points = -1
        max_idx = None
        
        for idx, pcd in enumerate(generated_pcds):
            points = np.asarray(pcd.points)
            num_points = len(points)
            if num_points > max_points:
                max_points = num_points
                max_idx = idx
        
        return max_idx
    
    def create_random_yaw_point_clouds(self, input_pcds: List[o3d.geometry.PointCloud], 
                                      generated_pcds: Optional[List[o3d.geometry.PointCloud]] = None) -> List[o3d.geometry.PointCloud]:
        """
        Create random rotated versions of input point clouds.
        Each part is centered at origin and has random rotations applied:
        - Yaw: random rotation around z-axis (0 to 2π)
        - Roll: random rotation around x-axis (±15 degrees)
        - Pitch: random rotation around y-axis (±15 degrees)
        
        If generated_pcds are provided, the input part corresponding to the generated part
        with the most points will NOT be rotated (only centered).
        
        Args:
            input_pcds: List of original input point clouds
            generated_pcds: Optional list of generated point clouds to determine which part to skip rotation
        
        Returns:
            List of point clouds with random rotations applied
        """
        # Find which part should not be rotated (if generated_pcds are available)
        skip_rotation_idx = None
        if generated_pcds:
            max_points_idx = self.find_max_points_generated_part_index(generated_pcds)
            if max_points_idx is not None and max_points_idx < len(input_pcds):
                skip_rotation_idx = max_points_idx
                logger.info(f"Skipping rotation for input part {skip_rotation_idx} (corresponds to generated part with most points: {len(np.asarray(generated_pcds[max_points_idx].points))} points)")
        
        random_yaw_pcds = []
        for idx, pcd in enumerate(input_pcds):
            skip_rotation = (skip_rotation_idx == idx)
            rotated_pcd = self.apply_random_yaw_rotation(pcd, skip_rotation=skip_rotation)
            random_yaw_pcds.append(rotated_pcd)
        return random_yaw_pcds
    
    def merge_input_point_clouds(self, input_pcds: List[o3d.geometry.PointCloud]) -> Optional[np.ndarray]:
        """
        Merge input point clouds and return all points as a numpy array.
        Used as fallback for center of mass calculation when GT is not available.
        
        Args:
            input_pcds: List of input point clouds
        
        Returns:
            Merged points as numpy array, or None if no points available
        """
        if not input_pcds:
            return None
        
        all_points = []
        for pcd in input_pcds:
            points = np.asarray(pcd.points)
            if len(points) > 0:
                all_points.append(points)
        
        if not all_points:
            return None
        
        merged_points = np.concatenate(all_points, axis=0)
        return merged_points
    
    def find_sequence_directories(self, input_dir: str, sequence: str) -> Tuple[Optional[str], Optional[str]]:
        """
        Find input and result sequence directories.
        
        Returns:
            Tuple of (input_sequence_dir, result_sequence_dir) or (None, None) if not found
            If result directory doesn't exist separately, result_sequence_dir may be the same as input_sequence_dir
        """
        sequence_parts = sequence.split('/')
        search_sequence = sequence_parts[-1] if sequence_parts else sequence
        
        # Find input sequence directory
        input_seq_dir = None
        candidate = os.path.join(self.input_dir, sequence)
        if os.path.exists(candidate):
            input_seq_dir = candidate
        
        if input_seq_dir is None and len(sequence_parts) > 1:
            candidate = self.input_dir
            for part in sequence_parts:
                candidate = os.path.join(candidate, part)
                if not os.path.exists(candidate):
                    candidate = None
                    break
            if candidate:
                input_seq_dir = candidate
        
        if input_seq_dir is None:
            def find_seq_dir(root_dir: str, target: str, max_depth: int = 3) -> Optional[str]:
                if max_depth <= 0:
                    return None
                if os.path.basename(root_dir) == target or target in os.path.basename(root_dir):
                    return root_dir
                if os.path.isdir(root_dir):
                    try:
                        for item in os.listdir(root_dir):
                            item_path = os.path.join(root_dir, item)
                            if os.path.isdir(item_path):
                                result = find_seq_dir(item_path, target, max_depth - 1)
                                if result:
                                    return result
                    except PermissionError:
                        pass
                return None
            
            input_seq_dir = find_seq_dir(self.input_dir, search_sequence)
        
        # Find result sequence directory
        result_seq_dir = None
        
        # First, try to find a separate result directory
        candidate = os.path.join(self.poses_dir, sequence)
        if os.path.exists(candidate):
            result_seq_dir = candidate
        
        if result_seq_dir is None and len(sequence_parts) > 1:
            candidate = self.poses_dir
            for part in sequence_parts:
                candidate = os.path.join(candidate, part)
                if not os.path.exists(candidate):
                    candidate = None
                    break
            if candidate:
                result_seq_dir = candidate
        
        if result_seq_dir is None:
            def find_result_seq_dir(root_dir: str, target: str, max_depth: int = 3) -> Optional[str]:
                if max_depth <= 0:
                    return None
                if os.path.basename(root_dir) == target or target in os.path.basename(root_dir):
                    return root_dir
                if os.path.isdir(root_dir):
                    try:
                        for item in os.listdir(root_dir):
                            item_path = os.path.join(root_dir, item)
                            if os.path.isdir(item_path):
                                result = find_result_seq_dir(item_path, target, max_depth - 1)
                                if result:
                                    return result
                    except PermissionError:
                        pass
                return None
            
            result_seq_dir = find_result_seq_dir(self.poses_dir, search_sequence)
        
        # If no separate result directory found, use input directory (for cases where transforms are alongside PLY files)
        if result_seq_dir is None and input_seq_dir is not None:
            # Check if transform files exist in the input directory
            # This handles cases like nss_ss_v1 where transforms are in the same directory as PLY files
            result_seq_dir = input_seq_dir
        
        return input_seq_dir, result_seq_dir
    
    def find_sample_directories(self, sequence_dir: str) -> List[str]:
        """
        Find all sample directories in a sequence directory.
        Samples can be:
        1. Directories starting with 'sample_' (original format)
        2. Directories containing PLY files directly (modelnet format)
        """
        sample_dirs = []
        if not os.path.exists(sequence_dir):
            return sample_dirs
        
        # First, look for sample_* directories (original format)
        for item in os.listdir(sequence_dir):
            item_path = os.path.join(sequence_dir, item)
            if os.path.isdir(item_path) and item.startswith('sample_'):
                sample_dirs.append(item_path)
        
        # If no sample_* directories found, look for directories containing PLY files (modelnet format)
        if not sample_dirs:
            for item in os.listdir(sequence_dir):
                item_path = os.path.join(sequence_dir, item)
                if os.path.isdir(item_path):
                    # Check if this directory contains PLY files
                    ply_files = glob.glob(os.path.join(item_path, "*.ply"))
                    if ply_files:
                        sample_dirs.append(item_path)
        
        return sorted(sample_dirs)
    
    def get_available_generations(self, result_sample_dir: str) -> List[str]:
        """
        Find all available generation names for a sample directory.
        Returns a list of generation names (e.g., ['generation_selected', 'generation00', 'generation01', ...])
        sorted by priority: generation_selected first, then generation00, generation01, etc.
        """
        if not os.path.exists(result_sample_dir):
            return []
        
        generations = set()
        
        # Find all transform files with generation patterns
        transform_patterns = [
            "*generation*_transform.txt",
            "*_generation*transform.txt",
        ]
        
        for pattern in transform_patterns:
            matches = glob.glob(os.path.join(result_sample_dir, pattern))
            for match in matches:
                basename = os.path.basename(match)
                # Extract generation name
                # Try to match patterns like: *generation_selected*transform.txt, *generation00*transform.txt
                gen_match = re.search(r'generation(_selected|\d+)', basename)
                if gen_match:
                    gen_name = gen_match.group(0)  # e.g., 'generation_selected' or 'generation00'
                    generations.add(gen_name)
        
        # Sort: generation_selected first, then generation00, generation01, etc.
        def sort_key(gen_name):
            if gen_name == "generation_selected":
                return (0, "")  # Highest priority
            else:
                # Extract number from generation00, generation01, etc.
                num_match = re.search(r'generation(\d+)', gen_name)
                if num_match:
                    return (1, int(num_match.group(1)))  # Sort by number
                return (2, gen_name)  # Other patterns last
        
        sorted_generations = sorted(generations, key=sort_key)
        return sorted_generations
    
    def find_matching_transform_file(self, ply_filename: str, result_sample_dir: str, generation: Optional[str] = None) -> Optional[str]:
        """
        Find matching transform file for a PLY file.
        If generation is specified, uses that generation. Otherwise, first tries to find 
        *generation_selected*.txt files, then falls back to generation00.
        """
        basename = os.path.splitext(os.path.basename(ply_filename))[0]
        part_match = None
        
        submap_match = re.search(r'submap_(\d+)', basename)
        if submap_match:
            part_num = int(submap_match.group(1))
            part_match = f"part{part_num:02d}"
        
        if part_match is None:
            part_match_obj = re.search(r'part_(\d+)', basename)
            if part_match_obj:
                part_num = int(part_match_obj.group(1))
                part_match = f"part{part_num:02d}"
        
        # Also handle simple "part(\d+)" pattern without underscore (e.g., "part1.ply")
        if part_match is None:
            part_match_obj = re.search(r'part(\d+)', basename)
            if part_match_obj:
                part_num = int(part_match_obj.group(1))
                part_match = f"part{part_num:02d}"
        
        # If generation is specified, use it directly
        if generation:
            if part_match:
                # Try multiple patterns for the specified generation
                patterns = [
                    os.path.join(result_sample_dir, f"*{generation}*{part_match}_transform.txt"),
                    os.path.join(result_sample_dir, f"*{generation}_{part_match}_transform.txt"),
                    os.path.join(result_sample_dir, f"*_{generation}_*{part_match}_transform.txt"),
                    os.path.join(result_sample_dir, f"*{part_match}_{generation}*transform.txt"),
                ]
                for pattern in patterns:
                    matches = glob.glob(pattern)
                    if matches:
                        return matches[0]
            
            sample_match = re.search(r'sample_(\d+)', basename)
            if sample_match:
                sample_id = sample_match.group(1)
                patterns = [
                    os.path.join(result_sample_dir, f"*sample{sample_id}*{generation}*transform.txt"),
                    os.path.join(result_sample_dir, f"*sample{sample_id}*_{generation}_*transform.txt"),
                ]
                for pattern in patterns:
                    matches = glob.glob(pattern)
                    if matches:
                        if part_match:
                            for match in matches:
                                if part_match in match:
                                    return match
                        if matches:
                            return matches[0]
            return None
        
        # If no generation specified, use default priority: generation_selected first, then generation00
        if part_match:
            # Try generation_selected pattern first (more flexible pattern)
            patterns_selected = [
                os.path.join(result_sample_dir, f"*generation_selected*{part_match}_transform.txt"),
                os.path.join(result_sample_dir, f"*_generation_selected_*{part_match}_transform.txt"),
                os.path.join(result_sample_dir, f"*{part_match}_generation_selected*transform.txt"),
            ]
            for pattern_selected in patterns_selected:
                matches_selected = glob.glob(pattern_selected)
                if matches_selected:
                    return matches_selected[0]
            
            # Fall back to generation00
            pattern = os.path.join(result_sample_dir, f"*generation00_{part_match}_transform.txt")
            matches = glob.glob(pattern)
            if matches:
                return matches[0]
        
        sample_match = re.search(r'sample_(\d+)', basename)
        if sample_match:
            sample_id = sample_match.group(1)
            # Try generation_selected pattern first (more flexible pattern)
            patterns_selected = [
                os.path.join(result_sample_dir, f"*sample{sample_id}*generation_selected*transform.txt"),
                os.path.join(result_sample_dir, f"*sample{sample_id}*_generation_selected_*transform.txt"),
            ]
            for pattern_selected in patterns_selected:
                matches_selected = glob.glob(pattern_selected)
                if matches_selected:
                    if part_match:
                        for match in matches_selected:
                            if part_match in match:
                                return match
                    if matches_selected:
                        return matches_selected[0]
            
            # Fall back to generation00
            pattern = os.path.join(result_sample_dir, f"*sample{sample_id}*generation00*transform.txt")
            matches = glob.glob(pattern)
            if matches:
                if part_match:
                    for match in matches:
                        if part_match in match:
                            return match
                return matches[0]
        
        return None
    
    def find_matching_transform_file_by_index(self, part_index: int, result_sample_dir: str, generation: Optional[str] = None) -> Optional[str]:
        """
        Find matching transform file by part index (for PLY files without part numbers in their names).
        
        Args:
            part_index: Zero-based index of the part (0, 1, 2, ...)
            result_sample_dir: Directory containing transform files
            generation: Optional generation name to use (e.g., 'generation00', 'generation_selected')
        
        Returns:
            Path to matching transform file, or None if not found
        """
        part_match = f"part{part_index:02d}"
        
        # If generation is specified, use it directly
        if generation:
            patterns = [
                os.path.join(result_sample_dir, f"*{generation}*{part_match}_transform.txt"),
                os.path.join(result_sample_dir, f"*{generation}_{part_match}_transform.txt"),
                os.path.join(result_sample_dir, f"*_{generation}_*{part_match}_transform.txt"),
                os.path.join(result_sample_dir, f"*{part_match}_{generation}*transform.txt"),
            ]
            for pattern in patterns:
                matches = glob.glob(pattern)
                if matches:
                    return matches[0]
            return None
        
        # If no generation specified, use default priority: generation_selected first, then generation00
        # Try generation_selected pattern first
        patterns_selected = [
            os.path.join(result_sample_dir, f"*generation_selected*{part_match}_transform.txt"),
            os.path.join(result_sample_dir, f"*_generation_selected_*{part_match}_transform.txt"),
            os.path.join(result_sample_dir, f"*{part_match}_generation_selected*transform.txt"),
        ]
        for pattern_selected in patterns_selected:
            matches_selected = glob.glob(pattern_selected)
            if matches_selected:
                return matches_selected[0]
        
        # Fall back to generation00
        pattern = os.path.join(result_sample_dir, f"*generation00*{part_match}_transform.txt")
        matches = glob.glob(pattern)
        if matches:
            return matches[0]
        
        return None
    
    def load_pose_from_file(self, pose_file: str) -> Optional[np.ndarray]:
        """Load 4x4 pose matrix from text file."""
        try:
            pose = np.loadtxt(pose_file)
            if pose.shape != (4, 4):
                if pose.shape == (16,):
                    pose = pose.reshape(4, 4)
                else:
                    logger.warning(f"Invalid pose shape {pose.shape} in {pose_file}")
                    return None
            return pose
        except Exception as e:
            logger.warning(f"Failed to load pose from {pose_file}: {e}")
            return None
    
    def load_json_metadata(self, result_sample_dir: str) -> Optional[Dict]:
        """
        Load JSON metadata file from result sample directory.
        Tries multiple patterns to find metadata files, prioritizing generation_selected.
        """
        # Priority order: generation_selected > generation00 > other generation files > metadata.json > any *.json
        
        # 1. Try generation_selected JSON files first (matching transform file priority)
        pattern_selected = os.path.join(result_sample_dir, "*generation_selected.json")
        matches_selected = glob.glob(pattern_selected)
        if matches_selected:
            try:
                with open(matches_selected[0], 'r') as f:
                    metadata = json.load(f)
                    logger.debug(f"Loaded metadata from generation_selected file: {os.path.basename(matches_selected[0])}")
                    return metadata
            except Exception as e:
                logger.debug(f"Failed to load generation_selected metadata file {matches_selected[0]}: {e}")
        
        # 2. Try generation00 JSON files
        pattern_generation00 = os.path.join(result_sample_dir, "*generation00.json")
        matches_generation00 = glob.glob(pattern_generation00)
        if matches_generation00:
            try:
                with open(matches_generation00[0], 'r') as f:
                    metadata = json.load(f)
                    logger.debug(f"Loaded metadata from generation00 file: {os.path.basename(matches_generation00[0])}")
                    return metadata
            except Exception as e:
                logger.debug(f"Failed to load generation00 metadata file {matches_generation00[0]}: {e}")
        
        # 3. Try any generation*.json files
        pattern_generation = os.path.join(result_sample_dir, "*generation*.json")
        matches_generation = glob.glob(pattern_generation)
        if matches_generation:
            # Sort to get the highest generation number
            matches_generation_sorted = sorted(matches_generation, reverse=True)
            try:
                with open(matches_generation_sorted[0], 'r') as f:
                    metadata = json.load(f)
                    logger.debug(f"Loaded metadata from generation file: {os.path.basename(matches_generation_sorted[0])}")
                    return metadata
            except Exception as e:
                logger.debug(f"Failed to load generation metadata file {matches_generation_sorted[0]}: {e}")
        
        # 4. Try exact metadata.json
        metadata_path = os.path.join(result_sample_dir, "metadata.json")
        if os.path.exists(metadata_path):
            try:
                with open(metadata_path, 'r') as f:
                    metadata = json.load(f)
                    logger.debug(f"Loaded metadata from metadata.json")
                    return metadata
            except Exception as e:
                logger.debug(f"Failed to load metadata.json: {e}")
        
        # 5. Try metadata_sample_*.json pattern
        pattern_metadata_sample = os.path.join(result_sample_dir, "metadata_sample_*.json")
        matches_metadata_sample = glob.glob(pattern_metadata_sample)
        if matches_metadata_sample:
            try:
                with open(matches_metadata_sample[0], 'r') as f:
                    metadata = json.load(f)
                    logger.debug(f"Loaded metadata from metadata_sample file: {os.path.basename(matches_metadata_sample[0])}")
                    return metadata
            except Exception as e:
                logger.debug(f"Failed to load metadata_sample file {matches_metadata_sample[0]}: {e}")
        
        # 6. Last resort: try any *.json file (but skip transform files)
        pattern_any_json = os.path.join(result_sample_dir, "*.json")
        matches_any_json = glob.glob(pattern_any_json)
        # Filter out transform-related files
        matches_any_json = [m for m in matches_any_json if "transform" not in os.path.basename(m).lower()]
        if matches_any_json:
            try:
                with open(matches_any_json[0], 'r') as f:
                    metadata = json.load(f)
                    logger.debug(f"Loaded metadata from JSON file: {os.path.basename(matches_any_json[0])}")
                    return metadata
            except Exception as e:
                logger.debug(f"Failed to load JSON file {matches_any_json[0]}: {e}")
        
        return None
    
    def print_metadata(self, metadata: Dict):
        """Print metadata in a formatted way."""
        if not metadata:
            return
        
        logger.info("=" * 60)
        logger.info("METADATA:")
        logger.info("=" * 60)
        
        def print_dict(d, indent=0):
            """Recursively print dictionary with indentation."""
            for key, value in d.items():
                if isinstance(value, dict):
                    logger.info("  " * indent + f"{key}:")
                    print_dict(value, indent + 1)
                elif isinstance(value, list):
                    if len(value) > 0 and isinstance(value[0], dict):
                        logger.info("  " * indent + f"{key}: [")
                        for i, item in enumerate(value[:5]):  # Limit to first 5 items
                            logger.info("  " * (indent + 1) + f"Item {i}:")
                            print_dict(item, indent + 2)
                        if len(value) > 5:
                            logger.info("  " * (indent + 1) + f"... ({len(value) - 5} more items)")
                        logger.info("  " * indent + "]")
                    else:
                        logger.info("  " * indent + f"{key}: {value}")
                else:
                    logger.info("  " * indent + f"{key}: {value}")
        
        print_dict(metadata)
        logger.info("=" * 60)
    
    def load_sample_point_clouds(self, input_sample_dir: str, result_sample_dir: str, return_transforms: bool = False, generation: Optional[str] = None) -> Union[Tuple[List[o3d.geometry.PointCloud], List[o3d.geometry.PointCloud], List[str]], Tuple[List[o3d.geometry.PointCloud], List[o3d.geometry.PointCloud], List[str], List[np.ndarray]]]:
        """
        Load both original and transformed point clouds for a single sample.
        
        Args:
            input_sample_dir: Directory containing input PLY files
            result_sample_dir: Directory containing transform files
            return_transforms: Whether to return transform matrices
            generation: Optional generation name to use (e.g., 'generation00', 'generation_selected')
        
        Returns:
            If return_transforms=False: Tuple of (list of original point clouds, list of registered point clouds, list of part names)
            If return_transforms=True: Tuple of (list of original point clouds, list of registered point clouds, list of part names, list of transforms)
            Returns empty lists if result directory doesn't exist or no valid point clouds found
        """
        sample_name = os.path.basename(input_sample_dir)
        
        # Handle case where result_sample_dir might be the same as input_sample_dir
        # (for datasets like nss_ss_v1 where transforms are alongside PLY files)
        if result_sample_dir == input_sample_dir:
            # Transforms are in the same directory as PLY files
            result_sample_dir = input_sample_dir
        elif not os.path.exists(result_sample_dir):
            # If separate result directory doesn't exist, check if transforms are in input directory
            # This handles cases where transforms are alongside PLY files
            if os.path.exists(input_sample_dir):
                # Check if transform files exist in input directory
                transform_files = glob.glob(os.path.join(input_sample_dir, "*transform.txt"))
                if transform_files:
                    # Transforms are in the same directory as PLY files
                    result_sample_dir = input_sample_dir
                else:
                    logger.warning(f"Result sample directory not found: {result_sample_dir}")
                    logger.warning(f"  Skipping sample {sample_name} - no registered results available")
                    if return_transforms:
                        return [], [], [], []
                    return [], [], []
        else:
            # Separate result directory exists (normal case)
            pass
        
        if not os.path.exists(input_sample_dir):
            logger.warning(f"Input sample directory not found: {input_sample_dir}")
            if return_transforms:
                return [], [], [], []
            return [], [], []
        
        # Load and print metadata
        metadata = self.load_json_metadata(result_sample_dir)
        if metadata:
            self.print_metadata(metadata)
        else:
            logger.debug(f"No metadata JSON file found in {result_sample_dir}")
        
        ply_files = natsort.natsorted(glob.glob(os.path.join(input_sample_dir, "*.ply")))
        if not ply_files:
            logger.warning(f"No PLY files found in {input_sample_dir}")
            if return_transforms:
                return [], [], [], []
            return [], [], []
        
        original_point_clouds = []
        registered_point_clouds = []
        part_names = []
        transforms_list = []  # Store transforms for potential reapplication
        part_colors = self.generate_part_colors(len(ply_files))
        skipped_count = 0
        
        for idx, ply_file in enumerate(ply_files):
            ply_filename = os.path.basename(ply_file)
            part_name = os.path.splitext(ply_filename)[0]
            
            # Try to find transform file by filename pattern first
            transform_file = self.find_matching_transform_file(ply_filename, result_sample_dir, generation=generation)
            
            # If not found by filename pattern, try matching by index (for files without part numbers in name)
            if transform_file is None:
                transform_file = self.find_matching_transform_file_by_index(idx, result_sample_dir, generation=generation)
            
            if transform_file is None:
                logger.debug(f"No transform file found for {ply_filename} (index {idx}), skipping")
                skipped_count += 1
                continue
            
            # Log which type of transform file was found
            if "generation_selected" in transform_file:
                logger.debug(f"Using generation_selected transform file: {os.path.basename(transform_file)}")
            else:
                logger.debug(f"Using generation00 transform file: {os.path.basename(transform_file)}")
            
            transform = self.load_pose_from_file(transform_file)
            if transform is None:
                skipped_count += 1
                continue
            
            try:
                pcd = o3d.io.read_point_cloud(ply_file)
                points = np.asarray(pcd.points)
                
                if len(points) == 0:
                    logger.debug(f"Empty point cloud for {part_name}, skipping")
                    skipped_count += 1
                    continue
                
                # Create original point cloud (before transformation)
                pcd_original = o3d.geometry.PointCloud()
                pcd_original.points = o3d.utility.Vector3dVector(points.copy())
                
                # Copy colors if available (before outlier removal)
                if pcd.has_colors():
                    pcd_original.colors = pcd.colors
                
                # Apply outlier removal if enabled
                pcd_original = self.apply_outlier_removal(pcd_original)
                points = np.asarray(pcd_original.points)  # Update points after outlier removal
                
                if len(points) == 0:
                    logger.debug(f"Point cloud for {part_name} is empty after outlier removal, skipping")
                    skipped_count += 1
                    continue
                
                # Add normals to original if available (after outlier removal)
                normals_orig = None
                if pcd_original.has_normals():
                    normals_orig = np.asarray(pcd_original.normals)
                    if len(normals_orig) == len(points):
                        pcd_original.normals = o3d.utility.Vector3dVector(normals_orig.copy())
                        normals_orig = normals_orig.copy()  # Keep a copy for later use
                    else:
                        normals_orig = None
                
                # Estimate normals for original if requested
                if normals_orig is None and self.estimate_normals:
                    pcd_temp = o3d.geometry.PointCloud()
                    pcd_temp.points = o3d.utility.Vector3dVector(points)
                    # Use adaptive radius based on point cloud scale
                    adaptive_radius = self.compute_adaptive_normal_radius(points)
                    pcd_temp.estimate_normals(
                        search_param=o3d.geometry.KDTreeSearchParamHybrid(
                            radius=adaptive_radius,
                            max_nn=30
                        )
                    )
                    pcd_temp.orient_normals_consistent_tangent_plane(k=15)
                    normals_orig = np.asarray(pcd_temp.normals)
                    pcd_original.normals = o3d.utility.Vector3dVector(normals_orig)
                
                # Transform points
                points_homogeneous = np.hstack([points, np.ones((len(points), 1))])
                points_transformed = (transform @ points_homogeneous.T).T[:, :3]
                
                # Create registered point cloud (after transformation)
                pcd_transformed = o3d.geometry.PointCloud()
                pcd_transformed.points = o3d.utility.Vector3dVector(points_transformed)
                
                # Transform normals if available
                normals_transformed = None
                if normals_orig is not None:
                    # Transform normals (only rotation, no translation)
                    normals_transformed = (transform[:3, :3] @ normals_orig.T).T
                    pcd_transformed.normals = o3d.utility.Vector3dVector(normals_transformed)
                
                # Estimate normals for transformed if requested and not already available
                if normals_transformed is None and self.estimate_normals:
                    pcd_temp = o3d.geometry.PointCloud()
                    pcd_temp.points = o3d.utility.Vector3dVector(points_transformed)
                    # Use adaptive radius based on point cloud scale
                    adaptive_radius = self.compute_adaptive_normal_radius(points_transformed)
                    pcd_temp.estimate_normals(
                        search_param=o3d.geometry.KDTreeSearchParamHybrid(
                            radius=adaptive_radius,
                            max_nn=30
                        )
                    )
                    pcd_temp.orient_normals_consistent_tangent_plane(k=15)
                    normals_transformed = np.asarray(pcd_temp.normals)
                    pcd_transformed.normals = o3d.utility.Vector3dVector(normals_transformed)
                
                # Set color for both
                color = part_colors[idx % len(part_colors)]
                pcd_original.paint_uniform_color(color)
                pcd_transformed.paint_uniform_color(color)
                
                original_point_clouds.append(pcd_original)
                registered_point_clouds.append(pcd_transformed)
                part_names.append(part_name)
                transforms_list.append(transform)  # Store transform for potential reapplication
            except Exception as e:
                logger.warning(f"Error processing {part_name}: {e}, skipping")
                skipped_count += 1
                continue
        
        if skipped_count > 0:
            logger.info(f"  Skipped {skipped_count} parts (missing transforms or errors)")
        
        if return_transforms:
            return original_point_clouds, registered_point_clouds, part_names, transforms_list
        else:
            return original_point_clouds, registered_point_clouds, part_names
    
    def load_generated_point_clouds(self, result_sample_dir: str, generation: Optional[str] = None) -> List[o3d.geometry.PointCloud]:
        """
        Load generated point cloud PLY files from the processed directory.
        If generation is specified, uses that generation. Otherwise, prioritizes generation_selected files, 
        then falls back to generation00.
        
        NOTE: Generated point clouds are loaded as-is without any transformation or shift.
        They are already in the center of mass frame and should not be modified.
        
        Args:
            result_sample_dir: Directory containing generated PLY files
            generation: Optional generation name to use (e.g., 'generation00', 'generation_selected')
        
        Returns:
            List of generated point clouds (already in center of mass frame, no transformation applied)
        """
        if not os.path.exists(result_sample_dir):
            return []
        
        generated_pcds = []
        
        # If generation is specified, use it directly
        if generation:
            patterns = [
                f"*{generation}_part*.ply",
                f"*{generation}*part*.ply",
            ]
        else:
            # Priority: generation_selected > generation00 > other generation files
            patterns = [
                "*generation_selected_part*.ply",
                "*generation00_part*.ply",
                "*generation*_part*.ply",
            ]
        
        for pattern in patterns:
            ply_files = sorted(glob.glob(os.path.join(result_sample_dir, pattern)))
            if ply_files:
                logger.debug(f"Found {len(ply_files)} generated PLY files matching pattern: {pattern}")
                
                # Extract part numbers and sort by part number
                ply_files_with_parts = []
                for ply_file in ply_files:
                    basename = os.path.basename(ply_file)
                    # Extract part number from filename (e.g., part00, part01, etc.)
                    part_match = re.search(r'part(\d+)', basename)
                    if part_match:
                        part_num = int(part_match.group(1))
                        ply_files_with_parts.append((part_num, ply_file))
                    else:
                        # If no part number found, use index
                        ply_files_with_parts.append((len(ply_files_with_parts), ply_file))
                
                # Sort by part number
                ply_files_with_parts.sort(key=lambda x: x[0])
                
                # Generate colors using the same scheme as original/registered
                num_parts = len(ply_files_with_parts)
                part_colors = self.generate_part_colors(num_parts)
                
                for idx, (part_num, ply_file) in enumerate(ply_files_with_parts):
                    try:
                        pcd = o3d.io.read_point_cloud(ply_file)
                        points = np.asarray(pcd.points)
                        
                        if len(points) == 0:
                            logger.debug(f"Empty generated point cloud: {os.path.basename(ply_file)}")
                            continue
                        
                        # These are already transformed, so use them as-is
                        # Add normals if available
                        if pcd.has_normals():
                            normals = np.asarray(pcd.normals)
                            if len(normals) == len(points):
                                pcd.normals = o3d.utility.Vector3dVector(normals)
                        
                        # Estimate normals if requested and not available
                        if not pcd.has_normals() and self.estimate_normals:
                            # Use adaptive radius based on point cloud scale
                            adaptive_radius = self.compute_adaptive_normal_radius(points)
                            pcd.estimate_normals(
                                search_param=o3d.geometry.KDTreeSearchParamHybrid(
                                    radius=adaptive_radius,
                                    max_nn=30
                                )
                            )
                            pcd.orient_normals_consistent_tangent_plane(k=15)
                        
                        # Use the same part-wise color scheme as original/registered
                        # Use the index in the sorted list to get the correct color
                        color = part_colors[idx % len(part_colors)]
                        pcd.paint_uniform_color(color)
                        
                        generated_pcds.append(pcd)
                    except Exception as e:
                        logger.warning(f"Error loading generated point cloud {ply_file}: {e}")
                        continue
                
                # Only use the first pattern that matches
                if generated_pcds:
                    break
        
        return generated_pcds
    
    def load_and_merge_gt_point_clouds(self, result_sample_dir: str) -> Optional[o3d.geometry.PointCloud]:
        """
        Load and merge GT (ground truth) point cloud PLY files from the processed directory.
        
        NOTE: GT point clouds are loaded and merged as-is without any transformation or shift.
        They are already in the center of mass frame and should not be modified.
        
        Returns:
            Merged GT point cloud (already in center of mass frame, no transformation applied), or None if no GT files found
        """
        if not os.path.exists(result_sample_dir):
            return None
        
        # Look for GT files
        gt_pattern = "*gt_part*.ply"
        gt_files = sorted(glob.glob(os.path.join(result_sample_dir, gt_pattern)))
        
        if not gt_files:
            logger.debug(f"No GT PLY files found matching pattern: {gt_pattern}")
            return None
        
        logger.debug(f"Found {len(gt_files)} GT PLY files")
        
        # Extract part numbers and sort by part number
        gt_files_with_parts = []
        for gt_file in gt_files:
            basename = os.path.basename(gt_file)
            # Extract part number from filename (e.g., part00, part01, etc.)
            part_match = re.search(r'part(\d+)', basename)
            if part_match:
                part_num = int(part_match.group(1))
                gt_files_with_parts.append((part_num, gt_file))
            else:
                # If no part number found, use index
                gt_files_with_parts.append((len(gt_files_with_parts), gt_file))
        
        # Sort by part number
        gt_files_with_parts.sort(key=lambda x: x[0])
        
        # Load and merge all GT point clouds
        all_points = []
        all_normals = []
        all_colors = []
        
        # Generate colors using the same scheme as original/registered
        num_parts = len(gt_files_with_parts)
        part_colors = self.generate_part_colors(num_parts)
        
        for idx, (part_num, gt_file) in enumerate(gt_files_with_parts):
            try:
                pcd = o3d.io.read_point_cloud(gt_file)
                points = np.asarray(pcd.points)
                
                if len(points) == 0:
                    logger.debug(f"Empty GT point cloud: {os.path.basename(gt_file)}")
                    continue
                
                # Get normals if available
                normals = None
                if pcd.has_normals():
                    normals_array = np.asarray(pcd.normals)
                    if len(normals_array) == len(points):
                        normals = normals_array
                
                # Estimate normals if requested and not available
                if normals is None and self.estimate_normals:
                    pcd_temp = o3d.geometry.PointCloud()
                    pcd_temp.points = o3d.utility.Vector3dVector(points)
                    # Use adaptive radius based on point cloud scale
                    adaptive_radius = self.compute_adaptive_normal_radius(points)
                    pcd_temp.estimate_normals(
                        search_param=o3d.geometry.KDTreeSearchParamHybrid(
                            radius=adaptive_radius,
                            max_nn=30
                        )
                    )
                    pcd_temp.orient_normals_consistent_tangent_plane(k=15)
                    normals = np.asarray(pcd_temp.normals)
                
                # Get colors if available, otherwise use part color
                if pcd.has_colors():
                    colors_array = np.asarray(pcd.colors)
                    if len(colors_array) == len(points):
                        colors = colors_array
                    else:
                        colors = np.tile(part_colors[idx % len(part_colors)], (len(points), 1))
                else:
                    colors = np.tile(part_colors[idx % len(part_colors)], (len(points), 1))
                
                all_points.append(points)
                if normals is not None:
                    all_normals.append(normals)
                all_colors.append(colors)
                
            except Exception as e:
                logger.warning(f"Error loading GT point cloud {gt_file}: {e}")
                continue
        
        if not all_points:
            return None
        
        # Merge all point clouds
        merged_pcd = o3d.geometry.PointCloud()
        merged_points = np.concatenate(all_points, axis=0)
        merged_pcd.points = o3d.utility.Vector3dVector(merged_points)
        
        if all_normals:
            merged_normals = np.concatenate(all_normals, axis=0)
            merged_pcd.normals = o3d.utility.Vector3dVector(merged_normals)
        
        merged_colors = np.concatenate(all_colors, axis=0)
        merged_pcd.colors = o3d.utility.Vector3dVector(merged_colors)
        
        logger.info(f"Merged {len(gt_files_with_parts)} GT parts into single point cloud with {len(merged_points)} points")
        
        return merged_pcd
    
    def generate_part_colors(self, num_parts: int) -> List[List[float]]:
        """Generate distinct colors for each part."""
        colors = []
        for i in range(num_parts):
            color_idx = i % len(CMAP_DEFAULT)
            colors.append(CMAP_DEFAULT[color_idx])
        return colors
    
    def get_sequence_samples(self, sequence: str) -> List[str]:
        """Get list of sample directory names for a sequence."""
        if sequence in self.sequence_samples_cache:
            return self.sequence_samples_cache[sequence]
        
        input_seq_dir, _ = self.find_sequence_directories(self.input_dir, sequence)
        if input_seq_dir is None:
            return []
        
        sample_dirs = self.find_sample_directories(input_seq_dir)
        sample_names = [os.path.basename(d) for d in sample_dirs]
        
        self.sequence_samples_cache[sequence] = sample_names
        return sample_names
    
    def get_available_sequences(self) -> List[str]:
        """Get list of available sequences."""
        sequences = []
        
        def has_direct_samples(dir_path: str) -> bool:
            """Check if directory directly contains sample directories (not nested)."""
            sample_dirs = self.find_sample_directories(dir_path)
            return len(sample_dirs) > 0
        
        def find_sequences_recursive(root_dir: str, current_path: str = "", max_depth: int = 5) -> List[str]:
            """Recursively find sequences (directories containing samples)."""
            if max_depth <= 0:
                return []
            
            found_sequences = []
            
            try:
                if not os.path.isdir(root_dir):
                    return []
                
                # Check if current directory contains samples
                if has_direct_samples(root_dir):
                    if current_path:
                        found_sequences.append(current_path)
                    else:
                        # Use basename if no path built yet
                        found_sequences.append(os.path.basename(root_dir))
                
                # Recursively search subdirectories
                for item in os.listdir(root_dir):
                    item_path = os.path.join(root_dir, item)
                    if os.path.isdir(item_path):
                        # Skip common non-sequence directories
                        if item in ['.', '..', '__pycache__', '.git', 'data_split']:
                            continue
                        
                        # Build path for nested structure
                        if current_path:
                            new_path = f"{current_path}/{item}"
                        else:
                            new_path = item
                        
                        # Recursively search
                        sub_sequences = find_sequences_recursive(item_path, new_path, max_depth - 1)
                        found_sequences.extend(sub_sequences)
                        
            except PermissionError:
                pass
            except Exception as e:
                logger.debug(f"Error searching {root_dir}: {e}")
            
            return found_sequences
        
        # First, check direct subdirectories for sequences (directly contain samples)
        for item in os.listdir(self.input_dir):
            item_path = os.path.join(self.input_dir, item)
            if os.path.isdir(item_path) and has_direct_samples(item_path):
                sequences.append(item)
        
        # If no sequences found, check one level deeper (for nested structure like scannet_test_v1/scannet_test_v1/)
        if not sequences:
            for item in os.listdir(self.input_dir):
                item_path = os.path.join(self.input_dir, item)
                if os.path.isdir(item_path):
                    # Check subdirectories one level deeper
                    try:
                        for subitem in os.listdir(item_path):
                            subitem_path = os.path.join(item_path, subitem)
                            if os.path.isdir(subitem_path) and has_direct_samples(subitem_path):
                                # Use nested path as sequence name
                                sequence_name = f"{item}/{subitem}" if item != subitem else subitem
                                if sequence_name not in sequences:
                                    sequences.append(sequence_name)
                    except PermissionError:
                        pass
        
        # If still no sequences found, do recursive search (for deeply nested structures like modelnet/modelnet/category/object_id)
        if not sequences:
            sequences = find_sequences_recursive(self.input_dir)
            # Remove duplicates while preserving order
            seen = set()
            unique_sequences = []
            for seq in sequences:
                if seq not in seen:
                    seen.add(seq)
                    unique_sequences.append(seq)
            sequences = unique_sequences
        
        return sorted(sequences)
    
    def find_next_sample(self) -> Tuple[Optional[str], Optional[int]]:
        """Find the next sample to load."""
        if self.current_sequence is None or self.current_sample_idx is None:
            return None, None
        
        samples = self.get_sequence_samples(self.current_sequence)
        if not samples:
            return None, None
        
        if self.current_sample_idx + 1 < len(samples):
            return self.current_sequence, self.current_sample_idx + 1
        
        # Move to next sequence
        sequences = self.get_available_sequences()
        try:
            current_seq_idx = sequences.index(self.current_sequence)
            if current_seq_idx + 1 < len(sequences):
                next_sequence = sequences[current_seq_idx + 1]
                next_samples = self.get_sequence_samples(next_sequence)
                if next_samples:
                    return next_sequence, 0
            # Loop back to first
            first_sequence = sequences[0]
            first_samples = self.get_sequence_samples(first_sequence)
            if first_samples:
                return first_sequence, 0
        except ValueError:
            pass
        
        return None, None
    
    def find_next_sequence(self) -> Tuple[Optional[str], Optional[int]]:
        """Find the first sample of the next sequence to load."""
        if self.current_sequence is None:
            return None, None
        
        sequences = self.get_available_sequences()
        if not sequences:
            return None, None
        
        try:
            current_seq_idx = sequences.index(self.current_sequence)
            if current_seq_idx + 1 < len(sequences):
                next_sequence = sequences[current_seq_idx + 1]
            else:
                # If we're at the last sequence, loop back to first
                next_sequence = sequences[0]
            
            # Get first sample of next sequence
            next_samples = self.get_sequence_samples(next_sequence)
            if next_samples:
                return next_sequence, 0
        except ValueError:
            pass
        
        return None, None
    
    def reload_current_sample_with_generation(self, generation: str) -> bool:
        """
        Reload the current sample with a specific generation.
        
        Args:
            generation: Generation name to load (e.g., 'generation00', 'generation_selected')
        
        Returns:
            True if successfully loaded, False otherwise
        """
        if self.current_sequence is None or self.current_sample_idx is None:
            return False
        
        try:
            samples = self.get_sequence_samples(self.current_sequence)
            if not samples or self.current_sample_idx >= len(samples):
                return False
            
            sample_name = samples[self.current_sample_idx]
            
            input_seq_dir, result_seq_dir = self.find_sequence_directories(self.input_dir, self.current_sequence)
            if input_seq_dir is None or result_seq_dir is None:
                return False
            
            input_sample_dir = os.path.join(input_seq_dir, sample_name)
            
            # Determine result sample directory name
            if sample_name.startswith('sample_'):
                result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
                if not os.path.exists(result_sample_dir) and result_seq_dir == input_seq_dir:
                    result_sample_dir = input_sample_dir
            else:
                result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
                if not os.path.exists(result_sample_dir):
                    result_sample_dir = os.path.join(result_seq_dir, sample_name)
                    if not os.path.exists(result_sample_dir) and result_seq_dir == input_seq_dir:
                        result_sample_dir = input_sample_dir
            
            # Load point clouds with the specified generation
            original_pcds, registered_pcds, part_names, transforms_list = self.load_sample_point_clouds(
                input_sample_dir, result_sample_dir, return_transforms=True, generation=generation
            )
            
            if not registered_pcds:
                return False
            
            # Load generated point clouds if enabled
            if hasattr(self, 'show_generated') and self.show_generated:
                self.generated_gt_pcd = self.load_and_merge_gt_point_clouds(result_sample_dir)
                self.generated_pcds = self.load_generated_point_clouds(result_sample_dir, generation=generation)
            else:
                self.generated_pcds = []
                self.generated_gt_pcd = None
            
            # Create random yaw rotated versions
            original_pcds_before_translation = []
            for pcd in original_pcds:
                pcd_copy = o3d.geometry.PointCloud()
                pcd_copy.points = o3d.utility.Vector3dVector(np.asarray(pcd.points).copy())
                if pcd.has_normals():
                    pcd_copy.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals).copy())
                if pcd.has_colors():
                    pcd_copy.colors = pcd.colors
                original_pcds_before_translation.append(pcd_copy)
            
            self.random_yaw_pcds = self.create_random_yaw_point_clouds(
                original_pcds_before_translation,
                generated_pcds=self.generated_pcds if hasattr(self, 'show_generated') and self.show_generated else None
            )
            
            # Compute center of mass and translate
            center_of_mass = None
            merged_input_points = self.merge_input_point_clouds(original_pcds)
            if merged_input_points is not None:
                _, center_of_mass = self.center_pcd(merged_input_points)
            
            if center_of_mass is not None:
                self.original_pcds = [self.translate_point_cloud(pcd, center_of_mass) for pcd in original_pcds]
                
                self.registered_pcds = []
                for idx, (pcd_original, transform) in enumerate(zip(self.original_pcds, transforms_list)):
                    points = np.asarray(pcd_original.points)
                    points_homogeneous = np.hstack([points, np.ones((len(points), 1))])
                    points_transformed = (transform @ points_homogeneous.T).T[:, :3]
                    
                    pcd_registered = o3d.geometry.PointCloud()
                    pcd_registered.points = o3d.utility.Vector3dVector(points_transformed)
                    
                    if pcd_original.has_normals():
                        normals_orig = np.asarray(pcd_original.normals)
                        normals_transformed = (transform[:3, :3] @ normals_orig.T).T
                        pcd_registered.normals = o3d.utility.Vector3dVector(normals_transformed)
                    
                    if pcd_original.has_colors():
                        pcd_registered.colors = pcd_original.colors
                    else:
                        color = self.generate_part_colors(len(self.original_pcds))[idx % len(self.original_pcds)]
                        pcd_registered.paint_uniform_color(color)
                    
                    self.registered_pcds.append(pcd_registered)
            else:
                self.original_pcds = original_pcds
                self.registered_pcds = registered_pcds
            
            # Select point clouds based on view mode
            if self.view_mode == "generated" and self.generated_pcds:
                point_clouds = self.generated_pcds.copy()
            elif self.view_mode == "gt" and self.generated_gt_pcd is not None:
                point_clouds = [self.generated_gt_pcd]
            elif self.view_mode == "random_yaw" and self.random_yaw_pcds:
                point_clouds = self.random_yaw_pcds.copy()
            elif self.view_mode == "original":
                point_clouds = self.original_pcds
            else:
                point_clouds = self.registered_pcds
            
            # Update visualization
            if self.vis is not None:
                self.vis.clear_geometries()
                
                for pcd in point_clouds:
                    self.vis.add_geometry(pcd, False)
                    self.vis.update_geometry(pcd)
                
                if not self.no_coordinate_frame:
                    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
                    self.vis.add_geometry(coordinate_frame, False)
                    self.vis.update_geometry(coordinate_frame)
                
                self._center_and_fit_view(point_clouds)
                self.vis.update_renderer()
            
            self.colored_pcds = point_clouds
            
            logger.info(f"Reloaded sample with generation: {generation}")
            return True
            
        except Exception as e:
            logger.warning(f"Error reloading sample with generation {generation}: {e}")
            return False
    
    def cycle_to_next_generation(self) -> bool:
        """
        Cycle to the next available generation for the current sample.
        If we've cycled through all generations, moves to the next sample.
        
        Returns:
            True if successfully cycled, False otherwise
        """
        if not self.cycle_generations or not self.available_generations:
            return False
        
        # Check if we're at the last generation
        if self.current_generation_idx + 1 >= len(self.available_generations):
            # Move to next sample instead
            logger.info("Cycled through all generations, moving to next sample...")
            return self.load_next_sample()
        
        # Move to next generation
        self.current_generation_idx += 1
        generation = self.available_generations[self.current_generation_idx]
        
        logger.info(f"Cycling to generation {self.current_generation_idx + 1}/{len(self.available_generations)}: {generation}")
        
        return self.reload_current_sample_with_generation(generation)
    
    def load_next_sample(self, max_attempts: int = 10):
        """Load and visualize the next sample."""
        if max_attempts <= 0:
            logger.error("Maximum attempts reached. Could not find a valid sample to load.")
            return False
        
        next_sequence, next_sample_idx = self.find_next_sample()
        
        if next_sequence is None or next_sample_idx is None:
            logger.warning("No next sample available")
            return False
        
        try:
            self.current_sequence = next_sequence
            self.current_sample_idx = next_sample_idx
            
            samples = self.get_sequence_samples(next_sequence)
            sample_name = samples[next_sample_idx]
            
            input_seq_dir, result_seq_dir = self.find_sequence_directories(self.input_dir, next_sequence)
            if input_seq_dir is None or result_seq_dir is None:
                logger.warning(f"Could not find sequence directories for {next_sequence}")
                # Try next sample
                return self.load_next_sample(max_attempts - 1)
            
            input_sample_dir = os.path.join(input_seq_dir, sample_name)
            
            # Determine result sample directory name
            # For sample_* format: use sample_*_processed
            # For other formats (like fracture_*): try both with and without _processed suffix
            # Also handle case where transforms are in the same directory as PLY files
            if sample_name.startswith('sample_'):
                result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
                # If _processed doesn't exist, check if transforms are in the same directory
                if not os.path.exists(result_sample_dir) and result_seq_dir == input_seq_dir:
                    result_sample_dir = input_sample_dir
            else:
                # Try with _processed suffix first, then without
                result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
                if not os.path.exists(result_sample_dir):
                    result_sample_dir = os.path.join(result_seq_dir, sample_name)
                    # If still doesn't exist and result_seq_dir == input_seq_dir, use input_sample_dir
                    if not os.path.exists(result_sample_dir) and result_seq_dir == input_seq_dir:
                        result_sample_dir = input_sample_dir
            
            # Initialize available generations if cycling is enabled
            if self.cycle_generations:
                self.available_generations = self.get_available_generations(result_sample_dir)
                if self.available_generations:
                    self.current_generation_idx = 0
                    logger.info(f"Found {len(self.available_generations)} available generation(s): {self.available_generations}")
                else:
                    logger.warning(f"No generations found for sample {sample_name}, disabling generation cycling")
                    self.cycle_generations = False
            
            # Load point clouds (both original and registered) and get transforms
            # Use current generation if cycling is enabled
            current_generation = None
            if self.cycle_generations and self.available_generations:
                current_generation = self.available_generations[self.current_generation_idx]
            
            original_pcds, registered_pcds, part_names, transforms_list = self.load_sample_point_clouds(
                input_sample_dir, result_sample_dir, return_transforms=True, generation=current_generation
            )
            
            if not registered_pcds:
                logger.warning(f"No point clouds loaded for sample {sample_name}")
                # Try to load next sample automatically
                logger.info("Attempting to load next available sample...")
                return self.load_next_sample(max_attempts - 1)  # Recursively try next sample with limit
            
            # Load generated point clouds if enabled
            # NOTE: Generated and GT point clouds are loaded as-is without any transformation
            # They are already in center of mass frame and will be used directly for visualization
            if hasattr(self, 'show_generated') and self.show_generated:
                # Load GT point cloud (for visualization only, not for center of mass calculation)
                self.generated_gt_pcd = self.load_and_merge_gt_point_clouds(result_sample_dir)
                if self.generated_gt_pcd is not None:
                    logger.info(f"Loaded GT point cloud (for visualization only, NOT used for center of mass)")
                
                # Load generated point clouds (no transformation applied)
                self.generated_pcds = self.load_generated_point_clouds(result_sample_dir, generation=current_generation)
                if self.generated_pcds:
                    logger.info(f"Loaded {len(self.generated_pcds)} generated point cloud(s) (no transformation applied)")
            else:
                self.generated_pcds = []
                self.generated_gt_pcd = None
            
            # Create random yaw rotated versions of input point clouds
            # Clone original point clouds BEFORE translation for random yaw view
            # (These are the untranslated original point clouds from load_sample_point_clouds)
            original_pcds_before_translation = []
            for pcd in original_pcds:
                # Create a deep copy
                pcd_copy = o3d.geometry.PointCloud()
                pcd_copy.points = o3d.utility.Vector3dVector(np.asarray(pcd.points).copy())
                if pcd.has_normals():
                    pcd_copy.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals).copy())
                if pcd.has_colors():
                    pcd_copy.colors = pcd.colors
                original_pcds_before_translation.append(pcd_copy)
            
            # Create random yaw rotated point clouds (centered at origin, random yaw per part)
            # Skip rotation for the part corresponding to the generated part with most points
            self.random_yaw_pcds = self.create_random_yaw_point_clouds(
                original_pcds_before_translation, 
                generated_pcds=self.generated_pcds if hasattr(self, 'show_generated') and self.show_generated else None
            )
            
            # Always use merged input point clouds for center of mass computation
            center_of_mass = None
            merged_input_points = self.merge_input_point_clouds(original_pcds)
            if merged_input_points is not None:
                _, center_of_mass = self.center_pcd(merged_input_points)
                logger.info(f"Using merged input point clouds center of mass: {center_of_mass}")
            
            # Translate original point clouds to center of mass frame first
            # Then re-apply transforms to get registered point clouds in center of mass frame
            # Generated and GT are already in center of mass frame, so don't translate them
            if center_of_mass is not None:
                # Translate original point clouds to center of mass frame
                self.original_pcds = [self.translate_point_cloud(pcd, center_of_mass) for pcd in original_pcds]
                
                # Re-apply transforms to translated original point clouds to get registered in center of mass frame
                self.registered_pcds = []
                for idx, (pcd_original, transform) in enumerate(zip(self.original_pcds, transforms_list)):
                    # Apply transform to the translated original point cloud
                    points = np.asarray(pcd_original.points)
                    points_homogeneous = np.hstack([points, np.ones((len(points), 1))])
                    points_transformed = (transform @ points_homogeneous.T).T[:, :3]
                    
                    pcd_registered = o3d.geometry.PointCloud()
                    pcd_registered.points = o3d.utility.Vector3dVector(points_transformed)
                    
                    # Transform normals if available
                    if pcd_original.has_normals():
                        normals_orig = np.asarray(pcd_original.normals)
                        normals_transformed = (transform[:3, :3] @ normals_orig.T).T
                        pcd_registered.normals = o3d.utility.Vector3dVector(normals_transformed)
                    
                    # Copy colors
                    if pcd_original.has_colors():
                        pcd_registered.colors = pcd_original.colors
                    else:
                        # Use same color as original
                        color = self.generate_part_colors(len(self.original_pcds))[idx % len(self.original_pcds)]
                        pcd_registered.paint_uniform_color(color)
                    
                    self.registered_pcds.append(pcd_registered)
            else:
                # If no center available, keep original registered point clouds
                self.original_pcds = original_pcds
                self.registered_pcds = registered_pcds
            
            # Select point clouds based on view mode (use translated versions)
            if self.view_mode == "generated" and self.generated_pcds:
                point_clouds = self.generated_pcds.copy()  # Only generated, no GT
            elif self.view_mode == "gt" and self.generated_gt_pcd is not None:
                point_clouds = [self.generated_gt_pcd]  # Only GT
            elif self.view_mode == "random_yaw" and self.random_yaw_pcds:
                point_clouds = self.random_yaw_pcds.copy()  # Random yaw rotated input point clouds
            elif self.view_mode == "original":
                point_clouds = self.original_pcds  # Use translated original point clouds
            else:
                point_clouds = self.registered_pcds  # Use translated registered point clouds
            
            # Clear current visualization
            if self.vis is not None:
                self.vis.clear_geometries()
            
            # Add new geometries
            if self.vis is not None:
                for pcd in point_clouds:
                    self.vis.add_geometry(pcd, False)
                    self.vis.update_geometry(pcd)
                
                if not self.no_coordinate_frame:
                    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
                    self.vis.add_geometry(coordinate_frame, False)
                    self.vis.update_geometry(coordinate_frame)
                
                self._center_and_fit_view(point_clouds)
                self.vis.update_renderer()
            
            self.colored_pcds = point_clouds
            
            logger.info("=" * 60)
            logger.info(f"LOADED NEXT SAMPLE: Sequence {next_sequence}, Sample {sample_name}")
            logger.info(f"  Parts: {len(part_names)}")
            view_name_map = {
                "original": "Original (untransformed)",
                "registered": "Registered (transformed)",
                "generated": "Generated (from processed directory)",
                "gt": "GT (ground truth)",
                "random_yaw": "Input with Random Yaw Rotation"
            }
            logger.info(f"  View: {view_name_map.get(self.view_mode, 'Unknown')}")
            logger.info("=" * 60)
            return True
            
        except FileNotFoundError as e:
            logger.warning(f"File not found: {e}")
            # Try next sample
            return self.load_next_sample(max_attempts - 1)
        except ValueError as e:
            logger.warning(f"Value error: {e}")
            # Try next sample
            return self.load_next_sample(max_attempts - 1)
        except Exception as e:
            logger.warning(f"Error loading sample: {e}")
            import traceback
            logger.debug(f"Full traceback:\n{traceback.format_exc()}")
            # Try next sample
            return self.load_next_sample(max_attempts - 1)
    
    def load_next_sequence(self, max_attempts: int = 50):
        """
        Load and visualize the first sample of the next sequence.
        Tries multiple samples in the target sequence before moving to the next sequence.
        """
        if max_attempts <= 0:
            logger.error("Maximum attempts reached. Could not find a valid sequence to load.")
            return False
        
        next_sequence, next_sample_idx = self.find_next_sequence()
        
        if next_sequence is None or next_sample_idx is None:
            logger.warning("No next sequence available")
            return False
        
        # Try multiple samples in the target sequence before giving up
        samples = self.get_sequence_samples(next_sequence)
        if not samples:
            logger.warning(f"No samples found in sequence {next_sequence}, trying next sequence...")
            # Update current sequence to skip this empty sequence
            self.current_sequence = next_sequence
            self.current_sample_idx = 0
            return self.load_next_sequence(max_attempts - 1)
        
        # Try up to all samples in the target sequence (or max_attempts, whichever is smaller)
        max_samples_to_try = min(len(samples), max_attempts)
        
        for attempt in range(max_samples_to_try):
            sample_idx = (next_sample_idx + attempt) % len(samples)
            sample_name = samples[sample_idx]
            
            try:
                input_seq_dir, result_seq_dir = self.find_sequence_directories(self.input_dir, next_sequence)
                if input_seq_dir is None or result_seq_dir is None:
                    logger.warning(f"Could not find sequence directories for {next_sequence}")
                    # Try next sequence
                    self.current_sequence = next_sequence
                    self.current_sample_idx = sample_idx
                    return self.load_next_sequence(max_attempts - max_samples_to_try)
                
                input_sample_dir = os.path.join(input_seq_dir, sample_name)
                
                # Determine result sample directory name
                # For sample_* format: use sample_*_processed
                # For other formats (like fracture_*): try both with and without _processed suffix
                if sample_name.startswith('sample_'):
                    result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
                else:
                    # Try with _processed suffix first, then without
                    result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
                    if not os.path.exists(result_sample_dir):
                        result_sample_dir = os.path.join(result_seq_dir, sample_name)
                
                # Check if result directory exists before trying to load
                if not os.path.exists(result_sample_dir):
                    logger.debug(f"Result sample directory not found: {result_sample_dir}, trying next sample...")
                    continue  # Try next sample in this sequence
                
                # Initialize available generations if cycling is enabled
                if self.cycle_generations:
                    self.available_generations = self.get_available_generations(result_sample_dir)
                    if self.available_generations:
                        self.current_generation_idx = 0
                        logger.info(f"Found {len(self.available_generations)} available generation(s): {self.available_generations}")
                    else:
                        logger.warning(f"No generations found for sample {sample_name}, disabling generation cycling")
                        self.cycle_generations = False
                
                # Load point clouds (both original and registered) and get transforms
                # Use current generation if cycling is enabled
                current_generation = None
                if self.cycle_generations and self.available_generations:
                    current_generation = self.available_generations[self.current_generation_idx]
                
                original_pcds, registered_pcds, part_names, transforms_list = self.load_sample_point_clouds(
                    input_sample_dir, result_sample_dir, return_transforms=True, generation=current_generation
                )
                
                if not registered_pcds:
                    logger.debug(f"No point clouds loaded for sample {sample_name}, trying next sample...")
                    continue  # Try next sample in this sequence
                
                # Successfully loaded point clouds - update state and proceed with visualization
                self.current_sequence = next_sequence
                self.current_sample_idx = sample_idx
                
                # Load generated point clouds if enabled
                # NOTE: Generated and GT point clouds are loaded as-is without any transformation
                # They are already in center of mass frame and will be used directly for visualization
                if hasattr(self, 'show_generated') and self.show_generated:
                    # Load GT point cloud (for visualization only, not for center of mass calculation)
                    self.generated_gt_pcd = self.load_and_merge_gt_point_clouds(result_sample_dir)
                    if self.generated_gt_pcd is not None:
                        logger.info(f"Loaded GT point cloud (for visualization only, NOT used for center of mass)")
                    
                    # Load generated point clouds (no transformation applied)
                    self.generated_pcds = self.load_generated_point_clouds(result_sample_dir, generation=current_generation)
                    if self.generated_pcds:
                        logger.info(f"Loaded {len(self.generated_pcds)} generated point cloud(s) (no transformation applied)")
                else:
                    self.generated_pcds = []
                    self.generated_gt_pcd = None
                
                # Always use merged input point clouds for center of mass computation
                center_of_mass = None
                merged_input_points = self.merge_input_point_clouds(original_pcds)
                if merged_input_points is not None:
                    _, center_of_mass = self.center_pcd(merged_input_points)
                    logger.info(f"Using merged input point clouds center of mass: {center_of_mass}")
                
                # Translate original point clouds to center of mass frame first
                # Then re-apply transforms to get registered point clouds in center of mass frame
                # Generated and GT are already in center of mass frame, so don't translate them
                if center_of_mass is not None:
                    # Translate original point clouds to center of mass frame
                    self.original_pcds = [self.translate_point_cloud(pcd, center_of_mass) for pcd in original_pcds]
                    
                    # Re-apply transforms to translated original point clouds to get registered in center of mass frame
                    self.registered_pcds = []
                    for idx, (pcd_original, transform) in enumerate(zip(self.original_pcds, transforms_list)):
                        # Apply transform to the translated original point cloud
                        points = np.asarray(pcd_original.points)
                        points_homogeneous = np.hstack([points, np.ones((len(points), 1))])
                        points_transformed = (transform @ points_homogeneous.T).T[:, :3]
                        
                        pcd_registered = o3d.geometry.PointCloud()
                        pcd_registered.points = o3d.utility.Vector3dVector(points_transformed)
                        
                        # Transform normals if available
                        if pcd_original.has_normals():
                            normals_orig = np.asarray(pcd_original.normals)
                            normals_transformed = (transform[:3, :3] @ normals_orig.T).T
                            pcd_registered.normals = o3d.utility.Vector3dVector(normals_transformed)
                        
                        # Copy colors
                        if pcd_original.has_colors():
                            pcd_registered.colors = pcd_original.colors
                        else:
                            # Use same color as original
                            color = self.generate_part_colors(len(self.original_pcds))[idx % len(self.original_pcds)]
                            pcd_registered.paint_uniform_color(color)
                        
                        self.registered_pcds.append(pcd_registered)
                else:
                    # If no center available, keep original registered point clouds
                    self.original_pcds = original_pcds
                    self.registered_pcds = registered_pcds
                
                # Create random yaw rotated versions of input point clouds
                # Clone original point clouds before translation for random yaw view
                original_pcds_before_translation = []
                for pcd in original_pcds:
                    # Create a deep copy
                    pcd_copy = o3d.geometry.PointCloud()
                    pcd_copy.points = o3d.utility.Vector3dVector(np.asarray(pcd.points).copy())
                    if pcd.has_normals():
                        pcd_copy.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals).copy())
                    if pcd.has_colors():
                        pcd_copy.colors = pcd.colors
                    original_pcds_before_translation.append(pcd_copy)
                
                # Create random yaw rotated point clouds (centered at origin, random yaw per part)
                # Skip rotation for the part corresponding to the generated part with most points
                self.random_yaw_pcds = self.create_random_yaw_point_clouds(
                    original_pcds_before_translation,
                    generated_pcds=self.generated_pcds if hasattr(self, 'show_generated') and self.show_generated else None
                )
                
                # Select point clouds based on view mode (use translated versions)
                if self.view_mode == "generated" and self.generated_pcds:
                    point_clouds = self.generated_pcds.copy()  # Only generated, no GT
                elif self.view_mode == "gt" and self.generated_gt_pcd is not None:
                    point_clouds = [self.generated_gt_pcd]  # Only GT
                elif self.view_mode == "random_yaw" and self.random_yaw_pcds:
                    point_clouds = self.random_yaw_pcds.copy()  # Random yaw rotated input point clouds
                elif self.view_mode == "original":
                    point_clouds = self.original_pcds  # Use translated original point clouds
                else:
                    point_clouds = self.registered_pcds  # Use translated registered point clouds
                
                # Clear current visualization
                if self.vis is not None:
                    self.vis.clear_geometries()
                
                # Add new geometries
                if self.vis is not None:
                    for pcd in point_clouds:
                        self.vis.add_geometry(pcd, False)
                        self.vis.update_geometry(pcd)
                    
                    if not self.no_coordinate_frame:
                        coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
                        self.vis.add_geometry(coordinate_frame, False)
                        self.vis.update_geometry(coordinate_frame)
                    
                    self._center_and_fit_view(point_clouds)
                    self.vis.update_renderer()
                
                self.colored_pcds = point_clouds
                
                logger.info("=" * 60)
                logger.info(f"JUMPED TO NEXT SEQUENCE: Sequence {next_sequence}, Sample {sample_name}")
                logger.info(f"  Parts: {len(part_names)}")
                view_name_map = {
                    "original": "Original (untransformed)",
                    "registered": "Registered (transformed)",
                    "generated": "Generated (from processed directory)",
                    "gt": "GT (ground truth)"
                }
                logger.info(f"  View: {view_name_map.get(self.view_mode, 'Unknown')}")
                logger.info("=" * 60)
                return True
                
            except Exception as e:
                logger.debug(f"Error loading sample {sample_name}: {e}")
                import traceback
                logger.debug(f"Full traceback:\n{traceback.format_exc()}")
                continue  # Try next sample in this sequence
        
        # If we've exhausted all samples in this sequence, try the next sequence
        logger.info(f"Tried {max_samples_to_try} sample(s) in sequence {next_sequence} but none had valid results. Moving to next sequence...")
        self.current_sequence = next_sequence
        self.current_sample_idx = len(samples) - 1  # Mark as having tried all samples
        return self.load_next_sequence(max_attempts - max_samples_to_try)
    
    def toggle_transform_view(self):
        """Cycle between original, registered, generated, GT, and random_yaw point cloud views."""
        # Cycle through: original -> registered -> generated -> gt -> random_yaw -> original
        if self.view_mode == "original":
            if self.registered_pcds:
                self.view_mode = "registered"
                point_clouds = self.registered_pcds
                view_name = "Registered (transformed)"
            elif self.generated_pcds:
                self.view_mode = "generated"
                point_clouds = self.generated_pcds.copy()
                view_name = "Generated (from processed directory)"
            elif self.generated_gt_pcd is not None:
                self.view_mode = "gt"
                point_clouds = [self.generated_gt_pcd]
                view_name = "GT (ground truth)"
            elif self.random_yaw_pcds:
                self.view_mode = "random_yaw"
                point_clouds = self.random_yaw_pcds.copy()
                view_name = "Input with Random Yaw Rotation"
            else:
                logger.warning("No other views available")
                return
        elif self.view_mode == "registered":
            if self.generated_pcds:
                self.view_mode = "generated"
                point_clouds = self.generated_pcds.copy()
                view_name = "Generated (from processed directory)"
            elif self.generated_gt_pcd is not None:
                self.view_mode = "gt"
                point_clouds = [self.generated_gt_pcd]
                view_name = "GT (ground truth)"
            elif self.random_yaw_pcds:
                self.view_mode = "random_yaw"
                point_clouds = self.random_yaw_pcds.copy()
                view_name = "Input with Random Yaw Rotation"
            elif self.original_pcds:
                self.view_mode = "original"
                point_clouds = self.original_pcds
                view_name = "Original (untransformed)"
            else:
                logger.warning("No other views available")
                return
        elif self.view_mode == "generated":
            if self.generated_gt_pcd is not None:
                self.view_mode = "gt"
                point_clouds = [self.generated_gt_pcd]
                view_name = "GT (ground truth)"
            elif self.random_yaw_pcds:
                self.view_mode = "random_yaw"
                point_clouds = self.random_yaw_pcds.copy()
                view_name = "Input with Random Yaw Rotation"
            elif self.original_pcds:
                self.view_mode = "original"
                point_clouds = self.original_pcds
                view_name = "Original (untransformed)"
            elif self.registered_pcds:
                self.view_mode = "registered"
                point_clouds = self.registered_pcds
                view_name = "Registered (transformed)"
            else:
                logger.warning("No other views available")
                return
        elif self.view_mode == "gt":
            if self.random_yaw_pcds:
                self.view_mode = "random_yaw"
                point_clouds = self.random_yaw_pcds.copy()
                view_name = "Input with Random Yaw Rotation"
            elif self.original_pcds:
                self.view_mode = "original"
                point_clouds = self.original_pcds
                view_name = "Original (untransformed)"
            elif self.registered_pcds:
                self.view_mode = "registered"
                point_clouds = self.registered_pcds
                view_name = "Registered (transformed)"
            elif self.generated_pcds:
                self.view_mode = "generated"
                point_clouds = self.generated_pcds.copy()
                view_name = "Generated (from processed directory)"
            else:
                logger.warning("No other views available")
                return
        else:  # random_yaw
            if self.original_pcds:
                self.view_mode = "original"
                point_clouds = self.original_pcds
                view_name = "Original (untransformed)"
            elif self.registered_pcds:
                self.view_mode = "registered"
                point_clouds = self.registered_pcds
                view_name = "Registered (transformed)"
            elif self.generated_pcds:
                self.view_mode = "generated"
                point_clouds = self.generated_pcds.copy()
                view_name = "Generated (from processed directory)"
            elif self.generated_gt_pcd is not None:
                self.view_mode = "gt"
                point_clouds = [self.generated_gt_pcd]
                view_name = "GT (ground truth)"
            else:
                logger.warning("No other views available")
                return
        
        logger.info(f"Switched to {view_name} view")
        
        # Update visualization if visualizer exists
        if self.vis is not None:
            # Clear all geometries
            self.vis.clear_geometries()
            
            # Add new point clouds
            for pcd in point_clouds:
                self.vis.add_geometry(pcd, False)
                self.vis.update_geometry(pcd)
            
            # Add coordinate frame if enabled
            if not self.no_coordinate_frame:
                coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
                self.vis.add_geometry(coordinate_frame, False)
                self.vis.update_geometry(coordinate_frame)
            
            # Don't reset view - preserve current camera position
            self.vis.update_renderer()
        
        self.colored_pcds = point_clouds
    
    def toggle_background_color(self):
        """Toggle between white and black background colors."""
        self.use_white_background = not self.use_white_background
        
        if self.use_white_background:
            new_background = self.white_background
            logger.info("Switched to white background")
        else:
            new_background = self.black_background
            logger.info("Switched to black background")
        
        self.background_color = new_background
        
        if self.vis is not None:
            render_option = self.vis.get_render_option()
            render_option.background_color = np.array(new_background)
    
    def _center_and_fit_view(self, point_clouds: List[o3d.geometry.PointCloud]):
        """Center and fit the view to show all point clouds properly."""
        if not point_clouds:
            return
        
        view_control = self.vis.get_view_control()
        
        all_points = []
        for pcd in point_clouds:
            points = np.asarray(pcd.points)
            if len(points) > 0:
                all_points.append(points)
        
        if not all_points:
            return
        
        all_points = np.concatenate(all_points, axis=0)
        center = np.mean(all_points, axis=0)
        min_coords = np.min(all_points, axis=0)
        max_coords = np.max(all_points, axis=0)
        extent = max_coords - min_coords
        max_extent = np.max(extent)
        
        front = np.array([0.0, 0.0, 1.0])
        up = np.array([0.0, 1.0, 0.0])
        lookat = center
        
        view_control.set_front(front)
        view_control.set_lookat(lookat)
        view_control.set_up(up)
        view_control.set_zoom(0.8)
    
    def visualize_sample(self, sequence: str, sample_idx: int = 0):
        """Visualize a specific sample."""
        self.current_sequence = sequence
        self.current_sample_idx = sample_idx
        
        samples = self.get_sequence_samples(sequence)
        if not samples:
            raise ValueError(f"No samples found for sequence {sequence}")
        
        if sample_idx >= len(samples):
            raise ValueError(f"Sample index {sample_idx} out of range for sequence {sequence}")
        
        sample_name = samples[sample_idx]
        
        input_seq_dir, result_seq_dir = self.find_sequence_directories(self.input_dir, sequence)
        if input_seq_dir is None or result_seq_dir is None:
            raise ValueError(f"Could not find sequence directories for {sequence}")
        
        input_sample_dir = os.path.join(input_seq_dir, sample_name)
        
        # Determine result sample directory name
        # For sample_* format: use sample_*_processed
        # For other formats (like fracture_*): try both with and without _processed suffix
        # Also handle case where transforms are in the same directory as PLY files
        if sample_name.startswith('sample_'):
            result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
            # If _processed doesn't exist, check if transforms are in the same directory
            if not os.path.exists(result_sample_dir) and result_seq_dir == input_seq_dir:
                result_sample_dir = input_sample_dir
        else:
            # Try with _processed suffix first, then without
            result_sample_dir = os.path.join(result_seq_dir, f"{sample_name}_processed")
            if not os.path.exists(result_sample_dir):
                result_sample_dir = os.path.join(result_seq_dir, sample_name)
                # If still doesn't exist and result_seq_dir == input_seq_dir, use input_sample_dir
                if not os.path.exists(result_sample_dir) and result_seq_dir == input_seq_dir:
                    result_sample_dir = input_sample_dir
        
        # Initialize available generations if cycling is enabled
        if self.cycle_generations:
            self.available_generations = self.get_available_generations(result_sample_dir)
            if self.available_generations:
                self.current_generation_idx = 0
                logger.info(f"Found {len(self.available_generations)} available generation(s): {self.available_generations}")
            else:
                logger.warning(f"No generations found for sample {sample_name}, disabling generation cycling")
                self.cycle_generations = False
        
        # Load point clouds (both original and registered) and get transforms
        # Use current generation if cycling is enabled
        current_generation = None
        if self.cycle_generations and self.available_generations:
            current_generation = self.available_generations[self.current_generation_idx]
        
        original_pcds, registered_pcds, part_names, transforms_list = self.load_sample_point_clouds(
            input_sample_dir, result_sample_dir, return_transforms=True, generation=current_generation
        )
        
        if not registered_pcds:
            logger.warning(f"No point clouds loaded for sample {sample_name}")
            logger.warning("This may be because:")
            logger.warning("  - Result sample directory doesn't exist")
            logger.warning("  - No transform files found")
            logger.warning("  - All PLY files failed to load")
            logger.info("Attempting to find next available sample...")
            # Try to find and load next sample
            next_sequence, next_sample_idx = self.find_next_sample()
            if next_sequence and next_sample_idx is not None:
                logger.info(f"Found next sample: {next_sequence}, index {next_sample_idx}")
                return self.visualize_sample(next_sequence, next_sample_idx)
            else:
                raise ValueError(f"No valid samples found. Could not load sample {sample_name} and no next sample available.")
        
        # Load generated point clouds if enabled
        # NOTE: Generated and GT point clouds are loaded as-is without any transformation
        # They are already in center of mass frame and will be used directly for visualization
        if hasattr(self, 'show_generated') and self.show_generated:
            # Load GT point cloud (for visualization only, not for center of mass calculation)
            self.generated_gt_pcd = self.load_and_merge_gt_point_clouds(result_sample_dir)
            if self.generated_gt_pcd is not None:
                logger.info(f"Loaded GT point cloud (for visualization only, NOT used for center of mass)")
            
            # Load generated point clouds (no transformation applied)
            self.generated_pcds = self.load_generated_point_clouds(result_sample_dir, generation=current_generation)
            if self.generated_pcds:
                logger.info(f"Loaded {len(self.generated_pcds)} generated point cloud(s) (no transformation applied)")
        else:
            self.generated_pcds = []
            self.generated_gt_pcd = None
        
        # Create random yaw rotated versions of input point clouds
        # Clone original point clouds BEFORE translation for random yaw view
        # (These are the untranslated original point clouds from load_sample_point_clouds)
        original_pcds_before_translation = []
        for pcd in original_pcds:
            # Create a deep copy
            pcd_copy = o3d.geometry.PointCloud()
            pcd_copy.points = o3d.utility.Vector3dVector(np.asarray(pcd.points).copy())
            if pcd.has_normals():
                pcd_copy.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals).copy())
            if pcd.has_colors():
                pcd_copy.colors = pcd.colors
            original_pcds_before_translation.append(pcd_copy)
        
        # Create random yaw rotated point clouds (centered at origin, random yaw per part)
        # Skip rotation for the part corresponding to the generated part with most points
        self.random_yaw_pcds = self.create_random_yaw_point_clouds(
            original_pcds_before_translation,
            generated_pcds=self.generated_pcds if hasattr(self, 'show_generated') and self.show_generated else None
        )
        
        # Always use merged input point clouds for center of mass computation
        center_of_mass = None
        merged_input_points = self.merge_input_point_clouds(original_pcds)
        if merged_input_points is not None:
            _, center_of_mass = self.center_pcd(merged_input_points)
            logger.info(f"Using merged input point clouds center of mass: {center_of_mass}")
        
        # Translate original point clouds to center of mass frame first
        # Then re-apply transforms to get registered point clouds in center of mass frame
        # Generated and GT are already in center of mass frame, so don't translate them
        if center_of_mass is not None:
            # Translate original point clouds to center of mass frame
            self.original_pcds = [self.translate_point_cloud(pcd, center_of_mass) for pcd in original_pcds]
            
            # Re-apply transforms to translated original point clouds to get registered in center of mass frame
            self.registered_pcds = []
            for idx, (pcd_original, transform) in enumerate(zip(self.original_pcds, transforms_list)):
                # Apply transform to the translated original point cloud
                points = np.asarray(pcd_original.points)
                points_homogeneous = np.hstack([points, np.ones((len(points), 1))])
                points_transformed = (transform @ points_homogeneous.T).T[:, :3]
                
                pcd_registered = o3d.geometry.PointCloud()
                pcd_registered.points = o3d.utility.Vector3dVector(points_transformed)
                
                # Transform normals if available
                if pcd_original.has_normals():
                    normals_orig = np.asarray(pcd_original.normals)
                    normals_transformed = (transform[:3, :3] @ normals_orig.T).T
                    pcd_registered.normals = o3d.utility.Vector3dVector(normals_transformed)
                
                # Copy colors
                if pcd_original.has_colors():
                    pcd_registered.colors = pcd_original.colors
                else:
                    # Use same color as original
                    color = self.generate_part_colors(len(self.original_pcds))[idx % len(self.original_pcds)]
                    pcd_registered.paint_uniform_color(color)
                
                self.registered_pcds.append(pcd_registered)
        else:
            # If no center available, keep original registered point clouds
            self.original_pcds = original_pcds
            self.registered_pcds = registered_pcds
        
        # Create random yaw rotated versions of input point clouds
        # Clone original point clouds before translation for random yaw view
        original_pcds_before_translation = []
        for pcd in original_pcds:
            # Create a deep copy
            pcd_copy = o3d.geometry.PointCloud()
            pcd_copy.points = o3d.utility.Vector3dVector(np.asarray(pcd.points).copy())
            if pcd.has_normals():
                pcd_copy.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals).copy())
            if pcd.has_colors():
                pcd_copy.colors = pcd.colors
            original_pcds_before_translation.append(pcd_copy)
        
        # Create random yaw rotated point clouds (centered at origin, random yaw per part)
        # Skip rotation for the part corresponding to the generated part with most points
        self.random_yaw_pcds = self.create_random_yaw_point_clouds(
            original_pcds_before_translation,
            generated_pcds=self.generated_pcds if hasattr(self, 'show_generated') and self.show_generated else None
        )
        
        # Select point clouds based on view mode (use translated versions)
        if self.view_mode == "generated" and self.generated_pcds:
            point_clouds = self.generated_pcds.copy()  # Only generated, no GT
        elif self.view_mode == "gt" and self.generated_gt_pcd is not None:
            point_clouds = [self.generated_gt_pcd]  # Only GT
        elif self.view_mode == "random_yaw" and self.random_yaw_pcds:
            point_clouds = self.random_yaw_pcds.copy()  # Random yaw rotated input point clouds
        elif self.view_mode == "original":
            point_clouds = self.original_pcds  # Use translated original point clouds
        else:
            point_clouds = self.registered_pcds  # Use translated registered point clouds
        
        logger.info("=" * 60)
        logger.info(f"SAMPLE POINT COUNT SUMMARY:")
        logger.info(f"  Sample: {sample_name}")
        logger.info(f"  Number of Parts: {len(point_clouds)}")
        total_points = sum(len(np.asarray(pcd.points)) for pcd in point_clouds)
        logger.info(f"  TOTAL POINTS: {total_points:6d}")
        logger.info("=" * 60)
        
        # Create visualizer
        try:
            self.vis = o3d.visualization.VisualizerWithKeyCallback()
            has_key_callback = True
        except AttributeError:
            self.vis = o3d.visualization.Visualizer()
            has_key_callback = False
        
        self.vis.create_window(window_name="Registered Point Cloud Visualizer")
        
        # Add point clouds
        for pcd in point_clouds:
            self.vis.add_geometry(pcd)
        
        # Add coordinate frame
        if not self.no_coordinate_frame:
            coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
            self.vis.add_geometry(coordinate_frame)
        
        # Set render options
        render_option = self.vis.get_render_option()
        render_option.background_color = np.array(self.background_color)
        render_option.point_size = self.point_size
        
        # Center and fit view
        self._center_and_fit_view(point_clouds)
        
        self.colored_pcds = point_clouds
        
        # Register keyboard callbacks
        if has_key_callback:
            try:
                def background_callback(vis):
                    self.toggle_background_color()
                    return False
                
                def next_sample_callback(vis):
                    # If generation cycling is enabled, cycle through generations instead of samples
                    if self.cycle_generations and self.available_generations:
                        success = self.cycle_to_next_generation()
                        if success:
                            logger.info(f"Spacebar pressed: Cycled to generation {self.available_generations[self.current_generation_idx]}")
                        else:
                            logger.warning("Spacebar pressed: Failed to cycle to next generation")
                    else:
                        success = self.load_next_sample()
                        if success:
                            logger.info("Spacebar pressed: Loaded next sample")
                        else:
                            logger.warning("Spacebar pressed: Failed to load next sample")
                    return False
                
                def toggle_transform_callback(vis):
                    self.toggle_transform_view()
                    return False
                
                def next_sequence_callback(vis):
                    success = self.load_next_sequence()
                    if success:
                        logger.info("N pressed: Jumped to next sequence")
                    else:
                        logger.warning("N pressed: Failed to jump to next sequence")
                    return False
                
                self.vis.register_key_callback(ord('B'), background_callback)
                self.vis.register_key_callback(ord(' '), next_sample_callback)
                self.vis.register_key_callback(ord('T'), toggle_transform_callback)
                self.vis.register_key_callback(ord('N'), next_sequence_callback)
                logger.info("Keyboard callbacks registered successfully")
            except Exception as e:
                logger.warning(f"Could not register keyboard callback: {e}")
                has_key_callback = False
        
        logger.info("Visualization ready. Press 'Q' to quit.")
        logger.info("Controls:")
        logger.info("  - Mouse: Rotate view")
        logger.info("  - Scroll: Zoom")
        logger.info("  - Shift+Mouse: Pan")
        logger.info("  - R: Reset view")
        logger.info("  - H: Print help")
        
        if has_key_callback:
            logger.info("  - B: Toggle background color")
            logger.info("  - T: Cycle between original, registered, generated, GT, and random_yaw views")
            if self.cycle_generations:
                logger.info("  - Spacebar: Cycle through generations for current sample")
            else:
                logger.info("  - Spacebar: Load next sample")
            logger.info("  - N: Jump to next sequence (first sample of next sequence)")
        
        view_name_map = {
            "original": "Original (untransformed)",
            "registered": "Registered (transformed)",
            "generated": "Generated (from processed directory)",
            "gt": "GT (ground truth)"
        }
        logger.info(f"  Current view: {view_name_map.get(self.view_mode, 'Unknown')}")
        
        logger.info("=" * 60)
        logger.info(f"CURRENT SAMPLE: Sequence {sequence}, Sample {sample_name}")
        logger.info("=" * 60)
        
        # Run visualizer
        self.vis.run()
        
        self.vis.destroy_window()
        self.vis = None


def main():
    parser = argparse.ArgumentParser(
        description="Visualize registered point clouds using estimated poses"
    )
    parser.add_argument("--input", "-i", type=str, required=True,
                       help="Input directory containing point cloud files")
    parser.add_argument("--results", "-r", type=str, required=True,
                       help="Directory containing result files (poses and generated point clouds)")
    parser.add_argument("--sequence", "-s", type=str, default=None,
                       help="Specific sequence to visualize (default: all sequences)")
    parser.add_argument("--sample_idx", type=int, default=0,
                       help="Sample index to start with (default: 0)")
    parser.add_argument("--max_points_per_fragment", type=int, default=10000,
                       help="Maximum points to load per fragment (default: 10000)")
    parser.add_argument("--generation", type=str, default="generation00",
                       help="Generation subdirectory name (default: generation00)")
    parser.add_argument("--point_size", type=float, default=3.0,
                       help="Point size in visualization (default: 3.0)")
    parser.add_argument("--background_color", type=float, nargs=3, default=[1.0, 1.0, 1.0],
                       help="Background color as R G B values in [0,1] (default: 1.0 1.0 1.0 - white)")
    parser.add_argument("--show_coordinate_frame", action='store_true',
                       help="Show coordinate frame (hidden by default)")
    parser.add_argument("--estimate_normals", "-n", action='store_true', default=False,
                       help="Estimate normals for point clouds that don't have them")
    parser.add_argument("--normal_estimation_radius", type=float, default=1.0,
                       help="Radius for normal estimation (default: 1.0)")
    parser.add_argument("--log_level", type=str, default="INFO",
                       choices=["DEBUG", "INFO", "WARNING", "ERROR"],
                       help="Set the logging level (default: INFO)")
    parser.add_argument("--show_generated", action='store_true', default=False,
                       help="Load and show generated point clouds from processed directory (generation_selected PLY files)")
    parser.add_argument("--remove_outliers", action='store_true', default=False,
                       help="Apply stochastic outlier removal to input point clouds")
    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.0,
                       help="Standard deviation ratio for outlier removal (default: 2.0)")
    parser.add_argument("--cycle_generations", action='store_true', default=False,
                       help="Cycle through different generations for each sample (press spacebar to cycle)")
    
    args = parser.parse_args()
    
    # Configure logging based on argument
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    logging.basicConfig(
        level=log_level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # Expand user paths
    input_dir = os.path.expanduser(args.input)
    results_dir = os.path.expanduser(args.results)
    
    if not os.path.exists(input_dir):
        logger.error(f"Input directory not found: {input_dir}")
        return
    
    if not os.path.exists(results_dir):
        logger.error(f"Results directory not found: {results_dir}")
        return
    
    # Create visualizer
    visualizer = RegisteredPointCloudVisualizer(
        input_dir=input_dir,
        poses_dir=results_dir,
        point_size=args.point_size,
        background_color=args.background_color,
        no_coordinate_frame=not args.show_coordinate_frame,
        estimate_normals=args.estimate_normals,
        normal_estimation_radius=args.normal_estimation_radius,
        generation=args.generation,
        max_points_per_fragment=args.max_points_per_fragment,
        remove_outliers=args.remove_outliers,
        outlier_nb_neighbors=args.outlier_nb_neighbors,
        outlier_std_ratio=args.outlier_std_ratio
    )
    
    # Set flag for showing generated point clouds
    visualizer.show_generated = args.show_generated
    
    # Set flag for cycling through generations
    visualizer.cycle_generations = args.cycle_generations
    
    # Get list of all available sequences from input directory
    available_sequences = visualizer.get_available_sequences()
    
    if not available_sequences:
        logger.error(f"No sequences found in input directory: {input_dir}")
        return
    
    logger.info(f"Found {len(available_sequences)} sequence(s) in input directory:")
    for seq in available_sequences:
        logger.info(f"  - {seq}")
    
    # If sequence not specified, randomly select one
    if args.sequence is None:
        import random
        args.sequence = random.choice(available_sequences)
        logger.info(f"No sequence specified, randomly selected: {args.sequence}")
    elif args.sequence not in available_sequences:
        logger.warning(f"Specified sequence '{args.sequence}' not found in available sequences")
        logger.warning(f"Available sequences: {available_sequences}")
        import random
        args.sequence = random.choice(available_sequences)
        logger.info(f"Randomly selected instead: {args.sequence}")
    
    # Visualize first sample
    try:
        visualizer.visualize_sample(args.sequence, args.sample_idx)
    except Exception as e:
        logger.error(f"Visualization failed: {e}")
        import traceback
        logger.error(traceback.format_exc())
    
    logger.info("Visualization complete!")


if __name__ == "__main__":
    main()