utils/vis/get_visualization_surfels.py

import sys
import json
import numpy as np
from PIL import Image


from torch.amp import autocast
import torch
import copy
from torch.nn import functional as F

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

sys.path.append("./extern/dust3r")
from dust3r.inference import inference, load_model
from dust3r.utils.image import load_images
from dust3r.image_pairs import make_pairs
from dust3r.cloud_opt import global_aligner, GlobalAlignerMode



def visualize_surfels(
        surfels, 
        draw_normals=False, 
        normal_scale=20, 
        disk_resolution=16, 
        disk_alpha=0.5
    ):
        """
        Visualize surfels as 2D disks oriented by their normals in 3D using matplotlib.

        Args:
            surfels (list of Surfel): Each Surfel has at least:
                - position: (x, y, z)
                - normal: (nx, ny, nz)
                - radius: scalar
                - color: (R, G, B) in [0..255] (optional)
            draw_normals (bool): If True, draws the surfel normals as quiver arrows.
            normal_scale (float): Scale factor for the normal arrows.
            disk_resolution (int): Number of segments to approximate each disk.
            disk_alpha (float): Alpha (transparency) for the filled disks.
        """

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        
        # Prepare arrays for optional quiver (if draw_normals=True)
        positions = []
        normals = []

        # We'll accumulate 3D polygons in a list for Poly3DCollection
        polygons = []
        polygon_colors = []

        for s in surfels:
            # --- Extract surfel data ---
            
            position = s.position
            normal = s.normal
            radius = s.radius
            
            if isinstance(position, torch.Tensor):
                x, y, z = position.detach().cpu().numpy()
                nx, ny, nz = normal.detach().cpu().numpy()
                radius = radius.detach().cpu().numpy()
            else:
                x, y, z = position
                nx, ny, nz = normal
                radius = radius
            
            
            # Convert color from [0..255] to [0..1], or use default
            if s.color is None:
                color = (0.2, 0.6, 1.0)  # Light blue
            else:
                r, g, b = s.color
                color = (r/255.0, g/255.0, b/255.0)

            # --- Build local coordinate axes for the disk ---
            normal = np.array([nx, ny, nz], dtype=float)
            norm_len = np.linalg.norm(normal)
            # Skip degenerate normals to avoid nan
            if norm_len < 1e-12:
                continue
            normal /= norm_len

            # Pick an 'up' vector that is not too close to the normal
            # so we can build a tangent plane
            up = np.array([0, 0, 1], dtype=float)
            if abs(normal.dot(up)) > 0.9:
                up = np.array([0, 1, 0], dtype=float)

            # xAxis = normal x up
            xAxis = np.cross(normal, up)
            xAxis /= np.linalg.norm(xAxis)
            # yAxis = normal x xAxis
            yAxis = np.cross(normal, xAxis)
            yAxis /= np.linalg.norm(yAxis)

            # --- Create a circle of 'disk_resolution' segments in local 2D coords ---
            angles = np.linspace(0, 2*np.pi, disk_resolution, endpoint=False)
            circle_points_3d = []
            for theta in angles:
                # local 2D circle: (r*cosθ, r*sinθ)
                px = radius * np.cos(theta)
                py = radius * np.sin(theta)

                # transform to 3D world space: position + px*xAxis + py*yAxis
                world_pt = np.array([x, y, z]) + px * xAxis + py * yAxis
                circle_points_3d.append(world_pt)

            # We have a list of [x, y, z]. For a filled polygon, Poly3DCollection
            # wants them as a single Nx3 array.
            circle_points_3d = np.array(circle_points_3d)
            polygons.append(circle_points_3d)
            polygon_colors.append(color)

            # Collect positions and normals for quiver (if used)
            positions.append([x, y, z])
            normals.append(normal)

        # --- Draw the disks as polygons ---
        poly_collection = Poly3DCollection(
            polygons,
            facecolors=polygon_colors,
            edgecolors='k',  # black edge
            linewidths=0.5,
            alpha=disk_alpha
        )
        ax.add_collection3d(poly_collection)

        # --- Optionally draw normal vectors (quiver) ---
        if draw_normals and len(positions) > 0:
            X = [p[0] for p in positions]
            Y = [p[1] for p in positions]
            Z = [p[2] for p in positions]

            Nx = [n[0] for n in normals]
            Ny = [n[1] for n in normals]
            Nz = [n[2] for n in normals]

            # Note: If your scene is large, you may want to increase `length`.
            ax.quiver(
                X, Y, Z, 
                Nx, Ny, Nz, 
                length=normal_scale, 
                color='red', 
                normalize=True
            )

        # --- Axis labels, aspect ratio, etc. ---
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        try:
            ax.set_box_aspect((1, 1, 1))
        except AttributeError:
            pass  # older MPL versions

        plt.title("Surfels as Disks (Oriented by Normal)")
        plt.show()
    


