image_processing.py

"""
Image processing utilities for the Image-to-3D pipeline.
"""

import os
import json
import numpy as np
import torch
from typing import Any, Dict, List, Tuple, Union
from PIL import Image
from torchvision import transforms as TF
from config import PROCESSING_PARAMS
import open3d as o3d
from scene_utils import filter_depth_discontinuities


def load_and_preprocess_images(image_list):
    images = []
    shapes = set()
    to_tensor = TF.ToTensor()
    
    for img in image_list:
        if img.mode == 'RGBA':
            background = Image.new('RGBA', img.size, (255, 255, 255, 255))
            img = Image.alpha_composite(background, img)
        
        img = img.convert("RGB")
        width, height = img.size
        new_width = PROCESSING_PARAMS["depth_image_size"]
        new_height = round(height * (new_width / width) / 14) * 14
        
        img = img.resize((new_width, new_height), Image.Resampling.BICUBIC)
        img = to_tensor(img)
        
        if new_height > PROCESSING_PARAMS["depth_image_size"]:
            start_y = (new_height - PROCESSING_PARAMS["depth_image_size"]) // 2
            img = img[:, start_y : start_y + PROCESSING_PARAMS["depth_image_size"], :]
        
        shapes.add((img.shape[1], img.shape[2]))
        images.append(img)
    
    if len(shapes) > 1:
        print(f"Warning: Found images with different shapes: {shapes}")
        max_height = max(shape[0] for shape in shapes)
        max_width = max(shape[1] for shape in shapes)
        
        padded_images = []
        for img in images:
            h_padding = max_height - img.shape[1]
            w_padding = max_width - img.shape[2]
            
            if h_padding > 0 or w_padding > 0:
                pad_top = h_padding // 2
                pad_bottom = h_padding - pad_top
                pad_left = w_padding // 2
                pad_right = w_padding - pad_left
                
                img = torch.nn.functional.pad(
                    img, (pad_left, pad_right, pad_top, pad_bottom), 
                    mode="constant", value=1.0
                )
            padded_images.append(img)
        images = padded_images
    
    images = torch.stack(images)
    
    if len(image_list) == 1:
        if images.dim() == 3:
            images = images.unsqueeze(0)
    
    return images


