rectified_point_flow/visualizer.py

"""Visualization utilities for point cloud registration."""

from pathlib import Path
import logging
from typing import Any, Optional

import lightning as L
import torch
from lightning.pytorch.callbacks import Callback

from .utils.render import visualize_point_clouds, img_tensor_to_pil, part_ids_to_colors, probs_to_colors
from .utils.point_clouds import ppp_to_ids

import matplotlib.pyplot as plt
import numpy as np
from itertools import combinations
import pandas as pd

logger = logging.getLogger("Visualizer")


class VisualizationCallback(Callback):
    """Base Lightning callback for visualizing point clouds during evaluation."""

    def __init__(
        self,
        save_dir: Optional[str] = None,
        renderer: str = "mitsuba",
        colormap: str = "default",
        scale_to_original_size: bool = False,
        center_points: bool = False,
        image_size: int = 512,
        point_radius: float = 0.015,
        camera_dist: float = 4.0,
        camera_elev: float = 20.0,
        camera_azim: float = 30.0,
        camera_fov: float = 45.0,
        max_samples_per_batch: Optional[int] = None,
    ):
        """Initialize base visualization callback.

        Args:
            save_dir (str): Directory to save images. If None, uses trainer.log_dir/visualizations.
            renderer (str): Renderer to use, can be "mitsuba" or "pytorch3d". Default: "mitsuba".
            colormap (str): Colormap to use. Default: "default".
            scale_to_original_size (bool): If True, scales the point clouds to the original size. 
                Otherwise, keep the scaling, i.e. [-1, 1]. Default: False.
            center_points: If True, centers the point cloud around the origin. Default: False.
            image_size (int): Output image resolution (square). Default: 512.
            point_radius (float): Radius of each rendered point in world units. Default: 0.015.
            camera_dist (float): Distance (m) of camera from origin. Default: 4.0.
            camera_elev (float): Elevation angle (deg). Default: 20.0.
            camera_azim (float): Azimuth angle (deg). Default: 30.0.
            camera_fov (float): Field of view (deg). Default: 45.0.
            max_samples_per_batch (int): Maximum samples to visualize per batch. None means all.
        """
        super().__init__()
        self.save_dir = save_dir
        self.renderer = renderer
        self.colormap = colormap
        self.scale_to_original_size = scale_to_original_size
        self.max_samples_per_batch = max_samples_per_batch

        self.vis_dir = None
        self._vis_kwargs = {
            "renderer": self.renderer,
            "center_points": center_points,
            "image_size": image_size,
            "point_radius": point_radius,
            "camera_dist": camera_dist,
            "camera_elev": camera_elev,
            "camera_azim": camera_azim,
            "camera_fov": camera_fov,
        }

    def setup(self, trainer: L.Trainer, pl_module: L.LightningModule, stage: str) -> None:
        if stage == "test" and self.renderer != "none":
            if self.save_dir is None:
                self.vis_dir = Path(trainer.log_dir) / "visualizations"
            else:
                self.vis_dir = Path(self.save_dir) / "visualizations"
            self.vis_dir.mkdir(parents=True, exist_ok=True)

    def _save_sample_images(
        self,
        points: torch.Tensor,
        colors: torch.Tensor,
        sample_name: str,
        subfolder: Optional[str] = None,
    ):
        """Save visualization images for a single sample.
        
        Args:
            points (torch.Tensor): Point cloud of shape (N, 3).
            colors (torch.Tensor): Colors of shape (N, 3).
            sample_name (str): sample name for filename.
            subfolder (str | None): Optional subfolder under visualization directory.
        """
        # Skip saving if renderer is set to "none"
        if self.renderer == "none":
            return
            
        try:
            image = visualize_point_clouds(
                points=points,
                colors=colors,
                **self._vis_kwargs
            )
            image_pil = img_tensor_to_pil(image)
            sample_name = sample_name.replace('/', '_')
            target_dir = self.vis_dir if subfolder is None else (self.vis_dir / str(subfolder))
            target_dir.mkdir(parents=True, exist_ok=True)
            image_pil.save(target_dir / f"{sample_name}.png")
        except Exception as e:
            logger.error(f"Error saving visualization for sample {sample_name}: {e}")

    def on_test_batch_end(
        self,
        trainer: L.Trainer,
        pl_module: L.LightningModule,
        outputs: Any,
        batch: dict[str, torch.Tensor],
        batch_idx: int,
        dataloader_idx: int = 0,
    ):
        """Override this method in subclasses for specific visualization logic."""
        raise NotImplementedError("Subclasses must implement on_test_batch_end")