def visualize_pointcloud(
        points,
        colors=None,
        title='Point Cloud',
        point_size=1,
        alpha=1.0,
        bg_color=(240/255, 223/255, 223/255)  # 新增参数，默认白色 (1,1,1)
    ):
    """
    可视化3D点云，同时支持每个点的RGB或RGBA颜色，并保证x, y, z三个轴等比例缩放。
    
    参数
    ----------
    points : np.ndarray 或 torch.Tensor
        形状为 [N, 3] 的数组或张量，每行表示一个3D点 (x, y, z)。
    colors : None, str, 或 np.ndarray
        - 如果为 None，则使用默认颜色 'blue'。
        - 如果为字符串，则所有点均使用该颜色。
        - 如果为数组，则形状应为 [N, 3] 或 [N, 4]，表示每个点的颜色，值的范围应为 [0, 1]（若为浮点数）。
    title : str, 可选
        图像标题，默认 'Point Cloud'。
    point_size : float, 可选
        点的大小，默认 1。
    alpha : float, 可选
        点的整体透明度，默认 1.0。
    bg_color : tuple, 可选
        背景颜色，格式为 (r, g, b)，每个值的范围为 [0, 1]，默认为白色 (1.0, 1.0, 1.0)。

    示例
    --------
    >>> import numpy as np
    >>> pts = np.random.rand(1000, 3)
    >>> cols = np.random.rand(1000, 3)
    >>> visualize_pointcloud(pts, colors=cols, title="随机点云", bg_color=(0.2, 0.2, 0.3))
    """
    
    # 如果是 Torch 张量，则转换为 NumPy 数组
    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()
    if isinstance(colors, torch.Tensor):
        colors = colors.detach().cpu().numpy()

    # 如果点云或颜色数据维度过高，则展平
    if len(points.shape) > 2:
        points = points.reshape(-1, 3)
    if colors is not None and isinstance(colors, np.ndarray) and len(colors.shape) > 2:
        colors = colors.reshape(-1, colors.shape[-1])
        
    # 验证点云形状
    if points.shape[1] != 3:
        raise ValueError("`points` array must have shape [N, 3].")
        
    # 处理颜色参数
    if colors is None:
        colors = 'blue'
    elif isinstance(colors, np.ndarray):
        colors = np.asarray(colors)
        if colors.shape[0] != points.shape[0]:
            raise ValueError("Colors array length must match the number of points.")
        if colors.shape[1] not in [3, 4]:
            raise ValueError("Colors array must have shape [N, 3] or [N, 4].")
    
    # 验证背景颜色参数
    if not isinstance(bg_color, tuple) or len(bg_color) != 3:
        raise ValueError("Background color must be a tuple of (r, g, b) with values between 0 and 1.")
    
    # 提取坐标
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]
    
    # 创建图像，并设置自定义背景颜色
    fig = plt.figure(figsize=(8, 6), facecolor=bg_color)
    ax = fig.add_subplot(111, projection='3d')
    ax.set_facecolor(bg_color)
    
    # 绘制散点图
    ax.scatter(x, y, z, c=colors, s=point_size, alpha=alpha)
    
    # 设置等比例缩放
    max_range = np.array([x.max() - x.min(), 
                          y.max() - y.min(),
                          z.max() - z.min()]).max() / 2.0
    mid_x = (x.max() + x.min()) * 0.5
    mid_y = (y.max() + y.min()) * 0.5
    mid_z = (z.max() + z.min()) * 0.5
    ax.set_xlim(mid_x - max_range, mid_x + max_range)
    ax.set_ylim(mid_y - max_range, mid_y + max_range)
    ax.set_zlim(mid_z - max_range, mid_z + max_range)
    
    # 隐藏刻度和标签
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_zlabel('')
    ax.grid(False)
    
    # 隐藏3D坐标轴的面板（pane）来去除轴的显示
    ax.xaxis.pane.set_visible(False)
    ax.yaxis.pane.set_visible(False)
    ax.zaxis.pane.set_visible(False)
    
    # 设置标题（如果需要显示标题）
    ax.set_title(title)
    
    plt.tight_layout()
    plt.show()