class ImageProcessor:
    """Handles image processing operations."""

    def __init__(self, model_manager, use_moge_v2=False):
        self.model_manager = model_manager
        self.use_moge_v2 = use_moge_v2

    @torch.no_grad()
    def get_moge_v2_pc(self, image, save_folder="./"):
        """Generate point cloud using MoGe-v2 model with depth conversion."""
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        if image.mode != "RGB":
            image = image.convert("RGB")

        # Get MoGe-v2 model
        moge_model = self.model_manager.get_model('moge_v2')

        # Prepare image for MoGe-v2
        image_np = np.array(image)
        input_tensor = torch.tensor(image_np / 255, dtype=torch.float32, device=moge_model.device).permute(2, 0, 1)

        # Run MoGe-v2 inference
        output = moge_model.infer(input_tensor, force_projection=False)

        # Extract outputs from MoGe-v2 (metric scale)
        depth_map = output['depth'].detach().cpu().numpy()  # (H, W)
        mask = output['mask'].detach().cpu().numpy()  # (H, W) - binary mask for valid pixels
        intrinsics_normalized = output['intrinsics'].detach().cpu().numpy()  # (3, 3) - normalized camera intrinsics

        max_depth = depth_map.max()
        min_depth = depth_map.min()
        if (max_depth-min_depth) > 200:
            mean_depth = depth_map.mean()
            std_depth = depth_map.std()
            if max_depth > mean_depth + std_depth*3.0:
                print(f"  Warning: Depth map max value {max_depth:.3f} is significantly larger than mean+std {mean_depth + std_depth*3.0:.3f}. Clipping depth values.")
                mask = mask & (depth_map <= (mean_depth + std_depth*3.0))

        H, W = depth_map.shape

        # Apply edge filtering to remove flying pixels at depth discontinuities
        edge_filter_params = PROCESSING_PARAMS.get("edge_filter_params", {
            "threshold": 0.1,
            "edge_dilation": 2,
            "kernel_size": 5
        })
        edge_mask = filter_depth_discontinuities(depth_map, **edge_filter_params)
        print(f"  MoGe edge filtering: kept {edge_mask.sum()} / {edge_mask.size} pixels ({edge_mask.sum()/edge_mask.size*100:.1f}%)")
        mask = mask & edge_mask

        # Convert normalized intrinsics to pixel space
        # MoGe-v2 outputs normalized intrinsics, we need to scale them to image size
        intrinsics = intrinsics_normalized.copy()
        intrinsics[0, 0] *= W  # fx
        intrinsics[1, 1] *= H  # fy
        intrinsics[0, 2] *= W  # cx
        intrinsics[1, 2] *= H  # cy

        # Create camera extrinsic (identity - camera at origin looking down +Z in OpenCV convention)
        # OpenCV convention: x right, y down, z forward
        extrinsic = np.eye(4, dtype=np.float32)

        # Convert depth map to point map using intrinsics
        # Create pixel grid
        y_coords, x_coords = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')

        # Backproject to camera coordinates using intrinsics
        fx, fy = intrinsics[0, 0], intrinsics[1, 1]
        cx, cy = intrinsics[0, 2], intrinsics[1, 2]

        # Camera coordinates (OpenCV convention: x right, y down, z forward)
        x_cam = (x_coords - cx) * depth_map / fx
        y_cam = (y_coords - cy) * depth_map / fy
        z_cam = depth_map

        # Stack to create point map [H, W, 3]
        points_map = np.stack([x_cam, y_cam, z_cam], axis=-1).astype(np.float32)

        # Extract valid points for point cloud
        valid_points = points_map[mask]
        valid_colors = image_np[mask] / 255.0

        # Create point cloud
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(valid_points)
        pcd.colors = o3d.utility.Vector3dVector(valid_colors)

        # Format output to match HunyuanWorld-Mirror structure
        # Add batch dimension [S, H, W, ...] where S=1 for single frame
        moge_output = {
            "depth": depth_map[np.newaxis, :, :, np.newaxis],  # [1, H, W, 1]
            "point_map": points_map[np.newaxis, :, :, :],  # [1, H, W, 3]
            "point_map_mask": mask[np.newaxis, :, :],  # [1, H, W]
            "point_map_mask_full": mask[np.newaxis, :, :],  # [1, H, W]
            "extrinsics": extrinsic[np.newaxis, :, :],  # [1, 4, 4] - c2w format
            "intrinsics": intrinsics[np.newaxis, :, :],  # [1, 3, 3]
        }
        return pcd, image, moge_output

    @torch.no_grad()
    def get_hunyuan_world_mirror_pc(self, image, save_folder="./", camera_pose=None, camera_intrinsics=None):
        """Generate point cloud using HunyuanWorld-Mirror model."""
        from src.models.utils.geometry import depth_to_world_coords_points
        from torchvision import transforms as TF

        hunyuan_model = self.model_manager.get_model('hunyuan_world_mirror')
        model_device = next(hunyuan_model.parameters()).device
    

        if isinstance(image, Image.Image):
            image = [image]
        imgs = load_and_preprocess_images(image).unsqueeze(0).to(model_device)

        # Run inference
        if camera_pose is None and camera_intrinsics is None:
            views = {"img": imgs}
            cond_flags = [0, 0, 0]  # [depth, rays, camera]
        else:
            camera_pose = torch.from_numpy(camera_pose).float().unsqueeze(0).to(model_device)
            camera_intrinsics = torch.from_numpy(camera_intrinsics).float().unsqueeze(0).to(model_device)

            views = {
                "img": imgs,
                "camera_pose": camera_pose,
                "camera_intrinsics": camera_intrinsics,
            }
            cond_flags = [0, 1, 1]  # [depth, rays, camera]

        with torch.autocast("cuda", dtype=torch.bfloat16):
            predictions = hunyuan_model(views=views, cond_flags=cond_flags)

        # Extract predictions
        # pts3d_preds = predictions["pts3d"][0]  # [S, H, W, 3]
        # pts3d_conf = predictions["pts3d_conf"][0]  # [S, H, W]
        depth_preds = predictions["depth"][0]  # [S, H, W, 1]
        depth_conf = predictions["depth_conf"][0]  # [S, H, W]
        camera_poses = predictions["camera_poses"][0]  # [S, 4, 4]
        camera_intrs = predictions["camera_intrs"][0]  # [S, 3, 3]

        point_map, _, _ = depth_to_world_coords_points(
            depth_preds[..., 0], camera_poses, camera_intrs
        )  # [S, H, W, 3]
        # Convert to numpy
        point_map_np = point_map.cpu().numpy()  # [S, H, W, 3]
        depth_np = depth_preds.cpu().numpy()  # [S, H, W, 1]
        depth_conf_np = depth_conf.cpu().numpy()  # [S, H, W]
        camera_poses_np = camera_poses.cpu().numpy()  # [S, 4, 4]
        camera_intrs_np = camera_intrs.cpu().numpy()  # [S, 3, 3]

        # Process each frame
        num_frames = point_map_np.shape[0]

        # Create confidence masks for each frame
        point_valid_masks = []
        point_map_masks_full = []

        for i in range(num_frames):
            depth_conf_frame = depth_conf_np[i]
            valid_conf = depth_conf_frame[~np.isnan(depth_conf_frame)]
            if valid_conf.size == 0:
                point_valid_mask = np.ones(depth_conf_frame.shape, dtype=bool)
            else:
                conf_percentile = float(PROCESSING_PARAMS.get("depth_conf_threshold", 0.0))
                if conf_percentile <= 0:
                    threshold = valid_conf.min()
                else:
                    conf_percentile = np.clip(conf_percentile, 0.0, 100.0)
                    threshold = np.percentile(valid_conf, conf_percentile)
                point_valid_mask = depth_conf_frame >= threshold

            point_valid_masks.append(point_valid_mask)
            point_map_masks_full.append(np.ones(depth_conf_frame.shape, dtype=bool))

        point_valid_masks = np.array(point_valid_masks)  # [S, H, W]
        point_map_masks_full = np.array(point_map_masks_full)  # [S, H, W]

        max_depth = depth_np[0].max()
        min_depth = depth_np[0].min()
        if (max_depth - min_depth) > 2:
            mean_depth = depth_np[0].mean()
            std_depth = depth_np[0].std()
            if max_depth > mean_depth + std_depth * 3.0:
                print(f"  Warning: Depth map max value {max_depth:.3f} is significantly larger than mean+std {mean_depth + std_depth * 3.0:.3f}. Clipping depth values.")
                point_valid_masks[0] = point_valid_masks[0] & (depth_np[0, ..., 0] <= (mean_depth + std_depth*3.0))

        # Apply edge filtering to remove flying pixels at depth discontinuities for each frame
        edge_filter_params = PROCESSING_PARAMS.get("edge_filter_params", {
            "threshold": 0.1,
            "edge_dilation": 2,
            "kernel_size": 5
        })
        for i in range(num_frames):
            edge_mask = filter_depth_discontinuities(depth_np[i, ..., 0], **edge_filter_params)
            print(f"  Hunyuan frame {i} edge filtering: kept {edge_mask.sum()} / {edge_mask.size} pixels ({edge_mask.sum()/edge_mask.size*100:.1f}%)")
            point_valid_masks[i] = point_valid_masks[i] & edge_mask

        # Extract valid points for point cloud (combine all frames)
        all_valid_points = []
        all_valid_colors = []
        rgb_all = imgs[0].permute((0, 2, 3, 1)).cpu().numpy()  # [S, H, W, 3]

        for i in range(num_frames):
            valid_points = point_map_np[i][point_valid_masks[i]]
            valid_colors = rgb_all[i][point_valid_masks[i]]
            all_valid_points.append(valid_points)
            all_valid_colors.append(valid_colors)

        all_valid_points = np.concatenate(all_valid_points, axis=0)
        all_valid_colors = np.concatenate(all_valid_colors, axis=0)

        # Create combined point cloud
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(all_valid_points)
        pcd.colors = o3d.utility.Vector3dVector(all_valid_colors)

        # Convert back to PIL image (use first frame)
        rgb = rgb_all[0]
        rgb_img = Image.fromarray((np.clip(rgb * 255, 0, 255)).astype(np.uint8))

        hunyuan_output = {
            "depth": depth_np,  # [S, H, W, 1]
            "point_map": point_map_np,  # [S, H, W, 3]
            "point_map_mask": point_valid_masks,  # [S, H, W]
            "point_map_mask_full": point_map_masks_full,  # [S, H, W]
            "extrinsics": camera_poses_np,  # [S, 4, 4] - c2w format from HunyuanWorld-Mirror
            "intrinsics": camera_intrs_np,  # [S, 3, 3]
        }

        return pcd, rgb_img, hunyuan_output
    
    def normalize_point_cloud(self, points: np.ndarray, camera_forward=None) -> Tuple[np.ndarray, Dict[str, Any]]:
        from scene_utils import align_point_cloud

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        transform_align = align_point_cloud(pcd, horizontal_axis='y', camera_forward=camera_forward)

        aligned_points = (points - transform_align[:3, 3]) @ transform_align[:3, :3].T
        min_bound = aligned_points.min(axis=0)
        max_bound = aligned_points.max(axis=0)
        center = (min_bound + max_bound) / 2
        scale = (max_bound - min_bound).max()

        scaled_points = (aligned_points - center) / (scale + 1e-8)

        norm_params = {
            "align_transform": transform_align.tolist(),
            "center": center.tolist(),
            "scale": float(scale),
        }

        return scaled_points, norm_params

    def generate_point_cloud_info(self, image: Image.Image, save_folder: str) -> Dict[str, Any]:
        """Generate depth, point cloud, and normalization metadata."""

        info: Dict[str, Any] = {}

        original_w, original_h = image.size

        self.model_manager.ensure_pointcloud_backbone(
            use_moge_v2=self.use_moge_v2,
        )

        use_moge_v2 = self.use_moge_v2 and self.model_manager.get_model('moge_v2') is not None

        if use_moge_v2:
            overall_pcd, image, pc_image_info = self.get_moge_v2_pc(image, save_folder=save_folder)
        else:
            overall_pcd, image, pc_image_info = self.get_hunyuan_world_mirror_pc(image, save_folder=save_folder)

        info.update(pc_image_info)

        point_map = info["point_map"]
        point_mask = info.get("point_map_mask")
        scale_points, norm_params = self.normalize_point_cloud(
            point_map[point_mask],
            camera_forward = pc_image_info['extrinsics'][0, :3, 2] if 'extrinsics' in pc_image_info else None,
        )

        overall_pcd.points = o3d.utility.Vector3dVector(scale_points)
        o3d.io.write_point_cloud(os.path.join(save_folder, "overall_pcd.ply"), overall_pcd)

        filtered_point_map = np.zeros_like(point_map)
        filtered_point_map[point_mask] = scale_points

        norm_path = os.path.join(save_folder, "normalization.json")
        with open(norm_path, "w") as f:
            json.dump(norm_params, f, indent=2)
        
        info["point_map_raw"] = point_map
        info["point_map"] = filtered_point_map
        info["normalization"] = norm_params
        info["image"] = image
        return info

    def to_trellis_camera(self, extrinsic: np.ndarray, intrinsic: np.ndarray, normalization: Dict[str, Any], image_size: Tuple[int, int]) -> Dict[str, Any]:
        """Convert extrinsic and intrinsic to Trellis camera format.
        """
        center = np.array(normalization["center"], dtype=np.float32)
        scale = float(normalization["scale"]) / 2.0
        transform_align = np.array(normalization["align_transform"], dtype=np.float32)

        extrinsic = extrinsic.copy()
        intrinsic = intrinsic.copy()

        # Normalize extrinsic
        # This is mathematically unusual but empirically correct for this pipeline
        normalize_translate = np.eye(4)
        normalize_translate[:3, 3] = -center
        T_total = normalize_translate @ transform_align
        extrinsic = (T_total @ extrinsic)[:3, :4]

        R = extrinsic[:3, :3]
        t = extrinsic[:3, 3]
        cam_pos = -R.T @ t

        # Scale camera position
        cam_pos_scaled = cam_pos / (scale + 1e-8)

        # Reconstruct extrinsic with scaled position
        t_scaled = -R @ cam_pos_scaled
        extrinsic_for_rendering = np.zeros((3, 4))
        extrinsic_for_rendering[:3, :3] = R
        extrinsic_for_rendering[:3, 3] = t_scaled

        # Convert to 4x4 and invert to get final w2c
        extrinsic_c2w = np.eye(4)
        extrinsic_c2w[:3, :4] = extrinsic_for_rendering
        extrinsic = np.linalg.inv(extrinsic_c2w)

        width, height = image_size

        intrinsic[0, 0] = intrinsic[0, 0] / width
        intrinsic[1, 1] = intrinsic[1, 1] / height
        intrinsic[0, 2] = intrinsic[0, 2] / width
        intrinsic[1, 2] = intrinsic[1, 2] / height

        return extrinsic, intrinsic, extrinsic_c2w[:3, 3]

    def from_trellis_camera(self, extrinsic_trellis: np.ndarray, normalization: Dict[str, Any],
                            intrinsic_trellis: np.ndarray=None, image_size: Tuple[int, int]=None) -> Tuple[np.ndarray, np.ndarray]:
        """Convert Trellis camera format back to original extrinsic and intrinsic parameters.

        This is the inverse operation of to_trellis_camera().
        """
        center = np.array(normalization["center"], dtype=np.float32)
        scale = float(normalization["scale"]) / 2.0
        transform_align = np.array(normalization["align_transform"], dtype=np.float32)

        intrinsic = None
        if intrinsic_trellis is not None and image_size is not None:
            intrinsic = intrinsic_trellis.copy()

            width, height = image_size

            # Step 1: Denormalize intrinsic (reverse normalization)
            intrinsic[0, 0] = intrinsic[0, 0] * width
            intrinsic[1, 1] = intrinsic[1, 1] * height
            intrinsic[0, 2] = intrinsic[0, 2] * width
            intrinsic[1, 2] = intrinsic[1, 2] * height

        extrinsic = extrinsic_trellis.copy()
        # Step 2: Convert world-to-camera back to camera-to-world
        extrinsic_c2w = np.linalg.inv(extrinsic)
        extrinsic_for_rendering = extrinsic_c2w[:3, :4]

        # Step 3: Extract R and scaled translation
        R = extrinsic_for_rendering[:3, :3]
        t_scaled = extrinsic_for_rendering[:3, 3]

        # Step 4: Recover scaled camera position
        cam_pos_scaled = -R.T @ t_scaled

        # Step 5: Unscale camera position
        cam_pos = cam_pos_scaled * (scale + 1e-8)

        # Step 6: Reconstruct extrinsic with unscaled position
        t = -R @ cam_pos
        extrinsic_normalized = np.zeros((3, 4))
        extrinsic_normalized[:3, :3] = R
        extrinsic_normalized[:3, 3] = t

        # Step 7: Denormalize extrinsic (reverse the transformation)
        normalize_translate = np.eye(4)
        normalize_translate[:3, 3] = -center
        T_total = normalize_translate @ transform_align
        T_total_inv = np.linalg.inv(T_total)

        extrinsic_homo = np.eye(4)
        extrinsic_homo[:3, :4] = extrinsic_normalized
        extrinsic_original_homo = T_total_inv @ extrinsic_homo

        return extrinsic_original_homo, intrinsic

    def generate_orbit_cameras(self, extrinsic_trellis: np.ndarray, num_views: int,
                              rotation_angle_degrees: float = None) -> List[np.ndarray]:
        """Generate camera extrinsics rotating around Z-axis while maintaining distance to origin.

        Args:
            extrinsic_trellis: Initial world-to-camera matrix (4x4) in Trellis format
            num_views: Number of camera views to generate
            rotation_angle_degrees: Total rotation angle in degrees. If None, uses 360 degrees for full orbit.
                                   Supports range from -360 to 360 degrees:
                                   - Positive values: counter-clockwise rotation (standard)
                                   - Negative values: clockwise rotation (reverse direction)
                                   - Example: -180 rotates 180 degrees in the opposite direction

        Returns:
            List of world-to-camera matrices (4x4) for each rotated camera position
        """
        if rotation_angle_degrees is None:
            rotation_angle_degrees = 360.0

        # Convert to camera-to-world for easier manipulation
        c2w = np.linalg.inv(extrinsic_trellis)

        # Extract camera position and rotation in world coordinates
        cam_pos = c2w[:3, 3]
        R_original = c2w[:3, :3]

        # Calculate the initial angle in XY plane
        initial_angle = np.arctan2(cam_pos[1], cam_pos[0])

        # Calculate angle step (supports negative values for reverse rotation)
        angle_step = np.deg2rad(rotation_angle_degrees) / max(num_views - 1, 1) if num_views > 1 else 0

        orbit_cameras = []

        for i in range(num_views):
            # Calculate rotation angle for this view
            # Negative angle_step will automatically rotate in the opposite direction
            angle = i * angle_step

            # Create rotation matrix around Z-axis
            cos_a = np.cos(angle)
            sin_a = np.sin(angle)
            R_z = np.array([
                [cos_a, -sin_a, 0],
                [sin_a, cos_a, 0],
                [0, 0, 1]
            ], dtype=np.float32)

            # Rotate camera position around Z-axis
            new_cam_pos = R_z @ cam_pos

            # Rotate camera orientation around Z-axis
            R_new = R_z @ R_original

            # Construct camera-to-world matrix
            c2w_new = np.eye(4)
            c2w_new[:3, :3] = R_new
            c2w_new[:3, 3] = new_cam_pos

            # Convert back to world-to-camera (Trellis format)
            w2c = np.linalg.inv(c2w_new)

            orbit_cameras.append(w2c)

        return orbit_cameras

    def process_image(self, image, save_folder, need_mask=True, **kwargs):
        """Process input image and extract all necessary information."""
        if not os.path.exists(save_folder):
            os.makedirs(save_folder)

        if need_mask:
            mask_info = self.generate_mask_info(image, save_folder)
            image_info = mask_info.copy() if mask_info else {}
        else:
            image_info = {}
        image_info["original_images"] = [image]

        point_cloud_info = self.generate_point_cloud_info(image, save_folder)
        image_info.update(point_cloud_info)
        trellis_extrinsic, trellis_intrinsic, trellis_cam_pos = self.to_trellis_camera(
            image_info["extrinsics"][0],
            image_info["intrinsics"][0],
            image_info["normalization"],
            image_info["image"].size,
        )
        image_info["trellis_extrinsic"] = trellis_extrinsic[None, ...]
        image_info["trellis_intrinsic"] = trellis_intrinsic[None, ...]
        image_info["trellis_camera_position"] = trellis_cam_pos[None, ...]
        # save trellis camera params
        trellis_camera_path = os.path.join(save_folder, "trellis_camera.json")
        trellis_camera_data = {
            "extrinsic": image_info["trellis_extrinsic"].tolist(),
            "intrinsic": image_info["trellis_intrinsic"].tolist(),
        }
        with open(trellis_camera_path, "w") as f:
            json.dump(trellis_camera_data, f, indent=2)

        num_views = int(kwargs.get("num_orbit_views", 121))
        rotation_angle_degrees = float(kwargs.get("orbit_rotation_angle_degrees", 30.0))
        orbit_cameras = self.generate_orbit_cameras(
            trellis_extrinsic,
            num_views=num_views,
            rotation_angle_degrees=rotation_angle_degrees,
        )
        image_info["orbit_cameras_extrinsic"] = orbit_cameras

        original_orbit_cameras = []
        for cam in orbit_cameras:
            orig_extrinsic, _ = self.from_trellis_camera(
                cam,
                image_info["normalization"]
            )
            original_orbit_cameras.append(orig_extrinsic)
        image_info["orbit_cameras_original_extrinsic"] = np.array(original_orbit_cameras)

        return image_info

    def filter_by_multiview_voting(self, points, colors, video_point_map, video_point_map_mask,
                                   extrinsics, intrinsics, image_shape,
                                   vote_threshold=0.5, depth_tolerance=0.1):
        """Filter point cloud using multi-view consistency voting.

        For each input point, check:
        1. How many views can see it (within frustum and depth range)
        2. For views that should see it, check if it's occluded by actual depth
        3. Only keep points that are consistent across multiple views

        Args:
            points: [N, 3] points to filter in world space (e.g., from stage 1 filtering)
            colors: [N, 3] or None, colors for the points
            video_point_map: [N_views, H, W, 3] point cloud in world space for depth reference
            video_point_map_mask: [N_views, H, W] valid point mask
            extrinsics: [N_views, 4, 4] camera-to-world matrices
            intrinsics: [N_views, 3, 3] camera intrinsic matrices
            image_shape: (H, W) image dimensions
            vote_threshold: minimum fraction of views that must agree (0-1)
            depth_tolerance: relative tolerance for depth comparison

        Returns:
            filtered_points: filtered points in world space
            filtered_colors: corresponding colors
            valid_mask: boolean mask for input points
        """
        num_views = len(video_point_map)
        h, w = image_shape[:2]

        print(f"Input points for voting: {len(points)}")

        # Build depth maps for each view for occlusion testing
        depth_maps = []
        for view_idx in range(num_views):
            depth_map = np.full((h, w), np.inf, dtype=np.float32)

            # Get w2c transformation
            extrinsic_w2c = np.linalg.inv(extrinsics[view_idx])
            R = extrinsic_w2c[:3, :3]
            t = extrinsic_w2c[:3, 3]

            # Get valid points for this view
            view_points = video_point_map[view_idx][video_point_map_mask[view_idx]]

            # Transform to camera space
            points_cam = (view_points @ R.T) + t

            # Project to image
            fx, fy = intrinsics[view_idx][0, 0], intrinsics[view_idx][1, 1]
            cx, cy = intrinsics[view_idx][0, 2], intrinsics[view_idx][1, 2]

            u = (points_cam[:, 0] * fx / (points_cam[:, 2] + 1e-8) + cx).astype(int)
            v = (points_cam[:, 1] * fy / (points_cam[:, 2] + 1e-8) + cy).astype(int)

            # Store minimum depth at each pixel (vectorized using np.minimum.at)
            valid_proj = (u >= 0) & (u < w) & (v >= 0) & (v < h) & (points_cam[:, 2] > 0)
            valid_indices = np.where(valid_proj)[0]

            if len(valid_indices) > 0:
                u_valid = u[valid_indices]
                v_valid = v[valid_indices]
                depths_valid = points_cam[valid_indices, 2]

                # Use numpy's at method for efficient in-place minimum
                # Note: this handles duplicate (v, u) coordinates correctly by keeping minimum
                np.minimum.at(depth_map, (v_valid, u_valid), depths_valid)

            depth_maps.append(depth_map)

        # Vote for each input point
        # Track both votes and the number of views that can theoretically see each point
        vote_counts = np.zeros(len(points), dtype=np.int32)
        visible_view_counts = np.zeros(len(points), dtype=np.int32)  # How many views can see this point

        for view_idx in range(num_views):
            print(f"Voting from view {view_idx + 1}/{num_views}...")

            # Get w2c transformation
            extrinsic_w2c = np.linalg.inv(extrinsics[view_idx])
            R = extrinsic_w2c[:3, :3]
            t = extrinsic_w2c[:3, 3]

            # Transform input points to this view's camera space
            points_cam = (points @ R.T) + t

            # Project to image
            fx, fy = intrinsics[view_idx][0, 0], intrinsics[view_idx][1, 1]
            cx, cy = intrinsics[view_idx][0, 2], intrinsics[view_idx][1, 2]

            u = (points_cam[:, 0] * fx / (points_cam[:, 2] + 1e-8) + cx).astype(int)
            v = (points_cam[:, 1] * fy / (points_cam[:, 2] + 1e-8) + cy).astype(int)
            depths = points_cam[:, 2]

            # Check if points are within image bounds and in front of camera
            in_frustum = (u >= 0) & (u < w) & (v >= 0) & (v < h) & (depths > 0)

            # For points in frustum, check occlusion using vectorized operations
            depth_map = depth_maps[view_idx]

            # Get actual depths at projected pixel locations (vectorized)
            frustum_indices = np.where(in_frustum)[0]
            if len(frustum_indices) == 0:
                continue

            u_frustum = u[frustum_indices]
            v_frustum = v[frustum_indices]
            depths_frustum = depths[frustum_indices]

            # Lookup actual depths from depth map (vectorized)
            actual_depths = depth_map[v_frustum, u_frustum]

            # Vectorized depth consistency checks
            has_actual_depth = ~np.isinf(actual_depths)

            # CORRECT voting logic:
            #
            # For a candidate point P (with predicted depth D_p):
            #   Project P to view V at pixel location (u,v)
            #   Get actual observed depth D_v at that pixel in view V
            #
            # Case 1: D_v ≈ D_p (depths match within tolerance)
            #   → View V confirms this point exists! → +1 vote
            #
            # Case 2: D_v > D_p (view V sees something FARTHER than P)
            #   → P should be in front and visible to V
            #   → But V sees something farther instead
            #   → P is likely an ERROR (doesn't actually exist) → -1 vote
            #
            # Case 3: D_v < D_p (view V sees something CLOSER than P)
            #   → P is occluded by foreground in view V
            #   → This is reasonable, P might still exist → 0 vote (neutral)
            #
            # Case 4: No observation at (u,v) in view V
            #   → Could be occlusion, edge, or out of view → 0 vote (neutral)

            # Calculate relative error for points with actual depth
            relative_error = np.abs(depths_frustum - actual_depths) / (actual_depths + 1e-8)

            # Case 1: Depth match → +1 vote (confirmed observation)
            depth_match = (relative_error <= depth_tolerance) & has_actual_depth

            # Case 2: Actual depth is FARTHER than predicted → -1 vote (ERROR!)
            # This is the key insight: if other view sees farther, this point shouldn't exist
            actual_farther = (actual_depths > depths_frustum * (1 + depth_tolerance)) & has_actual_depth

            # Case 3: Occluded by closer object (has_actual_depth and actual_depth < predicted)
            # These points are "potentially visible" - count them as visible views
            occluded = (actual_depths < depths_frustum * (1 - depth_tolerance)) & has_actual_depth

            # Count visible views: in frustum AND (either matched, or farther, or occluded, or no observation)
            # Basically, any point in frustum is "potentially visible" from this view
            visible_view_counts[frustum_indices] += 1

            # Apply votes (vectorized)
            vote_counts[frustum_indices[depth_match]] += 1
            vote_counts[frustum_indices[actual_farther]] -= 1

        # Filter based on votes with per-point visible view normalization
        #
        # Key insight: vote_threshold should be relative to how many views can actually see the point
        # Example:
        #   - Point A: visible in 8 views, needs 8 * 0.5 = 4 positive votes
        #   - Point B: visible in 2 views, needs 2 * 0.5 = 1 positive vote
        #
        # This prevents unfairly penalizing points that are only visible from few angles

        # Calculate required votes per point (relative to its visible view count)
        required_votes = np.maximum(1, (visible_view_counts * vote_threshold).astype(int))

        # Keep points that have:
        # 1. Positive net votes (vote_counts > 0), AND
        # 2. Enough votes relative to their visible view count
        valid_mask = (vote_counts >= required_votes) & (vote_counts > 0)

        filtered_points = points[valid_mask]
        filtered_colors = colors[valid_mask] if colors is not None else None

        print(f"Vote threshold: {vote_threshold} (relative to visible views per point)")
        print(f"Visible view distribution: min={visible_view_counts.min()}, max={visible_view_counts.max()}, mean={visible_view_counts.mean():.1f}")
        print(f"Vote distribution: min={vote_counts.min()}, max={vote_counts.max()}, mean={vote_counts.mean():.1f}, median={np.median(vote_counts):.1f}")
        print(f"Positive votes: {(vote_counts > 0).sum()}, Negative votes: {(vote_counts < 0).sum()}, Zero: {(vote_counts == 0).sum()}")
        print(f"Filtered points: {len(filtered_points)} / {len(points)} ({100*len(filtered_points)/len(points):.1f}%)")

        return filtered_points, filtered_colors, valid_mask

    def filter_by_first_frame_frustum_with_depth(self, video_point_map, video_point_map_mask,
                                                 first_frame_extrinsic, first_frame_intrinsic,
                                                image_shape, video_pcd_colors):
        """Filter using first frame's actual point cloud depth range.

        Args:
            first_frame_extrinsic: Camera-to-world (c2w) 4x4 matrix in original (non-normalized) world space
            first_frame_intrinsic: Camera intrinsic 3x3 matrix
            video_point_map: Point cloud in original (non-normalized) world space
        """

        # align_transform = np.array(normalization["align_transform"], dtype=np.float32)
        # scale = float(normalization["scale"])
        # center = np.array(normalization["center"], dtype=np.float32)

        # Convert c2w to w2c for easier transformation
        extrinsic_w2c = np.linalg.inv(first_frame_extrinsic)
        R = extrinsic_w2c[:3, :3]
        t = extrinsic_w2c[:3, 3]

        # Get first frame points and their depth range (points are in world space)
        first_frame_points = video_point_map[0][video_point_map_mask[0]].reshape(-1, 3)

        # Transform first frame points to camera space to get depth range
        # P_cam = R @ P_world + t
        first_frame_cam = (first_frame_points @ R.T) + t
        depth_min = first_frame_cam[:, 2].min()
        depth_max = first_frame_cam[:, 2].max()

        print(f"First frame depth range: {depth_min:.3f} to {depth_max:.3f}")

        # Get all points (in world space)
        all_points = video_point_map[video_point_map_mask].reshape(-1, 3)

        # Transform to first frame's camera space
        points_cam = (all_points @ R.T) + t

        # Filter by depth range
        depth_mask = (points_cam[:, 2] >= depth_min) & (points_cam[:, 2] <= depth_max)

        # Project to image space
        fx, fy = first_frame_intrinsic[0, 0], first_frame_intrinsic[1, 1]
        cx, cy = first_frame_intrinsic[0, 2], first_frame_intrinsic[1, 2]

        u = (points_cam[:, 0] * fx / (points_cam[:, 2] + 1e-8)) + cx
        v = (points_cam[:, 1] * fy / (points_cam[:, 2] + 1e-8)) + cy

        # Check if within image bounds
        h, w = image_shape[:2]
        image_mask = (u >= 0) & (u < w) & (v >= 0) & (v < h)

        # Combine all masks
        frustum_mask = depth_mask & image_mask

        # Now normalize the filtered points
        filtered_points_world = all_points[frustum_mask]
        # filtered_points_normalized = (filtered_points_world - align_transform[:3, 3]) @ align_transform[:3, :3].T
        # filtered_points_normalized = (filtered_points_normalized - center) / (scale + 1e-8)

        # Filter colors
        filtered_colors = None
        if video_pcd_colors is not None:
            filtered_colors = np.asarray(video_pcd_colors)[frustum_mask]

        print(f"Filtered points: {filtered_points_world.shape[0]} / {all_points.shape[0]} ({100*filtered_points_world.shape[0]/all_points.shape[0]:.1f}%)")

        return filtered_points_world, filtered_colors, frustum_mask

    def process_video(self, video_images=None, save_folder=None, **kwargs):
        """Process input video and extract all necessary information."""

        self.model_manager.ensure_pointcloud_backbone()
        
        selected_frames = kwargs.get("selected_frames", 24)
        if selected_frames < len(video_images):
            indices = np.linspace(0, len(video_images) - 1, selected_frames).astype(int)
            video_images = [video_images[i] for i in indices]


        video_pcd, _, video_pc_info = self.get_hunyuan_world_mirror_pc(video_images, save_folder=save_folder)

        # Use HunyuanWorld-Mirror's predicted camera parameters (c2w format)
        video_point_map = video_pc_info["point_map"]  # [S, H, W, 3]
        video_point_map_mask = video_pc_info.get("point_map_mask")  # [S, H, W]
        predicted_extrinsics = video_pc_info["extrinsics"]  # [S, 4, 4] c2w from HunyuanWorld-Mirror
        predicted_intrinsics = video_pc_info["intrinsics"]  # [S, 3, 3]

        # Two-stage filtering:
        # 1. First filter by first frame's frustum to remove points clearly outside view
        # 2. Then apply multi-view voting to remove inconsistent/erroneous points
        h, w = video_point_map[0].shape[:2]

        all_point_map_indices = np.where(np.ones_like(video_point_map_mask, dtype=bool))
        all_point_map_indices = np.stack(all_point_map_indices, axis=-1)
        all_point_map_indices = all_point_map_indices[video_point_map_mask.reshape(-1)]
        # Stage 1: First frame frustum filtering
        print(f"\n=== Stage 1: First frame frustum filtering ===")
        filtered_points, filtered_colors, filtered_masks = self.filter_by_first_frame_frustum_with_depth(
            video_point_map,
            video_point_map_mask,
            predicted_extrinsics[0],  # Use HunyuanWorld-Mirror's c2w prediction
            predicted_intrinsics[0],  # Use HunyuanWorld-Mirror's intrinsic prediction
            (h, w),
            video_pcd.colors
        )
        all_point_map_indices = all_point_map_indices[filtered_masks]

        # Stage 2: Multi-view voting (optional, enabled by default for better quality)
        use_multiview_voting = kwargs.get("use_multiview_voting", True)
        if use_multiview_voting and len(video_point_map) > 1:
            # With per-point visible view normalization, we can use higher thresholds
            # vote_threshold: fraction of VISIBLE views (for each point) that must give positive votes
            # Example with vote_threshold=0.5:
            #   - Point visible in 8 views needs 4 positive votes
            #   - Point visible in 2 views needs 1 positive vote
            # Higher threshold (e.g., 0.5-0.7) is more conservative (fewer points, higher quality)
            # Lower threshold (e.g., 0.3-0.4) is more permissive (more points, may include some errors)
            vote_threshold = kwargs.get("vote_threshold", 0.1)
            depth_tolerance = kwargs.get("depth_tolerance", 0.3)  # Relative depth tolerance (30%)

            print(f"\n=== Stage 2: Multi-view voting filter ===")
            print(f"Vote threshold: {vote_threshold}, Depth tolerance: {depth_tolerance}")

            # Apply multiview voting to the already-filtered points from stage 1
            filtered_points, filtered_colors, filtered_masks = self.filter_by_multiview_voting(
                filtered_points,
                filtered_colors,
                video_point_map,
                video_point_map_mask,
                predicted_extrinsics,
                predicted_intrinsics,
                (h, w),
                vote_threshold=vote_threshold,
                depth_tolerance=depth_tolerance
            )
            all_point_map_indices = all_point_map_indices[filtered_masks]

        filtered_points, norm_params = self.normalize_point_cloud(filtered_points,
                                                                  camera_forward=predicted_extrinsics[0, :3, 2])

        filtered_point_map = np.zeros_like(video_point_map)
        filtered_point_map_mask = np.zeros_like(video_point_map_mask, dtype=bool)
        filtered_point_map[all_point_map_indices[:, 0],
                            all_point_map_indices[:, 1],
                            all_point_map_indices[:, 2]] = filtered_points
        filtered_point_map_mask[all_point_map_indices[:, 0],
                                    all_point_map_indices[:, 1],
                                    all_point_map_indices[:, 2]] = True

        # Save filtered point cloud
        pcd_filtered = o3d.geometry.PointCloud()
        pcd_filtered.points = o3d.utility.Vector3dVector(filtered_points)
        if filtered_colors is not None:
            pcd_filtered.colors = o3d.utility.Vector3dVector(filtered_colors)
        o3d.io.write_point_cloud(os.path.join(save_folder, "video_first_frame_frustum_pcd.ply"), pcd_filtered)

        output_dict = {
            "extrinsics": predicted_extrinsics,
            "intrinsics": predicted_intrinsics,
            "point_map": filtered_point_map,
            "point_map_mask": filtered_point_map_mask,
            "all_viewpoints": filtered_points,
            "normalization": norm_params,
            "original_images": video_images,
        }

        trellis_extrinsics = []
        trellis_intrinsics = []
        trellis_cam_positions = []
        for i in range(len(video_images)):
            trellis_extrinsic, trellis_intrinsic, trellis_cam_pos = self.to_trellis_camera(
                np.linalg.inv(predicted_extrinsics[i]),
                predicted_intrinsics[i],
                norm_params,
                (h, w),
            )
            trellis_extrinsics.append(trellis_extrinsic)
            trellis_intrinsics.append(trellis_intrinsic)
            trellis_cam_positions.append(trellis_cam_pos)
        trellis_extrinsics = np.stack(trellis_extrinsics, axis=0)
        trellis_intrinsics = np.stack(trellis_intrinsics, axis=0)
        trellis_cam_positions = np.stack(trellis_cam_positions, axis=0)
        output_dict["trellis_extrinsic"] = trellis_extrinsics
        output_dict["trellis_intrinsic"] = trellis_intrinsics
        output_dict["trellis_camera_position"] = trellis_cam_positions
        # save trellis camera params
        trellis_camera_path = os.path.join(save_folder, "trellis_camera.json")
        trellis_camera_data = {
            "extrinsic": trellis_extrinsics.tolist(),
            "intrinsic": trellis_intrinsics.tolist(),
        }
        with open(trellis_camera_path, "w") as f:
            json.dump(trellis_camera_data, f, indent=2)

        num_views = int(kwargs.get("num_orbit_views", 121))
        rotation_angle_degrees = float(kwargs.get("orbit_rotation_angle_degrees", 0.0))
        
        if rotation_angle_degrees != 0:
            orbit_cameras = self.generate_orbit_cameras(
                trellis_extrinsics[0],
                num_views=num_views,
                rotation_angle_degrees=rotation_angle_degrees,
            )
            output_dict["orbit_cameras_extrinsic"] = orbit_cameras

            original_orbit_cameras = []
            for cam in orbit_cameras:
                orig_extrinsic, _ = self.from_trellis_camera(
                    cam,
                    output_dict["normalization"]
                )
                original_orbit_cameras.append(orig_extrinsic)
            output_dict["orbit_cameras_original_extrinsic"] = np.array(original_orbit_cameras)
        
        video_pc_info.update(output_dict)

        return video_pc_info