class FlowVisualizationCallback(VisualizationCallback):
    """Visualization callback for rectified point flow models."""

    def __init__(
        self,
        save_trajectory: bool = True,
        trajectory_gif_fps: int = 25,
        trajectory_gif_pause_last_frame: float = 1.0,
        filter_failures: bool = False,
        failure_metric: Optional[str] = None,
        failure_generation_idx: int = 0,
        folder_suffix: Optional[str] = None,
        save_individual_parts: bool = False,
        colorize_input_parts_with_pca: bool = False,
        colorize_generated_parts_with_pca: bool = False,
        **kwargs
    ):
        """Initialize flow visualization callback.

        Args:
            save_trajectory (bool): Whether to save trajectory as GIF. Default: True.
            trajectory_gif_fps (int): Frames per second for the GIF.
            trajectory_gif_pause_last_frame (float): Pause time for the last frame in seconds.
            filter_failures (bool): If True, only visualize samples that fail the specified metric.
                Default: False (visualize all samples).
            failure_metric (str | None): Metric name to check for failures. For recall metrics,
                a failure means the value is 0 (or < threshold). For error metrics, a failure
                means the value exceeds the threshold. Examples: "recall_at_5deg_2m (outdoor_bufferx)".
                If None and filter_failures=True, uses "recall_at_5deg_2m (outdoor_bufferx)".
            failure_generation_idx (int): Which generation index to use for failure checking.
                Default: 0 (first generation).
            folder_suffix (str | None): Optional suffix to append to dataset name in folder path.
                If provided, folders will be named like "{dataset_name}_{folder_suffix}".
                Default: None (use dataset name only).
            save_individual_parts (bool): If True, save individual visualizations for each part
                colored by part ID. Default: False.
            colorize_input_parts_with_pca (bool): If True and save_individual_parts=True, colorize
                input parts using PCA of point-wise local features from batch["features"].
                Requires features to be available in the batch. Default: False.
            colorize_generated_parts_with_pca (bool): If True and save_individual_parts=True, colorize
                generated parts using PCA of transformer features (after transformer, before MLP head).
                Requires transformer_features to be available in outputs. Default: False.
            **kwargs: Additional arguments passed to base class.
        """
        super().__init__(**kwargs)
        self.save_trajectory = save_trajectory
        self.trajectory_gif_fps = trajectory_gif_fps
        self.trajectory_gif_pause_last_frame = trajectory_gif_pause_last_frame
        self.filter_failures = filter_failures
        self.failure_metric = failure_metric or "recall_at_5deg_2m (outdoor_bufferx)"
        self.failure_generation_idx = failure_generation_idx
        self.folder_suffix = folder_suffix
        self.save_individual_parts = save_individual_parts
        self.colorize_input_parts_with_pca = colorize_input_parts_with_pca
        self.colorize_generated_parts_with_pca = colorize_generated_parts_with_pca
        
        # PCA basis storage (computed from first batch)
        self.input_features_pca_basis = None  # Dict with 'mean', 'eigenvecs', 'pca_min', 'pca_max'
        self.transformer_features_pca_basis = None  # Dict with 'mean', 'eigenvecs', 'pca_min', 'pca_max'
        self.pca_basis_computed = False  # Track if PCA basis has been computed

    def _compute_pca_basis(self, features: torch.Tensor) -> dict:
        """Compute PCA basis from features.
        
        Args:
            features: Point-wise features of shape (N, F) where F is feature dimension.
            
        Returns:
            Dict with 'mean', 'eigenvecs', 'pca_min', 'pca_max' for PCA transformation.
        """
        if features.shape[0] < 2:
            return None
        
        # Center the features
        features_mean = features.mean(dim=0, keepdim=True)  # (1, F)
        features_centered = features - features_mean
        
        # Compute covariance matrix
        cov = torch.mm(features_centered.t(), features_centered) / (features.shape[0] - 1)
        
        # Compute eigenvalues and eigenvectors
        try:
            eigenvals, eigenvecs = torch.linalg.eigh(cov)
            # Sort in descending order
            idx = eigenvals.argsort(descending=True)
            eigenvecs = eigenvecs[:, idx]
        except Exception as e:
            logger.warning(f"PCA basis computation failed: {e}.")
            return None
        
        # Project features onto first 3 principal components
        pca_features = torch.mm(features_centered, eigenvecs[:, :3])  # (N, 3)
        
        # Compute normalization parameters
        pca_min = pca_features.min(dim=0, keepdim=True)[0]  # (1, 3)
        pca_max = pca_features.max(dim=0, keepdim=True)[0]  # (1, 3)
        
        return {
            'mean': features_mean,
            'eigenvecs': eigenvecs[:, :3],  # (F, 3)
            'pca_min': pca_min,
            'pca_max': pca_max,
        }

    def _features_to_pca_colors(
        self, 
        features: torch.Tensor, 
        pca_basis: Optional[dict] = None
    ) -> torch.Tensor:
        """Convert point-wise features to RGB colors using PCA.
        
        Args:
            features: Point-wise features of shape (N, F) where F is feature dimension.
            pca_basis: Optional pre-computed PCA basis dict with 'mean', 'eigenvecs', 'pca_min', 'pca_max'.
                      If None, computes PCA per-part (old behavior).
            
        Returns:
            RGB colors of shape (N, 3) with values in [0, 1].
        """
        if features.shape[0] < 2:
            # Not enough points for PCA, return uniform color
            return torch.ones(features.shape[0], 3, device=features.device) * 0.5
        
        if pca_basis is not None:
            # Use pre-computed PCA basis
            features_mean = pca_basis['mean']  # (1, F)
            eigenvecs = pca_basis['eigenvecs']  # (F, 3)
            pca_min = pca_basis['pca_min']  # (1, 3)
            pca_max = pca_basis['pca_max']  # (1, 3)
            
            # Center the features using stored mean
            features_centered = features - features_mean
            
            # Project features onto first 3 principal components
            pca_features = torch.mm(features_centered, eigenvecs)  # (N, 3)
        else:
            # Compute PCA per-part (old behavior)
            # Center the features
            features_centered = features - features.mean(dim=0, keepdim=True)
            
            # Compute covariance matrix
            cov = torch.mm(features_centered.t(), features_centered) / (features.shape[0] - 1)
            
            # Compute eigenvalues and eigenvectors
            try:
                eigenvals, eigenvecs = torch.linalg.eigh(cov)
                # Sort in descending order
                idx = eigenvals.argsort(descending=True)
                eigenvecs = eigenvecs[:, idx]
            except Exception as e:
                logger.warning(f"PCA computation failed: {e}. Using uniform colors.")
                return torch.ones(features.shape[0], 3, device=features.device) * 0.5
            
            # Project features onto first 3 principal components
            pca_features = torch.mm(features_centered, eigenvecs[:, :3])  # (N, 3)
            
            # Normalize to [0, 1] range for RGB
            pca_min = pca_features.min(dim=0, keepdim=True)[0]
            pca_max = pca_features.max(dim=0, keepdim=True)[0]
        
        # Normalize to [0, 1] range for RGB
        pca_range = pca_max - pca_min
        
        # Avoid division by zero
        pca_range = torch.where(pca_range < 1e-6, torch.ones_like(pca_range), pca_range)
        
        colors_normalized = (pca_features - pca_min) / pca_range
        
        # Clamp to [0, 1] to ensure valid RGB values
        colors_normalized = torch.clamp(colors_normalized, 0.0, 1.0)
        
        return colors_normalized

    def _save_trajectory_gif(
        self,
        trajectory: torch.Tensor,
        colors: torch.Tensor,
        sample_name: str,
        subfolder: Optional[str] = None,
    ):
        """Save trajectory as GIF.
        
        Args:
            trajectory: Point clouds representing the trajectory steps of shape (num_steps, N, 3).
            colors: Colors of shape (N, 3). Same for all trajectory steps.
            sample_name (str): sample name for filename.
            subfolder (str | None): Optional subfolder under visualization directory.
        """
        # Skip saving if renderer is set to "none"
        if self.renderer == "none":
            return
            
        try:
            target_dir = self.vis_dir if subfolder is None else (self.vis_dir / str(subfolder))
            target_dir.mkdir(parents=True, exist_ok=True)
            gif_path = target_dir / f"{sample_name}.gif"

            # Render trajectory steps
            rendered_images = visualize_point_clouds(
                points=trajectory,                                          # (num_steps, N, 3)
                colors=colors,                                              # (N, 3)
                **self._vis_kwargs,
            )                                                               # (num_steps, H, W, 3)
            frames = []
            num_steps = trajectory.shape[0]
            for step in range(num_steps):
                frame_pil = img_tensor_to_pil(rendered_images[step])        # (H, W, 3)
                frames.append(frame_pil)
            
            # Frame duration and pause on last frame in ms
            duration = int(1000 / self.trajectory_gif_fps)
            durations = [duration] * len(frames)
            durations[-1] = int(duration + self.trajectory_gif_pause_last_frame * 1000)

            frames[0].save(
                gif_path,
                save_all=True,
                append_images=frames[1:],
                duration=durations,
                loop=0,  # Infinite loop
                optimize=True
            )
        except Exception as e:
            logger.error(f"Error saving trajectory GIF for sample {sample_name}: {e}")

    def on_test_batch_end(
        self,
        trainer: L.Trainer,
        pl_module: L.LightningModule,
        outputs: Any,
        batch: dict[str, torch.Tensor],
        batch_idx: int,
        dataloader_idx: int = 0,
    ):
        """Save flow visualizations at the end of each test batch."""
        if self.vis_dir is None:
            return

        points_per_part = batch["points_per_part"]                            # (bs, max_parts)
        B, _ = points_per_part.shape
        part_ids = ppp_to_ids(points_per_part)                                # (bs, N)
        
        # Handle dynamic batching - check if we have cu_seqlens_batch
        if "cu_seqlens_batch" in batch:
            # Dynamic batching format (TP, 3)
            cu_seqlens_batch = batch["cu_seqlens_batch"]                      # (B + 1,)
            pts_flat = batch["pointclouds"]                                   # (TP, 3)
            pts_gt_flat = batch["pointclouds_gt"]                            # (TP, 3)
            
            # Convert from dynamic batching to fixed batching
            pts_batches = []
            pts_gt_batches = []
            
            for b in range(B):
                start_idx = cu_seqlens_batch[b]
                end_idx = cu_seqlens_batch[b + 1]
                pts_batches.append(pts_flat[start_idx:end_idx])
                pts_gt_batches.append(pts_gt_flat[start_idx:end_idx])
            
            # Find max points for padding
            max_points = max(batch.shape[0] for batch in pts_batches)
            
            # Create padded tensors
            device = pts_flat.device
            pts = torch.zeros(B, max_points, 3, device=device, dtype=pts_flat.dtype)
            pts_gt = torch.zeros(B, max_points, 3, device=device, dtype=pts_gt_flat.dtype)
            
            for b, (pts_batch, pts_gt_batch) in enumerate(zip(pts_batches, pts_gt_batches)):
                pts[b, :pts_batch.shape[0]] = pts_batch
                pts_gt[b, :pts_gt_batch.shape[0]] = pts_gt_batch
        else:
            # Fixed batching format (B, N, 3)
            pts = batch["pointclouds"].view(B, -1, 3)                         # (bs, N, 3)
            pts_gt = batch["pointclouds_gt"].view(B, -1, 3)                   # (bs, N, 3)
        
        # K generations
        trajectories_list = outputs['trajectories']                           # (K, num_steps, num_points, 3)
        K = len(trajectories_list)
        
        # Handle trajectory reshaping for dynamic batching
        if "cu_seqlens_batch" in batch:
            # For dynamic batching, trajectories are in (num_steps, TP, 3) format
            pointclouds_pred_list = []
            for traj in trajectories_list:
                # Convert each trajectory from (num_steps, TP, 3) to (B, N, 3) for the final step
                final_step = traj[-1]  # (TP, 3)
                
                # Split into batches like we did above
                pred_batches = []
                for b in range(B):
                    start_idx = cu_seqlens_batch[b]
                    end_idx = cu_seqlens_batch[b + 1]
                    pred_batches.append(final_step[start_idx:end_idx])
                
                # Pad to max_points
                pred_padded = torch.zeros(B, max_points, 3, device=final_step.device, dtype=final_step.dtype)
                for b, pred_batch in enumerate(pred_batches):
                    pred_padded[b, :pred_batch.shape[0]] = pred_batch
                
                pointclouds_pred_list.append(pred_padded)
        else:
            # Fixed batching
            pointclouds_pred_list = [traj[-1].view(B, -1, 3) for traj in trajectories_list]

        if self.scale_to_original_size:
            scale = batch["scales"]                                           # (B,)
            pts = pts * scale[:, None, None]                                  # (bs, N, 3)
            pts_gt = pts_gt * scale[:, None, None]                           # (bs, N, 3)
            pointclouds_pred_list = [pred * scale[:, None, None] for pred in pointclouds_pred_list]

        # Extract features if available and PCA coloring is enabled
        features_batch = None
        if self.colorize_input_parts_with_pca and "features" in batch and batch["features"] is not None:
            features_raw = batch["features"]  # (B, N, F) or (TP, F) for dynamic batching
            if "cu_seqlens_batch" in batch:
                # Dynamic batching: features are in (TP, F) format
                # Reuse max_points computed earlier
                feature_batches = []
                for b in range(B):
                    start_idx = cu_seqlens_batch[b]
                    end_idx = cu_seqlens_batch[b + 1]
                    feature_batches.append(features_raw[start_idx:end_idx])
                
                # Pad to max_points (already computed above)
                F = features_raw.shape[-1]
                device = features_raw.device
                features_batch = torch.zeros(B, max_points, F, device=device, dtype=features_raw.dtype)
                
                for b, feat_batch in enumerate(feature_batches):
                    features_batch[b, :feat_batch.shape[0]] = feat_batch
            else:
                # Fixed batching: features are in (B, N, F) format
                features_batch = features_raw.view(B, -1, features_raw.shape[-1])

        # Extract transformer features if available and PCA coloring for generated parts is enabled
        transformer_features_list = None
        if self.colorize_generated_parts_with_pca and 'transformer_features' in outputs:
            transformer_features_list = outputs['transformer_features']  # List of (TP, embed_dim), one per generation

        # Compute PCA basis from first batch if not already computed
        if not self.pca_basis_computed:
            # Compute PCA basis for input features
            if self.colorize_input_parts_with_pca and features_batch is not None:
                # Collect all valid features from the batch
                all_input_features = []
                for b in range(B):
                    valid_n = int(points_per_part[b].sum().item())
                    if valid_n > 0:
                        all_input_features.append(features_batch[b][:valid_n])
                
                if all_input_features:
                    all_input_features = torch.cat(all_input_features, dim=0)  # (total_points, F)
                    self.input_features_pca_basis = self._compute_pca_basis(all_input_features)
                    if self.input_features_pca_basis is not None:
                        logger.info(f"Computed input features PCA basis from first batch: {all_input_features.shape[0]} points")
            
            # Compute PCA basis for transformer features
            if self.colorize_generated_parts_with_pca and transformer_features_list is not None and len(transformer_features_list) > 0:
                # Use first generation's transformer features
                transformer_features_first = transformer_features_list[0]  # (TP, embed_dim)
                
                if "cu_seqlens_batch" in batch:
                    # Dynamic batching: collect all features
                    all_transformer_features = transformer_features_first  # Already (TP, embed_dim)
                else:
                    # Fixed batching: collect all features
                    all_transformer_features = transformer_features_first.view(-1, transformer_features_first.shape[-1])  # (B*N, embed_dim)
                
                self.transformer_features_pca_basis = self._compute_pca_basis(all_transformer_features)
                if self.transformer_features_pca_basis is not None:
                    logger.info(f"Computed transformer features PCA basis from first batch: {all_transformer_features.shape[0]} points")
            
            self.pca_basis_computed = True

        # Get failure mask if filtering is enabled
        failure_mask = None
        if self.filter_failures and 'eval_results' in outputs:
            eval_results_list = outputs.get('eval_results', [])
            if eval_results_list and len(eval_results_list) > self.failure_generation_idx:
                eval_results = eval_results_list[self.failure_generation_idx]
                if self.failure_metric in eval_results:
                    metric_values = eval_results[self.failure_metric]  # (B,)
                    # For recall metrics: failure means value == 0 (or < 1.0)
                    # For error metrics: failure means value > threshold (but we'll handle recalls here)
                    if 'recall' in self.failure_metric.lower():
                        failure_mask = metric_values < 1.0  # Recall < 1.0 means failure
                    else:
                        # For error metrics, we'd need a threshold, but for now assume 0 means success
                        failure_mask = metric_values > 0
                else:
                    logger.warning(f"Metric '{self.failure_metric}' not found in eval_results. Available metrics: {list(eval_results.keys())}")

        # If filtering failures, count how many failures we've visualized
        failures_visualized = 0
        
        for i in range(B):
            # Check if we've reached the maximum number of samples to visualize
            # If filtering failures, count actual failures visualized; otherwise count samples checked
            if self.max_samples_per_batch is not None:
                if self.filter_failures and failure_mask is not None:
                    # When filtering failures, limit by number of failures visualized
                    if failures_visualized >= self.max_samples_per_batch:
                        break
                else:
                    # When not filtering, limit by number of samples checked
                    if i >= self.max_samples_per_batch:
                        break
            
            # Check if we should skip this sample based on failure filtering
            if failure_mask is not None and not failure_mask[i].item():
                continue  # Skip visualization for non-failure cases
            
            # Increment failure counter if we're visualizing a failure
            if failure_mask is not None:
                failures_visualized += 1
                
            dataset_name = str(batch["dataset_name"][i])
            sample_name = f"{dataset_name}_sample{int(batch['index'][i]):05d}"
            
            # Construct subfolder name with optional suffix
            subfolder_name = dataset_name
            if self.folder_suffix is not None:
                subfolder_name = f"{dataset_name}_{self.folder_suffix}"

            colors = part_ids_to_colors(
                part_ids[i], colormap=self.colormap, part_order="id"
            )
            # Trim padded zeros if any
            valid_n = int(points_per_part[i].sum().item())
            pts_i = pts[i][:valid_n]
            pts_gt_i = pts_gt[i][:valid_n]
            part_ids_i = part_ids[i][:valid_n]
            colors_i = colors[:valid_n]
            self._save_sample_images(
                points=pts_i,
                colors=colors_i,
                sample_name=f"{sample_name}_input",
                subfolder=subfolder_name,
            )
            self._save_sample_images(
                points=pts_gt_i,
                colors=colors_i,
                sample_name=f"{sample_name}_gt",
                subfolder=subfolder_name,
            )

            for n in range(K):
                pointclouds_pred = pointclouds_pred_list[n]
                pred_i = pointclouds_pred[i][:valid_n]
                self._save_sample_images(
                    points=pred_i,
                    colors=colors_i,
                    sample_name=f"{sample_name}_generation{n+1:02d}",
                    subfolder=subfolder_name,
                )

                if self.save_trajectory:
                    # Save end_point_trajectory (primary trajectory)
                    trajectory = trajectories_list[n]  # (num_steps, TP, 3) or (num_steps, B*N, 3)
                    
                    if "cu_seqlens_batch" in batch:
                        # Dynamic batching: convert (num_steps, TP, 3) to (num_steps, B, N, 3)
                        num_steps = trajectory.shape[0]
                        trajectory_batched = torch.zeros(num_steps, B, max_points, 3, 
                                                       device=trajectory.device, dtype=trajectory.dtype)
                        
                        for step in range(num_steps):
                            step_points = trajectory[step]  # (TP, 3)
                            for b in range(B):
                                start_idx = cu_seqlens_batch[b]
                                end_idx = cu_seqlens_batch[b + 1]
                                batch_points = step_points[start_idx:end_idx]
                                trajectory_batched[step, b, :batch_points.shape[0]] = batch_points
                        
                        trajectory_for_sample = trajectory_batched[:, i]  # (num_steps, N, 3)
                    else:
                        # Fixed batching: reshape and extract
                        num_steps = trajectory.shape[0]
                        trajectory = trajectory.reshape(num_steps, B, -1, 3)  # (num_steps, B, N, 3)
                        trajectory_for_sample = trajectory[:, i]  # (num_steps, N, 3)
                    
                    if self.scale_to_original_size:
                        trajectory_for_sample = trajectory_for_sample * scale[i, None, None]
                    
                    self._save_trajectory_gif(
                        trajectory=trajectory_for_sample,
                        colors=colors,
                        sample_name=f"{sample_name}_generation{n+1:02d}_end_point",
                        subfolder=subfolder_name,
                    )
                    
                    # Save original trajectory (x_t) if available
                    if 'trajectories_x_t' in outputs and outputs['trajectories_x_t'][n] is not None:
                        trajectory_x_t = outputs['trajectories_x_t'][n]  # (num_steps, TP, 3) or (num_steps, B*N, 3)
                        
                        if "cu_seqlens_batch" in batch:
                            # Dynamic batching: convert (num_steps, TP, 3) to (num_steps, B, N, 3)
                            num_steps = trajectory_x_t.shape[0]
                            trajectory_batched = torch.zeros(num_steps, B, max_points, 3, 
                                                           device=trajectory_x_t.device, dtype=trajectory_x_t.dtype)
                            
                            for step in range(num_steps):
                                step_points = trajectory_x_t[step]  # (TP, 3)
                                for b in range(B):
                                    start_idx = cu_seqlens_batch[b]
                                    end_idx = cu_seqlens_batch[b + 1]
                                    batch_points = step_points[start_idx:end_idx]
                                    trajectory_batched[step, b, :batch_points.shape[0]] = batch_points
                            
                            trajectory_x_t_for_sample = trajectory_batched[:, i]  # (num_steps, N, 3)
                        else:
                            # Fixed batching: reshape and extract
                            num_steps = trajectory_x_t.shape[0]
                            trajectory_x_t = trajectory_x_t.reshape(num_steps, B, -1, 3)  # (num_steps, B, N, 3)
                            trajectory_x_t_for_sample = trajectory_x_t[:, i]  # (num_steps, N, 3)
                        
                        if self.scale_to_original_size:
                            trajectory_x_t_for_sample = trajectory_x_t_for_sample * scale[i, None, None]
                        
                        self._save_trajectory_gif(
                            trajectory=trajectory_x_t_for_sample,
                            colors=colors,
                            sample_name=f"{sample_name}_generation{n+1:02d}_original",
                            subfolder=subfolder_name,
                        )
            
            # Save individual parts visualization if enabled
            if self.save_individual_parts:
                # Get unique part IDs for this sample
                unique_part_ids = torch.unique(part_ids_i)
                
                # Extract features for this sample if PCA coloring is enabled
                features_i = None
                if self.colorize_input_parts_with_pca and features_batch is not None:
                    features_i = features_batch[i][:valid_n]  # (valid_n, F)
                
                for part_id in unique_part_ids:
                    part_id_int = int(part_id.item())
                    # Create mask for this part
                    part_mask = part_ids_i == part_id
                    
                    if part_mask.sum() == 0:
                        continue  # Skip if no points for this part
                    
                    # Extract points for this part
                    part_points = pts_i[part_mask]
                    
                    # For generations and GT: always use part ID colors
                    part_color_id = colors_i[part_mask][0:1]  # Get first color (all same)
                    part_colors_id = part_color_id.expand(part_points.shape[0], -1)
                    
                    # Always save visualization with uniform ID colors for input
                    self._save_sample_images(
                        points=part_points,
                        colors=part_colors_id,
                        sample_name=f"{sample_name}_input_part{part_id_int:02d}",
                        subfolder=subfolder_name,
                    )
                    
                    # Additionally save PCA-colored version if enabled
                    if self.colorize_input_parts_with_pca and features_i is not None:
                        # Use PCA-based coloring from features
                        part_features = features_i[part_mask]  # (num_points_in_part, F)
                        part_colors_input_pca = self._features_to_pca_colors(
                            part_features, 
                            pca_basis=self.input_features_pca_basis
                        )
                        self._save_sample_images(
                            points=part_points,
                            colors=part_colors_input_pca,
                            sample_name=f"{sample_name}_input_part{part_id_int:02d}_pca",
                            subfolder=subfolder_name,
                        )
                    
                    # # Also save GT for this part
                    # part_points_gt = pts_gt_i[part_mask]
                    # self._save_sample_images(
                    #     points=part_points_gt,
                    #     colors=part_colors_id,
                    #     sample_name=f"{sample_name}_gt_part{part_id_int:02d}",
                    #     subfolder=subfolder_name,
                    # )
                    
                    # Save each generation for this part
                    for n in range(K):
                        pointclouds_pred = pointclouds_pred_list[n]
                        pred_i = pointclouds_pred[i][:valid_n]
                        part_points_pred = pred_i[part_mask]
                        
                        # Always save visualization with uniform ID colors for generated part
                        self._save_sample_images(
                            points=part_points_pred,
                            colors=part_colors_id,
                            sample_name=f"{sample_name}_generation{n+1:02d}_part{part_id_int:02d}",
                            subfolder=subfolder_name,
                        )
                        
                        # Additionally save PCA-colored version if enabled
                        if self.colorize_generated_parts_with_pca and transformer_features_list is not None and n < len(transformer_features_list):
                            # Extract transformer features for this generation and sample
                            transformer_features_n = transformer_features_list[n]  # (TP, embed_dim)
                            
                            if "cu_seqlens_batch" in batch:
                                # Dynamic batching: split transformer features by batch
                                start_idx = cu_seqlens_batch[i]
                                end_idx = cu_seqlens_batch[i + 1]
                                transformer_features_i = transformer_features_n[start_idx:end_idx]  # (N_i, embed_dim)
                            else:
                                # Fixed batching: reshape and extract
                                transformer_features_i = transformer_features_n.view(B, -1, transformer_features_n.shape[-1])[i][:valid_n]  # (valid_n, embed_dim)
                            
                            # Extract features for this part
                            part_transformer_features = transformer_features_i[part_mask]  # (num_points_in_part, embed_dim)
                            
                            if part_transformer_features.shape[0] > 0:
                                # Use PCA-based coloring from transformer features
                                part_colors_gen_pca = self._features_to_pca_colors(
                                    part_transformer_features,
                                    pca_basis=self.transformer_features_pca_basis
                                )
                                self._save_sample_images(
                                    points=part_points_pred,
                                    colors=part_colors_gen_pca,
                                    sample_name=f"{sample_name}_generation{n+1:02d}_part{part_id_int:02d}_pca",
                                    subfolder=subfolder_name,
                                )