# def visualize_pointcloud(
#         points,
#         colors=None,
#         title='Point Cloud',
#         point_size=1,
#         alpha=1.0
#     ):
#     """
#     可视化3D点云，同时支持每个点的RGB或RGBA颜色，并保证x, y, z三个轴等比例缩放。
    
#     参数
#     ----------
#     points : np.ndarray 或 torch.Tensor
#         形状为 [N, 3] 的数组或张量，每行表示一个3D点 (x, y, z)。
#     colors : None, str, 或 np.ndarray
#         - 如果为 None，则使用默认颜色 'blue'。
#         - 如果为字符串，则所有点均使用该颜色。
#         - 如果为数组，则形状应为 [N, 3] 或 [N, 4]，表示每个点的颜色，值的范围应为 [0, 1]（若为浮点数）。
#     title : str, 可选
#         图像标题，默认 'Point Cloud'。
#     point_size : float, 可选
#         点的大小，默认 1。
#     alpha : float, 可选
#         点的整体透明度，默认 1.0。

#     示例
#     --------
#     >>> import numpy as np
#     >>> pts = np.random.rand(1000, 3)
#     >>> cols = np.random.rand(1000, 3)
#     >>> visualize_pointcloud(pts, colors=cols, title="随机点云")
#     """
    
#     # 如果是 Torch 张量，则转换为 NumPy 数组
#     if isinstance(points, torch.Tensor):
#         points = points.detach().cpu().numpy()
#     if isinstance(colors, torch.Tensor):
#         colors = colors.detach().cpu().numpy()

#     # 如果点云或颜色数据维度过高，则展平
#     if len(points.shape) > 2:
#         points = points.reshape(-1, 3)
#     if colors is not None and isinstance(colors, np.ndarray) and len(colors.shape) > 2:
#         colors = colors.reshape(-1, colors.shape[-1])
        
#     # 验证点云形状
#     if points.shape[1] != 3:
#         raise ValueError("`points` array must have shape [N, 3].")
        
#     # 处理颜色参数
#     if colors is None:
#         colors = 'blue'
#     elif isinstance(colors, np.ndarray):
#         colors = np.asarray(colors)
#         if colors.shape[0] != points.shape[0]:
#             raise ValueError("Colors array length must match the number of points.")
#         if colors.shape[1] not in [3, 4]:
#             raise ValueError("Colors array must have shape [N, 3] or [N, 4].")
    
#     # 提取坐标
#     x = points[:, 0]
#     y = points[:, 1]
#     z = points[:, 2]
    
#     # 创建图像，并设置背景为白色
#     fig = plt.figure(figsize=(8, 6), facecolor='white')
#     ax = fig.add_subplot(111, projection='3d')
#     ax.set_facecolor('white')
    
#     # 绘制散点图
#     ax.scatter(x, y, z, c=colors, s=point_size, alpha=alpha)
    
#     # 设置等比例缩放
#     max_range = np.array([x.max() - x.min(), 
#                           y.max() - y.min(),
#                           z.max() - z.min()]).max() / 2.0
#     mid_x = (x.max() + x.min()) * 0.5
#     mid_y = (y.max() + y.min()) * 0.5
#     mid_z = (z.max() + z.min()) * 0.5
#     ax.set_xlim(mid_x - max_range, mid_x + max_range)
#     ax.set_ylim(mid_y - max_range, mid_y + max_range)
#     ax.set_zlim(mid_z - max_range, mid_z + max_range)
    
#     # 隐藏刻度和标签
#     ax.set_xticks([])
#     ax.set_yticks([])
#     ax.set_zticks([])
#     ax.set_xlabel('')
#     ax.set_ylabel('')
#     ax.set_zlabel('')
#     ax.grid(False)
    
#     # 隐藏3D坐标轴的面板（pane）来去除轴的显示
#     ax.xaxis.pane.set_visible(False)
#     ax.yaxis.pane.set_visible(False)
#     ax.zaxis.pane.set_visible(False)
    
#     # 设置标题（如果需要显示标题）
#     ax.set_title(title)
    
#     plt.tight_layout()
#     plt.show()
class Surfel:
    def __init__(self, position, normal, radius=1.0, color=None):
        """
        position: (x, y, z)
        normal:   (nx, ny, nz)
        radius:   scalar
        color:    (r, g, b) or None
        """
        self.position = position
        self.normal = normal
        self.radius = radius
        self.color = color

    def __repr__(self):
        return (f"Surfel(position={self.position}, "
                f"normal={self.normal}, radius={self.radius}, "
                f"color={self.color})")




class Octree:
    def __init__(self, points, indices=None, bbox=None, max_points=10):
        """
        构建八叉树：
          - points: 所有点的 numpy 数组，形状为 (N, 3)
          - indices: 当前节点中点的索引列表
          - bbox: 当前节点的包围盒，形式为 (center, half_size)，其中半径为正方体半边长
          - max_points: 叶子节点允许的最大点数
        """
        self.points = points
        if indices is None:
            indices = np.arange(points.shape[0])
        self.indices = indices

        # 如果没有给定包围盒，则计算所有点的包围盒，保证是一个正方体
        if bbox is None:
            min_bound = points.min(axis=0)
            max_bound = points.max(axis=0)
            center = (min_bound + max_bound) / 2
            half_size = np.max(max_bound - min_bound) / 2
            bbox = (center, half_size)
        self.center, self.half_size = bbox

        self.children = []  # 存储子节点
        self.max_points = max_points

        if len(self.indices) > self.max_points:
            self.subdivide()

    def subdivide(self):
        """将当前节点划分为8个子节点"""
        cx, cy, cz = self.center
        hs = self.half_size / 2
        # 八个象限的偏移量
        offsets = np.array([[dx, dy, dz] for dx in (-hs, hs) 
                                       for dy in (-hs, hs) 
                                       for dz in (-hs, hs)])
        for offset in offsets:
            child_center = self.center + offset
            child_indices = []
            # 检查每个点是否在子节点的包围盒内
            for idx in self.indices:
                p = self.points[idx]
                if np.all(np.abs(p - child_center) <= hs):
                    child_indices.append(idx)
            child_indices = np.array(child_indices)
            if len(child_indices) > 0:
                child = Octree(self.points, indices=child_indices, bbox=(child_center, hs), max_points=self.max_points)
                self.children.append(child)
        # 划分后，内部节点不再直接保存点索引
        self.indices = None

    def sphere_intersects_node(self, center, r):
        """
        判断以center为球心, r为半径的球是否与当前节点的轴对齐包围盒相交。
        算法：计算球心到盒子的距离（只考虑超出盒子边界的部分），若小于r，则相交。
        """
        diff = np.abs(center - self.center)
        max_diff = diff - self.half_size
        max_diff = np.maximum(max_diff, 0)
        dist_sq = np.sum(max_diff**2)
        return dist_sq <= r*r

    def query_ball_point(self, point, r):
        """
        查询距离给定点 point 小于 r 的所有点索引。
        """
        results = []
        if not self.sphere_intersects_node(point, r):
            return results
        # 如果当前节点没有子节点，则为叶子节点
        if len(self.children) == 0:
            if self.indices is not None:
                for idx in self.indices:
                    if np.linalg.norm(self.points[idx] - point) <= r:
                        results.append(idx)
            return results
        else:
            for child in self.children:
                results.extend(child.query_ball_point(point, r))
            return results
        