class OverlapVisualizationCallback(VisualizationCallback):
    """Visualization callback for overlap prediction models."""

    def __init__(
        self,
        save_pairwise_overlap: bool = False,
        save_pairwise_csv: bool = False,
        save_pairwise_pointclouds: bool = True,
        **kwargs
    ):
        """Initialize overlap visualization callback.

        Args:
            save_pairwise_overlap (bool): Whether to save pairwise overlap visualizations (histograms and heatmaps). Default: False.
            save_pairwise_csv (bool): Whether to save pairwise overlap data as CSV. Default: False.
            save_pairwise_pointclouds (bool): Whether to save pairwise overlap point cloud visualizations. Default: True.
            **kwargs: Additional arguments passed to base class.
        """
        super().__init__(**kwargs)
        self.save_pairwise_overlap = save_pairwise_overlap
        self.save_pairwise_csv = save_pairwise_csv
        self.save_pairwise_pointclouds = save_pairwise_pointclouds

    def _save_pairwise_overlap_plot(
        self,
        overlap_probs: torch.Tensor,
        part_ids: torch.Tensor,
        points_per_part: torch.Tensor,
        sample_name: str,
        subfolder: Optional[str] = None,
    ):
        """Save pairwise overlap probability plot for each part combination.
        
        Args:
            overlap_probs: Overlap probabilities of shape (N, C) where C is number of parts.
            part_ids: Part IDs for each point of shape (N,).
            points_per_part: Number of points per part of shape (max_parts,).
            sample_name (str): Sample name for filename.
        """
        # Skip saving if renderer is set to "none"
        if self.renderer == "none":
            return
            
        try:
            # Move all tensors to CPU to avoid device mismatch issues
            overlap_probs = overlap_probs.cpu()
            part_ids = part_ids.cpu()
            points_per_part = points_per_part.cpu()
            
            # Get valid parts (parts with non-zero point counts)
            valid_parts = points_per_part > 0
            valid_part_indices = torch.where(valid_parts)[0]
            num_valid_parts = len(valid_part_indices)
            
            if num_valid_parts < 2:
                logger.warning(f"Sample {sample_name} has less than 2 valid parts, skipping pairwise overlap plot")
                return
            
            # Calculate number of pairwise combinations
            num_combinations = num_valid_parts * (num_valid_parts - 1) // 2
            
            # Create subplot grid
            cols = min(3, num_combinations)  # Max 3 columns
            rows = (num_combinations + cols - 1) // cols
            
            fig, axes = plt.subplots(rows, cols, figsize=(4*cols, 4*rows))
            # Ensure axes is always a list
            if num_combinations == 1:
                axes = [axes]
            elif rows == 1 and cols == 1:
                axes = [axes]
            elif rows == 1:
                axes = axes if isinstance(axes, list) else [axes]
            else:
                axes = axes.flatten()
            
            # Convert to list if it's a numpy array
            if hasattr(axes, 'flatten'):
                axes = axes.flatten().tolist()
            elif not isinstance(axes, list):
                axes = [axes]
            
            # Generate all pairwise combinations
            part_combinations = list(combinations(valid_part_indices.tolist(), 2))
            
            for idx, (part_i, part_j) in enumerate(part_combinations):
                ax = axes[idx]
                
                # Get points that belong to part_i
                part_i_mask = part_ids == part_i
                part_i_points = overlap_probs[part_i_mask]  # (N_i, C)
                
                if len(part_i_points) == 0:
                    continue
                
                # Get overlap probabilities for part_j from part_i points
                overlap_probs_i_to_j = part_i_points[:, part_j]  # (N_i,)
                
                # Create histogram of overlap probabilities
                ax.hist(overlap_probs_i_to_j.numpy(), bins=20, alpha=0.7, 
                       color='blue', edgecolor='black')
                ax.set_xlabel(f'Overlap Probability (Part {part_i} → Part {part_j})')
                ax.set_ylabel('Number of Points')
                ax.set_title(f'Part {part_i} → Part {part_j}\nMean: {overlap_probs_i_to_j.mean():.3f}')
                ax.grid(True, alpha=0.3)
                
                # Set x-axis limits to [0, 1]
                ax.set_xlim(0, 1)
            
            # Hide unused subplots
            for idx in range(num_combinations, len(axes)):
                axes[idx].set_visible(False)
            
            plt.tight_layout()
            
            # Save the plot
            target_dir = self.vis_dir if subfolder is None else (self.vis_dir / str(subfolder))
            target_dir.mkdir(parents=True, exist_ok=True)
            plot_path = target_dir / f"{sample_name}_pairwise_overlap.png"
            plt.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close()
            
        except Exception as e:
            logger.error(f"Error saving pairwise overlap plot for sample {sample_name}: {e}")

    def _save_pairwise_overlap_heatmap(
        self,
        overlap_probs: torch.Tensor,
        part_ids: torch.Tensor,
        points_per_part: torch.Tensor,
        sample_name: str,
        subfolder: Optional[str] = None,
    ):
        """Save pairwise overlap probability heatmap for all part combinations.
        
        Args:
            overlap_probs: Overlap probabilities of shape (N, C) where C is number of parts.
            part_ids: Part IDs for each point of shape (N,).
            points_per_part: Number of points per part of shape (max_parts,).
            sample_name (str): Sample name for filename.
        """
        # Skip saving if renderer is set to "none"
        if self.renderer == "none":
            return
        try:
            # Move all tensors to CPU to avoid device mismatch issues
            overlap_probs = overlap_probs.cpu()
            part_ids = part_ids.cpu()
            points_per_part = points_per_part.cpu()
            
            # Get valid parts (parts with non-zero point counts)
            valid_parts = points_per_part > 0
            valid_part_indices = torch.where(valid_parts)[0]
            num_valid_parts = len(valid_part_indices)
            
            if num_valid_parts < 2:
                logger.warning(f"Sample {sample_name} has less than 2 valid parts, skipping pairwise overlap heatmap")
                return
            
            # Create overlap matrix
            overlap_matrix = torch.zeros(num_valid_parts, num_valid_parts)
            
            # Calculate mean overlap probability for each pair
            for i, part_i in enumerate(valid_part_indices):
                for j, part_j in enumerate(valid_part_indices):
                    if i == j:
                        # Self-overlap is always 1.0 (diagonal)
                        overlap_matrix[i, j] = 1.0
                    else:
                        # Get points that belong to part_i
                        part_i_mask = part_ids == part_i
                        part_i_points = overlap_probs[part_i_mask]  # (N_i, C)
                        
                        if len(part_i_points) > 0:
                            # Get overlap probabilities for part_j from part_i points
                            overlap_probs_i_to_j = part_i_points[:, part_j]  # (N_i,)
                            overlap_matrix[i, j] = overlap_probs_i_to_j.mean()
            
            # Create heatmap
            fig, ax = plt.subplots(figsize=(8, 6))
            
            # Convert to numpy for plotting
            overlap_matrix_np = overlap_matrix.numpy()
            
            # Create heatmap
            im = ax.imshow(overlap_matrix_np, cmap='Blues', vmin=0, vmax=1)
            
            # Add colorbar
            cbar = plt.colorbar(im, ax=ax)
            cbar.set_label('Mean Overlap Probability', rotation=270, labelpad=15)
            
            # Set labels
            ax.set_xlabel('Part ID')
            ax.set_ylabel('Part ID')
            ax.set_title(f'Pairwise Overlap Probability Matrix\n{sample_name}')
            
            # Set tick labels
            ax.set_xticks(range(num_valid_parts))
            ax.set_yticks(range(num_valid_parts))
            ax.set_xticklabels([f'Part {valid_part_indices[i]}' for i in range(num_valid_parts)])
            ax.set_yticklabels([f'Part {valid_part_indices[i]}' for i in range(num_valid_parts)])
            
            # Rotate x-axis labels for better readability
            plt.setp(ax.get_xticklabels(), rotation=45, ha='right')
            
            # Add text annotations
            for i in range(num_valid_parts):
                for j in range(num_valid_parts):
                    text = ax.text(j, i, f'{overlap_matrix_np[i, j]:.3f}',
                                 ha='center', va='center', fontsize=8,
                                 color='white' if overlap_matrix_np[i, j] > 0.5 else 'black')
            
            plt.tight_layout()
            
            # Save the plot
            target_dir = self.vis_dir if subfolder is None else (self.vis_dir / str(subfolder))
            target_dir.mkdir(parents=True, exist_ok=True)
            plot_path = target_dir / f"{sample_name}_pairwise_overlap_heatmap.png"
            plt.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close()
            
        except Exception as e:
            logger.error(f"Error saving pairwise overlap heatmap for sample {sample_name}: {e}")

    def _save_pairwise_overlap_csv(
        self,
        overlap_probs: torch.Tensor,
        part_ids: torch.Tensor,
        points_per_part: torch.Tensor,
        sample_name: str,
        subfolder: Optional[str] = None,
    ):
        """Save pairwise overlap probability data as CSV for further analysis.
        
        Args:
            overlap_probs: Overlap probabilities of shape (N, C) where C is number of parts.
            part_ids: Part IDs for each point of shape (N,).
            points_per_part: Number of points per part of shape (max_parts,).
            sample_name (str): Sample name for filename.
        """
        # Skip saving if renderer is set to "none"
        if self.renderer == "none":
            return
            
        try:
            # Move all tensors to CPU to avoid device mismatch issues
            overlap_probs = overlap_probs.cpu()
            part_ids = part_ids.cpu()
            points_per_part = points_per_part.cpu()
            
            # Get valid parts (parts with non-zero point counts)
            valid_parts = points_per_part > 0
            valid_part_indices = torch.where(valid_parts)[0]
            num_valid_parts = len(valid_part_indices)
            
            if num_valid_parts < 2:
                logger.warning(f"Sample {sample_name} has less than 2 valid parts, skipping pairwise overlap CSV")
                return
            
            # Collect all pairwise overlap data
            pairwise_data = []
            
            for i, part_i in enumerate(valid_part_indices):
                for j, part_j in enumerate(valid_part_indices):
                    if i == j:
                        # Self-overlap is always 1.0
                        continue
                    
                    # Get points that belong to part_i
                    part_i_mask = part_ids == part_i
                    part_i_points = overlap_probs[part_i_mask]  # (N_i, C)
                    
                    if len(part_i_points) > 0:
                        # Get overlap probabilities for part_j from part_i points
                        overlap_probs_i_to_j = part_i_points[:, part_j]  # (N_i,)
                        
                        # Add each point's overlap probability to the data
                        for point_idx, prob in enumerate(overlap_probs_i_to_j):
                            pairwise_data.append({
                                'source_part': int(part_i),
                                'target_part': int(part_j),
                                'point_idx': point_idx,
                                'overlap_probability': float(prob),
                                'source_part_points': len(part_i_points)
                            })
            
            # Create DataFrame and save to CSV
            if pairwise_data:
                df = pd.DataFrame(pairwise_data)
                target_dir = self.vis_dir if subfolder is None else (self.vis_dir / str(subfolder))
                target_dir.mkdir(parents=True, exist_ok=True)
                csv_path = target_dir / f"{sample_name}_pairwise_overlap.csv"
                df.to_csv(csv_path, index=False)
                
                # Also save summary statistics
                summary_data = []
                for i, part_i in enumerate(valid_part_indices):
                    for j, part_j in enumerate(valid_part_indices):
                        if i == j:
                            continue
                        
                        part_data = df[(df['source_part'] == part_i) & (df['target_part'] == part_j)]
                        if len(part_data) > 0:
                            summary_data.append({
                                'source_part': int(part_i),
                                'target_part': int(part_j),
                                'mean_overlap': float(part_data['overlap_probability'].mean()),
                                'std_overlap': float(part_data['overlap_probability'].std()),
                                'min_overlap': float(part_data['overlap_probability'].min()),
                                'max_overlap': float(part_data['overlap_probability'].max()),
                                'num_points': len(part_data)
                            })
                
                if summary_data:
                    summary_df = pd.DataFrame(summary_data)
                    summary_csv_path = target_dir / f"{sample_name}_pairwise_overlap_summary.csv"
                    summary_df.to_csv(summary_csv_path, index=False)
            
        except Exception as e:
            logger.error(f"Error saving pairwise overlap CSV for sample {sample_name}: {e}")

    def _save_pairwise_overlap_pointclouds(
        self,
        overlap_probs: torch.Tensor,
        part_ids: torch.Tensor,
        points_per_part: torch.Tensor,
        pts_gt: torch.Tensor,
        sample_name: str,
        subfolder: Optional[str] = None,
    ):
        """Save point cloud visualizations for each pairwise overlap.
        
        For each pair of parts (A, B), creates a single visualization showing both parts:
        - Points of part A colored with gradient based on part A's color
        - Points of part B colored with gradient based on part B's color
        
        Args:
            overlap_probs: Overlap probabilities of shape (N, C) where C is number of parts.
            part_ids: Part IDs for each point of shape (N,).
            points_per_part: Number of points per part of shape (max_parts,).
            pts_gt: Ground truth point cloud of shape (N, 3).
            sample_name (str): Sample name for filename.
        """
        # Skip saving if renderer is set to "none"
        if self.renderer == "none":
            return
            
        try:
            # Move all tensors to CPU to avoid device mismatch issues
            overlap_probs = overlap_probs.cpu()
            part_ids = part_ids.cpu()
            points_per_part = points_per_part.cpu()
            pts_gt = pts_gt.cpu()
            
            # Get valid parts (parts with non-zero point counts)
            valid_parts = points_per_part > 0
            valid_part_indices = torch.where(valid_parts)[0]
            num_valid_parts = len(valid_part_indices)
            
            if num_valid_parts < 2:
                logger.warning(f"Sample {sample_name} has less than 2 valid parts, skipping pairwise overlap pointclouds")
                return
            
            # Generate all pairwise combinations
            part_combinations = list(combinations(valid_part_indices.tolist(), 2))
            
            for part_i, part_j in part_combinations:
                # Get points for both parts
                part_i_mask = part_ids == part_i
                part_j_mask = part_ids == part_j
                
                if part_i_mask.sum() > 0 and part_j_mask.sum() > 0:
                    # Get points for part_i
                    part_i_points = pts_gt[part_i_mask]  # (N_i, 3)
                    
                    # Get points for part_j
                    part_j_points = pts_gt[part_j_mask]  # (N_j, 3)
                    
                    # Combine points for both parts
                    combined_points = torch.cat([part_i_points, part_j_points], dim=0)  # (N_i + N_j, 3)
                    
                    # Get overlap probabilities for coloring
                    part_i_overlap_probs = overlap_probs[part_i_mask, part_j]  # (N_i,) - overlap from part_i to part_j
                    part_j_overlap_probs = overlap_probs[part_j_mask, part_i]  # (N_j,) - overlap from part_j to part_i
                    
                    # Use probs_to_colors with different colormaps for each part
                    part_i_colors = probs_to_colors(part_i_overlap_probs, colormap="matplotlib:Blues")  # Part A uses Blues
                    part_j_colors = probs_to_colors(part_j_overlap_probs, colormap="matplotlib:Reds")   # Part B uses Reds
                    
                    # Combine colors
                    combined_colors = torch.cat([part_i_colors, part_j_colors], dim=0)  # (N_i + N_j, 3)
                    
                    # Save combined visualization
                    self._save_sample_images(
                        points=combined_points,
                        colors=combined_colors,
                        sample_name=f"{sample_name}_pair{part_i}_pair{part_j}",
                        subfolder=subfolder,
                    )
            
        except Exception as e:
            logger.error(f"Error saving pairwise overlap pointclouds for sample {sample_name}: {e}")

    def on_test_batch_end(
        self,
        trainer: L.Trainer,
        pl_module: L.LightningModule,
        outputs: Any,
        batch: dict[str, torch.Tensor],
        batch_idx: int,
        dataloader_idx: int = 0,
    ):
        """Save overlap visualizations at the end of each test batch."""
        if self.vis_dir is None:
            return

        overlap_probs = outputs["overlap_probs"]                              # (total_points, C)
        points_per_part = batch["points_per_part"]                            # (bs, max_parts)
        B, _ = points_per_part.shape
        part_ids = ppp_to_ids(points_per_part)                                # (bs, N)
        
        # Handle dynamic batching - check if we have cu_seqlens_batch
        if "cu_seqlens_batch" in batch:
            # Dynamic batching format (TP, 3)
            cu_seqlens_batch = batch["cu_seqlens_batch"]                      # (B + 1,)
            pts_flat = batch["pointclouds"]                                   # (TP, 3)
            pts_gt_flat = batch["pointclouds_gt"]                            # (TP, 3)
            
            # Convert from dynamic batching to fixed batching
            pts_batches = []
            pts_gt_batches = []
            
            for b in range(B):
                start_idx = cu_seqlens_batch[b]
                end_idx = cu_seqlens_batch[b + 1]
                pts_batches.append(pts_flat[start_idx:end_idx])
                pts_gt_batches.append(pts_gt_flat[start_idx:end_idx])
            
            # Find max points for padding
            max_points = max(batch.shape[0] for batch in pts_batches)
            N = max_points
            
            # Create padded tensors
            device = pts_flat.device
            pts = torch.zeros(B, max_points, 3, device=device, dtype=pts_flat.dtype)
            pts_gt = torch.zeros(B, max_points, 3, device=device, dtype=pts_gt_flat.dtype)
            
            for b, (pts_batch, pts_gt_batch) in enumerate(zip(pts_batches, pts_gt_batches)):
                pts[b, :pts_batch.shape[0]] = pts_batch
                pts_gt[b, :pts_gt_batch.shape[0]] = pts_gt_batch
        else:
            # Fixed batching format (B, N, 3)
            pts = batch["pointclouds"].view(B, -1, 3)                         # (bs, N, 3)
            pts_gt = batch["pointclouds_gt"].reshape(B, -1, 3)                # (bs, N, 3)
            N = pts_gt.shape[1]
        
        # Reshape overlap_probs to (B, N, C)
        overlap_probs = overlap_probs.reshape(B, N, -1)

        # print(part_ids[0])
        # print("# Parts: ", max(part_ids[0]).item()+1)
        # # print(overlap_probs.shape)
        # print(overlap_probs[0])
        # check passed, make sense

        overlap_prob, _ = torch.max(overlap_probs, dim=2)                     # (bs, N)

        # Scale to original size
        if self.scale_to_original_size:
            scale = batch["scales"]                                           # (bs,)
            pts = pts * scale[:, None, None]                                  # (bs, N, 3)
            pts_gt = pts_gt * scale[:, None, None]

        for i in range(B):
            # Check if we've reached the maximum number of samples to visualize
            if self.max_samples_per_batch is not None and i >= self.max_samples_per_batch:
                break
                
            dataset_name = str(batch["dataset_name"][i])
            sample_name = f"{dataset_name}_sample{int(batch['index'][i]):05d}"
            
            # Save input image
            colors = part_ids_to_colors(
                part_ids[i], colormap="default", part_order="id"
            )
            
            # Save gt image
            self._save_sample_images(
                points=pts_gt[i],
                colors=colors,
                sample_name=f"{sample_name}_gt",
                subfolder=dataset_name,
            )
            
            # Save overlap visualization
            colors = probs_to_colors(overlap_prob[i], colormap="matplotlib:Blues")
            self._save_sample_images(
                points=pts_gt[i],
                colors=colors,
                sample_name=f"{sample_name}_overlap",
                subfolder=dataset_name,
            )
            
            # Save pairwise overlap plots
            if self.save_pairwise_overlap:
                self._save_pairwise_overlap_plot(
                    overlap_probs=overlap_probs[i],  # (N, C)
                    part_ids=part_ids[i],           # (N,)
                    points_per_part=points_per_part[i],  # (max_parts,)
                    sample_name=sample_name,
                    subfolder=dataset_name,
                )

            # Save pairwise overlap heatmap
            if self.save_pairwise_overlap:
                self._save_pairwise_overlap_heatmap(
                    overlap_probs=overlap_probs[i],  # (N, C)
                    part_ids=part_ids[i],           # (N,)
                    points_per_part=points_per_part[i],  # (max_parts,)
                    sample_name=sample_name,
                    subfolder=dataset_name,
                )

            # Save pairwise overlap CSV
            if self.save_pairwise_csv:
                self._save_pairwise_overlap_csv(
                    overlap_probs=overlap_probs[i],  # (N, C)
                    part_ids=part_ids[i],           # (N,)
                    points_per_part=points_per_part[i],  # (max_parts,)
                    sample_name=sample_name,
                    subfolder=dataset_name,
                )

            # Save pairwise overlap pointclouds
            if self.save_pairwise_pointclouds:
                self._save_pairwise_overlap_pointclouds(
                    overlap_probs=overlap_probs[i],  # (N, C)
                    part_ids=part_ids[i],           # (N,)
                    points_per_part=points_per_part[i],  # (max_parts,)
                    pts_gt=pts_gt[i],               # (N, 3)
                    sample_name=sample_name,
                    subfolder=dataset_name,
                )