def estimate_normal_from_pointmap(pointmap: torch.Tensor) -> torch.Tensor:
    """
    Estimate surface normals from a 3D point map by computing cross products of
    neighboring points, using PyTorch tensors.

    Parameters
    ----------
    pointmap : torch.Tensor
        A PyTorch tensor of shape [H, W, 3] containing 3D points in camera coordinates.
        Each point is represented as (X, Y, Z). This tensor can be on CPU or GPU.

    Returns
    -------
    torch.Tensor
        A PyTorch tensor of shape [H, W, 3] containing estimated surface normals.
        Each normal is a unit vector (X, Y, Z).
        Points where normals cannot be computed (e.g. boundaries) will be zero vectors.
    """
    # pointmap is shape (H, W, 3)
    h, w = pointmap.shape[:2]
    device = pointmap.device  # Keep the device (CPU/GPU) consistent
    dtype = pointmap.dtype
    
    # Initialize the normal map
    normal_map = torch.zeros((h, w, 3), device=device, dtype=dtype)
    
    for y in range(h):
        for x in range(w):
            # Check if neighbors are within bounds
            if x+1 >= w or y+1 >= h:
                continue
            
            p_center = pointmap[y, x]
            p_right  = pointmap[y, x+1]
            p_down   = pointmap[y+1, x]
            
            # Compute vectors
            v1 = p_right - p_center
            v2 = p_down - p_center
            
            v1 = v1 / torch.linalg.norm(v1)
            v2 = v2 / torch.linalg.norm(v2)
            
            # Cross product in camera coordinates
            n_c = torch.cross(v1, v2)
            # n_c *= 1e10
            
            # Compute norm of the normal vector
            norm_len = torch.linalg.norm(n_c)
            
            if norm_len < 1e-8:
                continue
            
            # Normalize and store
            normal_map[y, x] = n_c / norm_len
    
    return normal_map


def load_multiple_images(image_names, image_size=512, dtype=torch.float32):
    images = load_images(image_names, size=image_size, force_1024=True, dtype=dtype)
    img_ori = (images[0]['img_ori'].squeeze(0).permute(1,2,0)+1.)/2. # Just for reference
    return images, img_ori


def load_initial_images(image_name):
    images = load_images([image_name], size=512, force_1024=True)
    img_ori = (images[0]['img_ori'].squeeze(0).permute(1,2,0)+1.)/2. # [H, W, 3], range [0,1]
    if len(images) == 1:
        images = [images[0], copy.deepcopy(images[0])]
        images[1]['idx'] = 1
    return images, img_ori


def merge_surfels(
    new_surfels: list,
    current_timestamp: str,
    existing_surfels: list,
    existing_surfel_to_timestamp: dict,
    position_threshold: float = 0.025,
    normal_threshold: float = 0.7,
    max_points_per_node: int = 10  # 八叉树叶子节点允许的最大点数
):
    """
    将新的 surfel 合并到已有 surfel 列表中，使用八叉树来加速空间查找。
    
    Args:
        new_surfels (list[Surfel]): 待合并的新 surfel 列表。
        current_timestamp (str): 当前的时间戳。
        existing_surfels (list[Surfel]): 已存在的 surfel 列表。
        existing_surfel_to_timestamp (dict): 每个 surfel 索引到时间戳的映射。
        position_threshold (float): 判断两个 surfel 空间距离是否足够近的阈值。
        normal_threshold (float): 判断两个 surfel 法向是否对齐的阈值。
        max_points_per_node (int): 构建八叉树时，每个叶子节点最大允许的点数。
        
    Returns:
        (list[Surfel], dict): 
            - 未能匹配的 surfel 列表，需要追加到已有 surfel 列表中。
            - 更新后的 existing_surfel_to_timestamp 映射。
    """
    # 安全检查
    assert len(existing_surfels) == len(existing_surfel_to_timestamp), (
        "existing_surfels 和 existing_surfel_to_timestamp 长度不匹配。"
    )
    
    # 构造已有 surfel 的位置和法向数组
    positions = np.array([s.position for s in existing_surfels])  # Shape: (N, 3)
    normals = np.array([s.normal for s in existing_surfels])      # Shape: (N, 3)
    
    
    # 用于存储未匹配到已有 surfel 的新 surfel
    filtered_surfels = []
    
    merge_count = 0
    for new_surfel in new_surfels:
        is_merged = False
        for idx in range(len(positions)):
            if np.linalg.norm(positions[idx] - new_surfel.position) < position_threshold:
                if np.dot(normals[idx], new_surfel.normal) > normal_threshold:
                    existing_surfel_to_timestamp[idx].append(current_timestamp)
                    is_merged = True
                    merge_count += 1
                    break
        
        
        if not is_merged:
            filtered_surfels.append(new_surfel)
    
    # 返回未匹配的 surfel 列表及更新后的时间戳映射
    print(f"merge_count: {merge_count}")
    return filtered_surfels, existing_surfel_to_timestamp

def pointmap_to_surfels(pointmap: torch.Tensor,
                        focal_lengths: torch.Tensor,
                        depth_map: torch.Tensor,
                        poses: torch.Tensor, # shape: (4, 4)
                        radius_scale: float = 0.5,
                        depth_threshold: float = 1.0,
                        estimate_normals: bool = True):
    surfels = []
    if len(focal_lengths) == 2:
        focal_lengths = torch.mean(focal_lengths, dim=0)
    H, W = pointmap.shape[:2]
    # 1) Estimate normals
    if estimate_normals:
        normal_map = estimate_normal_from_pointmap(pointmap)
    else:
        normal_map = torch.zeros_like(pointmap)
    depth_remove_count = 0
    for v in range(H-1):
        for u in range(W-1):
            if depth_map[v, u] > depth_threshold:
                depth_remove_count += 1
                continue
            position = pointmap[v, u].detach().cpu().numpy() # in global coords
            normal = normal_map[v, u].detach().cpu().numpy() # in global coords
            depth = depth_map[v, u].detach().cpu().numpy() # in local coords
            view_direction = position - poses[0:3, 3].detach().cpu().numpy()
            view_direction = view_direction / np.linalg.norm(view_direction)
            if np.dot(view_direction, normal) < 0:
                normal = -normal
            adjustment_value = 0.2 + 0.8 * np.abs(np.dot(view_direction, normal))
            radius = (radius_scale * depth/focal_lengths/adjustment_value).detach().cpu().numpy()
            surfels.append(Surfel(position, normal, radius))
    print(f"depth_remove_count: {depth_remove_count}")
    return surfels



def run_dust3r(input_images,
               dust3r,
               batch_size = 1,
               niter = 1000,
               lr = 0.01,
               schedule = 'linear',
               clean_pc = False,
               focal_lengths = None,
               poses = None,
               device = 'cuda',
               background_mask = None,
               use_amp = False  # <<< AMP CHANGE: add a flag to enable/disable AMP
               ):

    # We wrap the entire inference and alignment in autocast so that
    # forward passes and any internal backward passes happen in mixed precision.
    with autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
        pairs = make_pairs(input_images, scene_graph='complete', prefilter=None, symmetrize=True)
        output = inference(pairs, dust3r, device, batch_size=batch_size)

        mode = GlobalAlignerMode.PointCloudDifferentFocalOptimizer 
        scene = global_aligner(output, device=device, mode=mode)
        if focal_lengths is not None:
            scene.preset_focal(focal_lengths)
        if poses is not None:
            scene.preset_pose(poses)
        if mode == GlobalAlignerMode.PointCloudDifferentFocalOptimizer:
            # Depending on how dust3r internally does optimization,
            # it may or may not require gradient scaling.
            # If you need it, you can do something more manual with GradScaler.
            loss = scene.compute_global_alignment(init='mst', niter=niter, schedule=schedule, lr=lr)
        else:
            loss = None

        # If you want to clean up the pointcloud after alignment
        if clean_pc:
            scene = scene.clean_pointcloud()

    return scene, loss


if __name__ == "__main__":
    load_image_size = 512
    load_dtype = torch.float16
    device = 'cuda'
    model_path = "checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"
    selected_frame_paths = ["assets/jesus/jesus_0.jpg",
                            "assets/jesus/jesus_1.jpg",
                            "assets/jesus/jesus_2.jpg"
                            ]
    
    # pil_image = Image.open("./assets/radcliffe_camera_bg.png").resize((512, 288))
    # r, g, b, a = pil_image.split()
    # background_mask = a
    # background_mask = (1 - torch.tensor(np.array(background_mask))).unsqueeze(0).repeat(2, 1, 1).bool()
    
    all_surfels = []
    surfel_to_timestamp = {}
    
    dust3r = load_model(model_path, device=device)
    dust3r.eval()
    dust3r = dust3r.to(device)
    dust3r = dust3r.half()
    if len(selected_frame_paths) == 1:
        selected_frame_paths = selected_frame_paths * 2
    frame_images, frame_img_ori = load_multiple_images(selected_frame_paths, 
                                                       image_size=load_image_size, 
                                                       dtype=load_dtype)
    
    
    scene, loss = run_dust3r(frame_images, dust3r, device=device, use_amp=True)
    # --- 1) Extract outputs ---
    # pointcloud shape: [N, H, W, 3]
    shrink_factor = 0.15
    pointcloud = torch.stack(scene.get_pts3d())
    # poses shape: [N, 4, 4]
    # optimized_poses = scene.get_im_poses()
    # focal_lengths shape: [N]
    focal_lengths = scene.get_focals()
    

    # adjustion_transformation_matrix = SpatialConstructor.estimate_pose_alignment(optimized_poses, original_camera_poses) # optimized_poses -> original_camera_poses matrix
    # adjusted_optimized_poses = adjustion_transformation_matrix @ optimized_poses


    # --- 2) Resize pointcloud ---
    # Permute for resizing -> [N, 3, H, W]
    pointcloud = pointcloud.permute(0, 3, 1, 2)

    # Resize using bilinear interpolation
    pointcloud = F.interpolate(
        pointcloud, 
        scale_factor=shrink_factor, 
        mode='bilinear'
    )
    # Permute back -> [N, H', W', 3]
    pointcloud = pointcloud.permute(0, 2, 3, 1)[-1:]
    # transform pointcloud
    # pointcloud = torch.stack([SpatialConstructor.transform_pointmap(pointcloud[i], adjustion_transformation_matrix) for i in range(pointcloud.shape[0])])



    rgbs = scene.imgs
    rgbs = torch.tensor(np.array(rgbs))
    rgbs = rgbs.permute(0, 3, 1, 2)
    rgbs = F.interpolate(rgbs, scale_factor=shrink_factor, mode='bilinear')
    rgbs = rgbs.permute(0, 2, 3, 1)[-1:]
    visualize_pointcloud(pointcloud, rgbs, point_size=4)
    # --- 3) Resize depth map ---
    # depth_map shape: [N, H, W]
    depth_map = torch.stack(scene.get_depthmaps())

    # Add channel dimension -> [N, 1, H, W]
    depth_map = depth_map.unsqueeze(1)

    depth_map = F.interpolate(
        depth_map, 
        scale_factor=shrink_factor, 
        mode='bilinear'
    )
    
    poses = scene.get_im_poses()[-1:]

    # Remove channel dimension -> [N, H', W']
    depth_map = depth_map.squeeze(1)[-1:]

    for frame_idx in range(len(pointcloud)):
        # if frame_idx > 1:
        #     break
    # Create surfels for the current frame
        surfels = pointmap_to_surfels(
            pointmap=pointcloud[frame_idx],
            focal_lengths=focal_lengths[frame_idx] * shrink_factor,
            depth_map=depth_map[frame_idx],
            poses=poses[frame_idx],
            estimate_normals=True,
            radius_scale=0.5,
            depth_threshold=0.48
        )

        # Merge with existing surfels if not the first frame
        if frame_idx > 0:
            surfels, surfel_to_timestamp = merge_surfels(
                new_surfels=surfels,
                current_timestamp=frame_idx,
                existing_surfels=all_surfels,
                existing_surfel_to_timestamp=surfel_to_timestamp,
                position_threshold=0.01,
                normal_threshold=0.7
                    )
        
                    # Update timestamp mapping
        num_surfels = len(surfels)
        surfel_start_index = len(all_surfels)
        for surfel_index in range(num_surfels):
            # Each newly created surfel gets mapped to this frame index
            # surfel_to_timestamp[surfel_start_index + surfel_index] = [frame_idx]
            surfel_to_timestamp[surfel_start_index + surfel_index] = [2]
        all_surfels.extend(surfels)
    
    positions = np.array([s.position for s in all_surfels], dtype=np.float32)
    normals   = np.array([s.normal   for s in all_surfels], dtype=np.float32)
    radii     = np.array([s.radius   for s in all_surfels], dtype=np.float32)
    colors    = np.array([s.color    for s in all_surfels], dtype=np.float32)
    
    
    visualize_surfels(all_surfels)
    

    # np.savez(f"./surfels_added_first2.npz",
    #         positions=positions,
    #         normals=normals,
    #         radii=radii,
    #         colors=colors)
    # with open("surfel_to_timestamp_first2.json", "w") as f:
    #     json.dump(surfel_to_timestamp, f)
    
    
    np.savez(f"./surfels_added_only3.npz",
            positions=positions,
            normals=normals,
            radii=radii,
            colors=colors)
    with open("surfel_to_timestamp_only3.json", "w") as f:
        json.dump(surfel_to_timestamp, f)
    
    stop = 1