gaussian_viewer.py

"""Unified Gaussian viewer.

Merges `gaussian_optim_viewer.py` (light, supports live training attach +
client/server) with `local_viewer_decomp.py` (heavy decomposition editor).

Editing mode is derived from `ViewerMode`:
  - LOCAL  -> editing enabled. Loads three Gaussian sets, builds the
    ComplexEdit widget, enables the Depth render mode, and supports
    interactive multi-view re-optimization driven by selection-box edits.
  - CLIENT / SERVER -> lightweight optim-viewer behavior. No editing widgets.
    Intended for live-attach to a remote training session.
"""

import os
import threading
import numpy as np
from OpenGL.GL import *
from threading import Lock
from argparse import ArgumentParser

from gs_viewer import Viewer
from gs_viewer.viewer_types import ViewerMode
from gs_viewer.widgets.image import TorchImage
from gs_viewer.widgets.cameras.fps import FPSCamera
from gs_viewer.widgets.monitor import PerformanceMonitor
from gs_viewer.widgets.ellipsoid_viewer import EllipsoidViewer
from gs_viewer.widgets.pixel_inspector import PixelInspector

from imgui_bundle import imgui_ctx, imgui, ImVec2, ImVec4, imguizmo

gizmo = imguizmo.im_guizmo
Matrix16 = gizmo.Matrix16


class Dummy(object):
    pass


class GaussianViewer(Viewer):
    def __init__(self, mode: ViewerMode, poses=None, res_x: int = None, res_y: int = None,
                 *, editing: bool = None):
        # Must be set before super().__init__: the base Viewer.__init__ calls
        # self.import_server_modules(), which reads self.editing_enabled.
        # Default: editing follows LOCAL mode, but factories may override
        # (e.g. live-attach during training passes editing=False).
        self.editing_enabled = (mode is ViewerMode.LOCAL) if editing is None else editing
        super().__init__(mode)
        self.window_title = "Gaussian Viewer"
        self.gaussian_lock = Lock()

        # ---- shared state (always initialized) ----
        self.background_color = [0.0, 0.0, 0.0]
        self.gaussian_color = "DIFFUSE"
        self.exposure = 1.0
        self.gamma = 2.2
        self.scaling_modifier = 1.0
        self.opacity_modifier = 1.0
        self.colors_precomp = None
        self.use_shs = True

        self.enable_pixel_inspector = False
        self.mouse_pos = ImVec2(0, 0)
        self.albedo_pixel_col = "unavailable"
        self.shading_pixel_col = "unavailable"
        self.glossy_pixel_col = "unavailable"
        self.show_img_training = False

        # Optional gaussian sets - filled in by factories
        self.albedo_gaussians = None
        self.shading_gaussians = None
        self.residual_gaussians = None

        # Tracks which gaussian set is currently uploaded to the ellipsoid viewer
        # so we can re-upload when the user switches the gaussian_color selection.
        self._ellipsoid_uploaded_color = None

        # Camera / projection defaults (overridden in editing mode from poses)
        self.poses = poses
        self.res_x = res_x if res_x is not None else 1297
        self.res_y = res_y if res_y is not None else 840
        self.fov_y = 47 if poses is None else poses[0].FoVy
        self.fov_x = 47 if poses is None else poses[0].FoVx
        self.znear = 0.001
        self.zfar = 100.0

        # ---- editing-only state (only set up when editing_enabled) ----
        if self.editing_enabled:
            self._init_editing_state()

    def _init_editing_state(self):
        """State used by the ComplexEdit + multi-view re-optimization flow."""
        # ---- selection-box / re-optimization state ----
        self.is_optimizing = False
        self.can_display_edit = False
        self.edit_use_shading_depth = False
        self.edition_gaussians = None
        self.prev_albedo_gaussians_params = None
        self.sampled_views = []

        self.track_edit_zone = False
        self.edit_start_pos = None
        self.edit_end_pos = None
        self.lower_coordinate = None
        self.higher_coordinate = None

        self.opt = None  # lazily built OptimizationParams when reoptimization first runs

        # Gizmo / depth controls
        self.tool_id = 0
        self.tools = ["move", "rotate"]
        self.display_gizmo = True
        self.show_depth = False
        self.display_normalized_depth = False
        self._image_top_left = None  # set in show_gui's Point View block, used by the gizmo

        # 50x50 grid of pixel coordinates around the image center, used by
        # sample_close_viewpoints() to locate the focal scene point.
        nb = 50
        cx, cy = self.res_x // 2, self.res_y // 2
        sx = np.arange(cx - nb // 2, cx + nb // 2)
        sy = np.arange(cy - nb // 2, cy + nb // 2)
        grid_x, grid_y = np.meshgrid(sx, sy, indexing='ij')
        self.middle_coordinates = torch.from_numpy(
            np.column_stack((grid_x.ravel(), grid_y.ravel())).astype(np.int32))

    def import_server_modules(self):
        global torch
        import torch

        global GaussianModel
        from scene import GaussianModel

        global PipelineParams, ModelParams, OptimizationParams
        from arguments import PipelineParams, ModelParams, OptimizationParams

        global MiniCam
        from scene.cameras import MiniCam

        global render
        from gaussian_renderer import render

        global get_expon_lr_func
        from utils.general_utils import get_expon_lr_func

        if self.editing_enabled:
            global sample_hemisphere_local, look_at_batch_torch
            from utils.graphics_utils import sample_hemisphere_local, look_at_batch_torch

            global l1_loss, fused_ssim
            from utils.loss_utils import l1_loss
            from fused_ssim import fused_ssim

            global distCUDA2
            from simple_knn._C import distCUDA2

    @classmethod
    def from_ply(cls, model_path, iter, mode: ViewerMode, *, albedo_pc_name: str = None):
        params = cls._read_cfg_args(model_path)
        dataset, pipe = cls._dataset_pipe_from_cfg(params)

        viewer = cls(mode)

        viewer.dataset = dataset
        viewer.pipe = pipe

        bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
        viewer.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

        ply_dir = os.path.join(model_path, "point_cloud", f"iteration_{iter}")
        cls._load_gaussian_sets(viewer, ply_dir, dataset.sh_degree, albedo_pc_name)

        return viewer

    @staticmethod
    def _read_cfg_args(model_path: str) -> dict:
        """Parse the Namespace dump train.py writes to <model_path>/cfg_args."""
        with open(os.path.join(model_path, "cfg_args")) as f:
            text = f.read()
        text = text[text.index("(") + 1: text.rindex(")")]
        out = {}
        for token in text.split(","):
            token = token.strip()
            if "=" not in token:
                continue
            key, value = token.split("=", 1)
            out[key.strip()] = value.strip().strip("'").strip('"')
        return out

    @staticmethod
    def _dataset_pipe_from_cfg(params: dict):
        dataset = Dummy()
        dataset.white_background = params.get("white_background", "False") == "True"
        dataset.sh_degree = int(params.get("sh_degree", "0"))
        dataset.train_test_exp = params.get("train_test_exp", "False") == "True"

        pipe = Dummy()
        pipe.debug = params.get("debug", "False") == "True"
        pipe.antialiasing = params.get("antialiasing", "False") == "True"
        pipe.compute_cov3D_python = params.get("compute_cov3D_python", "False") == "True"
        pipe.convert_SHs_python = params.get("convert_SHs_python", "False") == "True"

        return dataset, pipe

    @staticmethod
    def _load_gaussian_sets(viewer, ply_dir: str, sh_degree: int, albedo_pc_name: str = None):
        albedo_name = (albedo_pc_name or "albedo") + ".ply"
        albedo_path = os.path.join(ply_dir, albedo_name)
        if os.path.isfile(albedo_path):
            viewer.albedo_gaussians = GaussianModel(sh_degree)
            viewer.albedo_gaussians.load_ply(albedo_path)
            print(f"Loaded albedo gaussians from {albedo_path}")

        shading_path = os.path.join(ply_dir, "shading.ply")
        if os.path.isfile(shading_path):
            viewer.shading_gaussians = GaussianModel(sh_degree)
            viewer.shading_gaussians.load_ply(shading_path)
            print(f"Loaded shading gaussians from {shading_path}")

        residual_path = os.path.join(ply_dir, "residual.safetensors")
        if os.path.isfile(residual_path):
            viewer.residual_gaussians = GaussianModel(sh_degree=3)
            viewer.residual_gaussians.load_safetensors(residual_path)
            print(f"Loaded residual gaussians from {residual_path}")

        if viewer.albedo_gaussians is None and viewer.shading_gaussians is None and viewer.residual_gaussians is None:
            raise FileNotFoundError(f"No gaussian set found under {ply_dir}.")

    @classmethod
    def from_gaussians(cls, dataset, pipe, mode: ViewerMode, *,
                       albedo_gaussians=None,
                       shading_gaussians=None,
                       residual_gaussians=None):
        """Live-attach factory used during optimization. Editing is disabled —
        the gaussian sets may be partially populated mid-training."""
        viewer = cls(mode, editing=False)
        viewer.dataset = dataset
        viewer.pipe = pipe
        viewer.albedo_gaussians = albedo_gaussians
        viewer.shading_gaussians = shading_gaussians
        viewer.residual_gaussians = residual_gaussians
        viewer.background = torch.tensor([0, 0, 0], dtype=torch.float32, device="cuda")
        return viewer

    @staticmethod
    def _ready(g):
        """True iff a GaussianModel is present AND has at least one point.

        During training a GaussianModel can be assigned to the viewer before
        it is populated (zero gaussians). Rasterizing it would crash, so all
        render-path dispatches go through this guard.
        """
        return g is not None and g.get_xyz.shape[0] > 0

    def set_albedo_gaussians(self, albedo_gaussians):
        self.albedo_gaussians = albedo_gaussians

    def set_shading_gaussians(self, shading_gaussians):
        self.shading_gaussians = shading_gaussians

    def set_residual_gaussians(self, residual_gaussians):
        self.residual_gaussians = residual_gaussians

    def _gaussians_for_ellipsoid(self, color: str):
        """Pick which GaussianModel to display in the ellipsoid viewer for a given render color.
        Returns (gaussians, use_raw_features_dc) or (None, _) if unavailable."""
        if color == "SHADING" and self.shading_gaussians is not None:
            return self.shading_gaussians, True
        if color == "RESIDUAL" and self.residual_gaussians is not None:
            return self.residual_gaussians, False
        if self.albedo_gaussians is not None:
            return self.albedo_gaussians, False
        return (self.shading_gaussians or self.residual_gaussians), False

    def _ensure_ellipsoid_upload(self):
        """Upload the gaussian set matching the current selection to the ellipsoid viewer,
        re-uploading only when the selection changes."""
        gaussians, use_raw = self._gaussians_for_ellipsoid(self.gaussian_color)
        if gaussians is None:
            return
        if self._ellipsoid_uploaded_color == self.gaussian_color and \
           self.ellipsoid_viewer.num_gaussians is not None:
            return
        colors = gaussians.get_raw_features_dc if use_raw else gaussians.get_features_dc
        self.ellipsoid_viewer.upload(
            gaussians.get_xyz.detach().cpu().numpy(),
            gaussians.get_rotation.detach().cpu().numpy(),
            gaussians.get_scaling.detach().cpu().numpy(),
            gaussians.get_opacity.detach().cpu().numpy(),
            colors.detach().reshape(-1, 3).cpu().numpy(),
        )
        self._ellipsoid_uploaded_color = self.gaussian_color

    def _mouse_pos_in_bounds(self, bounds: list) -> bool:
        return ((self.mouse_pos.x >= 0 and self.mouse_pos.x < bounds[0]) and
                (self.mouse_pos.y >= 0 and self.mouse_pos.y < bounds[1]))

    def create_widgets(self):
        # ---- always-on widgets ----
        self.camera = FPSCamera(self.mode, self.res_x, self.res_y, self.fov_y, self.znear, self.zfar)
        self.point_view = TorchImage(self.mode)
        self.ellipsoid_viewer = EllipsoidViewer(self.mode)
        self.monitor = PerformanceMonitor(self.mode, ["Render"], add_other=False)
        self.pixel_inspector = PixelInspector()

        self.render_modes = ["Splats", "Ellipsoids"]
        self.render_mode = 0

        # ---- editing-only widgets ----
        if self.editing_enabled:
            self._create_editing_widgets()

    def _create_editing_widgets(self):
        """Build the ComplexEdit widget and extend the render-mode list."""
        from gs_viewer.widgets.complex_edit import ComplexEdit

        self.complex_edit = ComplexEdit(res_x=self.res_x, res_y=self.res_y, mode=self.mode)

    def get_positions_under_selection(self):
        """Return the world-space xyz of the shading gaussians visible under the selection box."""
        cam = self.camera
        world_to_view = torch.from_numpy(cam.to_camera).cuda().transpose(0, 1)
        full_proj_transform = torch.from_numpy(cam.full_projection).cuda().transpose(0, 1)
        mini = MiniCam(cam.res_x, cam.res_y, cam.fov_y, cam.fov_x,
                       cam.z_near, cam.z_far, world_to_view, full_proj_transform)

        primary = self.shading_gaussians or self.albedo_gaussians
        rendered = render(mini, primary, self.pipe, self.background,
                          scaling_modifier=self.scaling_modifier,
                          override_color=self.colors_precomp,
                          opacity_modifier=self.opacity_modifier,
                          colors_activation=False)
        mainIds = rendered["mainGaussID"].squeeze(0)

        sx = np.arange(self.lower_coordinate[0], self.higher_coordinate[0])
        sy = np.arange(self.lower_coordinate[1], self.higher_coordinate[1])
        grid_x, grid_y = np.meshgrid(sx, sy, indexing='ij')
        coords = np.column_stack((grid_x.ravel(), grid_y.ravel()))
        return primary.get_xyz[mainIds[coords[:, 1], coords[:, 0]]]

    def sample_close_viewpoints(self, camera, num_samples: int, max_deg: float):
        """Sample num_samples MiniCams near the given viewpoint, all looking at the same scene point."""
        if num_samples == 0:
            return [camera]

        sampled_points = sample_hemisphere_local(n=max_deg, size=num_samples)
        global_samples = ((sampled_points + torch.Tensor([0, 0, 1]).to("cuda")) - camera.world_view_transform[3, :3]) @ camera.world_view_transform[:3, :3].T

        rendered = render(camera, self.albedo_gaussians, self.pipe, self.background,
                          override_color=self.colors_precomp)
        mainIds = rendered["mainGaussID"].squeeze(0)
        center_3d = self.albedo_gaussians.get_xyz[
            mainIds[self.middle_coordinates[:, 1], self.middle_coordinates[:, 0]]
        ].mean(axis=0)

        tmp_up = camera.world_view_transform[:3, 1]
        rotations, translation = look_at_batch_torch(
            origins=global_samples,
            targets=center_3d.unsqueeze(0),
            ups=tmp_up.unsqueeze(0))

        return [
            MiniCam.from_params(
                R=R.detach().cpu().numpy(), T=t.detach().cpu().numpy(),
                FoVx=camera.FoVx, FoVy=camera.FoVy,
                width=camera.image_width, height=camera.image_height)
            for R, t in zip(rotations, translation)
        ] + [camera]

    def init_edition_gaussians(self, alpha_mask, frac_pixels: float = 0.4, use_shading_depth: bool = True):
        """Spawn new edit gaussians by unprojecting alpha-masked pixels at the rendered depth."""
        cam = self.sampled_views[-1]
        opengl_cam = FPSCamera(self.mode, res_x=cam.image_width, res_y=cam.image_height,
                               fov_y=np.rad2deg(cam.FoVy), z_near=cam.znear, z_far=cam.zfar)
        Rt = cam.world_view_transform.transpose(0, 1)
        opengl_cam.update_pose(torch.linalg.inv(Rt).detach().cpu().numpy())

        edit_pos_buffer = self.complex_edit.get_edit_pos_buffer_offline(opengl_cam)

        nb_samples = int(alpha_mask.sum().item() * frac_pixels)
        if nb_samples <= 0:
            print("Edit alpha mask is empty; nothing to spawn.")
            return
        avail_pixels = torch.nonzero(alpha_mask, as_tuple=False)
        selected_indices = torch.multinomial(
            torch.ones(avail_pixels.shape[0], device="cuda"), nb_samples, replacement=False)
        selected_pixels = avail_pixels[selected_indices]

        selection_mask = torch.zeros_like(alpha_mask, dtype=torch.bool)
        selection_mask[selected_pixels[:, 0], selected_pixels[:, 1]] = True

        pointmap = torch.from_numpy(edit_pos_buffer).to("cuda")[selection_mask]
        n_new = pointmap.shape[0]
        print(f"{n_new} new gaussians spawned for the edit.")

        self.edition_gaussians = GaussianModel(sh_degree=0)
        dist2 = torch.clamp_min(distCUDA2(pointmap), 1e-7)
        scales = torch.log(torch.sqrt(dist2))[..., None].repeat(1, 3)
        rots = torch.zeros((pointmap.shape[0], 4), device="cuda")
        rots[:, 0] = 1

        self.edition_gaussians.use_color_activation = False
        self.edition_gaussians.from_params({
            "_xyz": pointmap,
            "_features_dc": torch.ones((pointmap.shape[0], 1, 3), device="cuda") * 0.5,
            "_scaling": scales,
            "_rotation": rots,
            "_opacity": torch.ones((pointmap.shape[0], 1), device="cuda"),
        })

        with torch.no_grad():
            self.albedo_gaussians.add_pretrained_gaussians(self.edition_gaussians)

    def get_composited_edit_with_pretrained(self, viewpoints, use_shading_depth: bool = False):
        """For each viewpoint, composite the ComplexEdit RGBA over the pretrained albedo render."""
        gt_images, gt_invdepths, gt_masks = [], [], []
        main_view_alpha_edit = None

        for ind, cam in enumerate(viewpoints):
            render_pkg = render(cam, self.albedo_gaussians, self.pipe, self.background,
                                override_color=self.colors_precomp)
            if use_shading_depth and self._ready(self.shading_gaussians):
                invdepth = render(cam, self.shading_gaussians, self.pipe, self.background,
                                  override_color=self.colors_precomp,
                                  colors_activation=False)["depth"]
            else:
                invdepth = render_pkg["depth"]

            rendered_depth_np = render_pkg["depth"].detach().cpu().numpy().squeeze(0)
            rendered_albedo = render_pkg["render"]

            opengl_cam = FPSCamera(self.mode, res_x=cam.image_width, res_y=cam.image_height,
                                   fov_y=np.rad2deg(cam.FoVy), z_near=cam.znear, z_far=cam.zfar)
            Rt = cam.world_view_transform.transpose(0, 1)
            opengl_cam.update_pose(torch.linalg.inv(Rt).detach().cpu().numpy())

            edit_image = torch.from_numpy(
                self.complex_edit.get_render_offline(opengl_cam, np.ascontiguousarray(rendered_depth_np)))
            rgb_edit = edit_image[..., :3].to("cuda").permute(2, 0, 1)
            alpha_edit = edit_image[..., 3].to("cuda")

            gt_image = alpha_edit * rgb_edit + (1.0 - alpha_edit) * rendered_albedo.detach().clone()
            gt_masks.append(alpha_edit)

            if ind == len(viewpoints) - 1:
                main_view_alpha_edit = alpha_edit

            gt_invdepths.append(invdepth)
            gt_images.append(gt_image)

        return gt_images, gt_invdepths, gt_masks, main_view_alpha_edit

    def reoptimize_merge_multiview(self, gt_images, gt_invdepths, iterations: int,
                                   lambda_dssim: float = 0.2, depth_reg: bool = False,
                                   gt_masks=None):
        """Re-fit the albedo gaussians (with the freshly spawned edit gaussians) over multi-view GT."""
        try:
            self.albedo_gaussians.enable_params()
            self.albedo_gaussians.training_setup(training_args=self._opt())

            self.can_display_edit = True
            depth_l1_weight = get_expon_lr_func(1.0, 0.01, max_steps=iterations)
            loss = None

            for i in range(iterations):
                self.complex_edit.current_iteration = i + 1
                pov_id = np.random.randint(0, len(self.sampled_views))
                sampled_cam = self.sampled_views[pov_id]
                gt_image = gt_images[pov_id]
                gt_invdepth = gt_invdepths[pov_id]

                bg = torch.rand((3), device="cuda")
                render_pkg = render(sampled_cam, self.albedo_gaussians, self.pipe, bg,
                                    override_color=self.colors_precomp)
                rendered = render_pkg["render"]
                invDepth = render_pkg["depth"]
                radii = render_pkg["radii"]

                if gt_masks is not None:
                    m = gt_masks[pov_id]
                    invDepth = invDepth[m]
                    rendered = rendered[m]
                    gt_invdepth = gt_invdepth[m]
                    gt_image = gt_image[m]

                L1 = l1_loss(gt_image, rendered)
                loss = (1.0 - lambda_dssim) * L1 + lambda_dssim * (1.0 - fused_ssim(gt_images[pov_id].unsqueeze(0), rendered.unsqueeze(0))) \
                    + 0.1 * torch.mean(self.albedo_gaussians.get_scaling)

                if depth_reg:
                    Ll1depth_pure = torch.abs(invDepth - gt_invdepth).mean()
                    loss += (depth_l1_weight(i) * Ll1depth_pure).item()

                loss.backward()
                with torch.no_grad():
                    visible = radii > 0
                    self.albedo_gaussians.optimizer.step(visible, radii.shape[0])
                    self.albedo_gaussians.optimizer.zero_grad(set_to_none=True)

            print(f"Reoptimized {iterations} iterations, final loss: {loss.item() if loss is not None else 'n/a'}")
        finally:
            self.is_optimizing = False
            self.complex_edit.start_optimization = False

    def _opt(self):
        """Lazily build a default OptimizationParams the first time reoptimization runs."""
        if self.opt is None:
            tmp = ArgumentParser()
            self.opt = OptimizationParams(tmp).extract(tmp.parse_args([]))
        return self.opt

    def step(self):
        camera = self.camera
        world_to_view = torch.from_numpy(camera.to_camera).cuda().transpose(0, 1)
        full_proj_transform = torch.from_numpy(camera.full_projection).cuda().transpose(0, 1)
        mini = MiniCam(camera.res_x, camera.res_y, camera.fov_y, camera.fov_x,
                       camera.z_near, camera.z_far, world_to_view, full_proj_transform)

        self._ensure_ellipsoid_upload()

        render_time = 0.0

        if self.render_mode == 0:
            if self.editing_enabled and self.show_depth:
                net_image, render_time = self._render_with_timing(self._render_depth, mini)
            else:
                net_image, render_time = self._render_with_timing(self._render_splats, mini)
            self.point_view.step(net_image)
        elif self.render_mode == 1:
            self.ellipsoid_viewer.step(self.camera)
            render_time = glGetQueryObjectuiv(self.ellipsoid_viewer.query, GL_QUERY_RESULT) / 1e6
        elif self.editing_enabled:
            # MainDepth: not yet ported. Fall back to splats.
            if not getattr(self, "_warned_unimpl_render_mode", False):
                print(f"Render mode {self.render_modes[self.render_mode]!r} not implemented yet; falling back to Splats.")
                self._warned_unimpl_render_mode = True
            net_image, render_time = self._render_with_timing(self._render_splats, mini)
            self.point_view.step(net_image)

        self.monitor.step([render_time])

        if self.editing_enabled:
            self._editing_per_frame_update()

    def _render_with_timing(self, render_fn, mini):
        """Common splat/depth render scaffolding: cuda-event timing + exposure + permute."""
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        with torch.no_grad():
            with self.gaussian_lock:
                net_image = render_fn(mini)
            net_image = ((net_image * self.exposure) ** (1.0 / self.gamma)).permute(1, 2, 0)
        end.record()
        end.synchronize()
        return net_image, start.elapsed_time(end)

    def _render_depth(self, camera):
        """Rasterized depth from the gaussian set selected by `gaussian_color`.

        SHADING uses the shading gaussians; everything else uses the albedo
        gaussians. Toggle `display_normalized_depth` to switch between raw
        alpha-blended depth and depth normalized by the alpha mask.
        """
        if self.gaussian_color == "SHADING" and self._ready(self.shading_gaussians):
            primary = self.shading_gaussians
        else:
            primary = next(
                (g for g in (self.albedo_gaussians, self.shading_gaussians, self.residual_gaussians)
                 if self._ready(g)), None)
        if primary is None:
            return torch.zeros((3, camera.image_height, camera.image_width), device="cuda")
        render_pkg = render(camera, primary, self.pipe, self.background,
                            scaling_modifier=self.scaling_modifier,
                            override_color=self.colors_precomp,
                            opacity_modifier=self.opacity_modifier,
                            colors_activation=False)
        depth = render_pkg["depth"]
        if self.display_normalized_depth:
            depth = depth / render_pkg["alpha_mask"].clamp(min=1e-6)
        return depth.repeat(3, 1, 1)

    def _editing_per_frame_update(self):
        """Per-frame editing logic: reoptimization trigger.

        Selection-box input lives in show_gui() because it needs the Point View
        imgui window context. ComplexEdit.step + texture composition live in
        _render_splats so they share the splat render's depth + mainGaussID.
        """
        if (self.complex_edit.output_image is not None
                and not self.is_optimizing
                and self.complex_edit.start_optimization):
            self._launch_reoptimization()

    def _draw_gizmo(self):
        """Render an imguizmo manipulator over the Point View image to translate/rotate the edit plane."""
        if self._image_top_left is None:
            return

        gizmo.set_drawlist()
        gizmo.set_rect(self._image_top_left.x, self._image_top_left.y,
                       self.camera.res_x, self.camera.res_y)

        to_camera = self.camera.to_camera.copy()
        to_camera[1] *= -1  # OpenGL y-flip
        view_mat = Matrix16((to_camera.T).flatten().tolist())
        proj_mat = Matrix16((self.camera.projection.T).flatten().tolist())

        R = self.complex_edit.plane_rotation[:3, :3] @ self.complex_edit.user_rotation
        t = self.complex_edit.plane_translation + self.complex_edit.user_translation

        pose = Matrix16([
            R[0, 0], R[1, 0], R[2, 0], 0.0,
            R[0, 1], R[1, 1], R[2, 1], 0.0,
            R[0, 2], R[1, 2], R[2, 2], 0.0,
            t[0],    t[1],    t[2],    1.0,
        ])

        tool = self.tools[self.tool_id]
        gizmo_op = {
            "move": gizmo.OPERATION.translate,
            "scale": gizmo.OPERATION.scaleu,
            "rotate": gizmo.OPERATION.rotate,
        }[tool]
        gizmo.manipulate(view_mat, proj_mat, gizmo_op, gizmo.MODE.local, pose, None, None, None, None)

        M = np.array(pose.values).reshape(4, 4).T
        R_new = M[:3, :3]
        t_new = M[:3, 3]
        self.complex_edit.user_rotation = self.complex_edit.plane_rotation.T @ R_new
        self.complex_edit.user_translation = t_new - self.complex_edit.plane_translation

    def _handle_selection_box(self):
        """Track shift+left-mouse drag and on release feed the selected xyz to ComplexEdit.

        Must be called inside the Point View imgui window context — uses the same
        window-origin / padding / text-line offsets as the mouse_pos computation
        so the on-screen rectangle aligns exactly with the click position.
        """
        self.track_edit_zone = imgui.is_key_down(imgui.Key.left_shift)
        if self.track_edit_zone:
            imgui.set_mouse_cursor(imgui.MouseCursor_.hand)

        if imgui.is_mouse_down(imgui.MouseButton_.left) and self.track_edit_zone:
            if self.edit_start_pos is None:
                self.edit_start_pos = self.mouse_pos
            else:
                # convert image-space coords back to screen-space using the same
                # offsets that mouse_pos was derived from (we are in the Point
                # View window context here)
                origin = imgui.get_window_pos() + imgui.get_style().window_padding
                text_h = imgui.get_text_line_height_with_spacing()
                draw_list = imgui.get_window_draw_list()
                color = imgui.get_color_u32((1.0, 1.0, 0.0, 1.0))
                a = ImVec2(origin.x + self.edit_start_pos.x, origin.y + self.edit_start_pos.y + text_h)
                b = ImVec2(origin.x + self.mouse_pos.x, origin.y + self.mouse_pos.y + text_h)
                draw_list.add_rect(a, b, color, 0.0, 0, 2.0)
        elif self.edit_start_pos is not None:
            # release: compute the selection bbox and notify ComplexEdit
            self.edit_end_pos = self.mouse_pos
            np_array = np.stack(
                (np.array([[self.edit_start_pos.x, self.edit_start_pos.y]]),
                 np.array([[self.edit_end_pos.x, self.edit_end_pos.y]])),
                axis=1).squeeze(0)
            self.lower_coordinate = np_array.min(axis=0).astype(int)
            self.higher_coordinate = np_array.max(axis=0).astype(int)
            extents = self.higher_coordinate - self.lower_coordinate
            self.edit_start_pos = None
            self.edit_end_pos = None
            if extents[0] > 1 and extents[1] > 1 and self._ready(self.shading_gaussians):
                xyz_under = self.get_positions_under_selection()
                self.complex_edit._add_texture_callback(
                    xyz_under, width=int(extents[0]), height=int(extents[1]))

    def _launch_reoptimization(self):
        """Sample views, build composited GT, init edit gaussians, and run reoptimize_merge_multiview off-thread."""
        self.is_optimizing = True
        self.can_display_edit = False
        print("Starting multi-view re-optimization of the edit ...")

        # Build the MiniCam for the current viewpoint (used as the main view).
        cam = self.camera
        world_to_view = torch.from_numpy(cam.to_camera).cuda().transpose(0, 1)
        full_proj_transform = torch.from_numpy(cam.full_projection).cuda().transpose(0, 1)
        main_view = MiniCam(cam.res_x, cam.res_y, cam.fov_y, cam.fov_x,
                            cam.z_near, cam.z_far, world_to_view, full_proj_transform)

        if not self.sampled_views:
            self.sampled_views = self.sample_close_viewpoints(
                main_view, self.complex_edit.num_sampled_cameras, 0.7)

        gt_images, gt_invdepths, _gt_masks, main_view_alpha_edit = (
            self.get_composited_edit_with_pretrained(
                self.sampled_views, use_shading_depth=self.edit_use_shading_depth))

        self.prev_albedo_gaussians_params = self.albedo_gaussians.copy_params()
        self.init_edition_gaussians(main_view_alpha_edit,
                                    use_shading_depth=self.edit_use_shading_depth)

        threading.Thread(
            target=self.reoptimize_merge_multiview,
            args=(gt_images, gt_invdepths,
                  self.complex_edit.optimization_iterations,
                  0.2,
                  self.complex_edit.use_depth_reg),
            daemon=True,
        ).start()

    def _composite_complex_edit(self, camera, net_image):
        """Run a ComplexEdit step keyed off the shading set, then composite its textured plane onto net_image."""
        shading_pkg = render(camera, self.shading_gaussians, self.pipe, self.background,
                             scaling_modifier=self.scaling_modifier,
                             override_color=self.colors_precomp,
                             opacity_modifier=self.opacity_modifier,
                             colors_activation=False)
        mainIds = shading_pkg["mainGaussID"].squeeze(0)


        inv_depth = (1.0 / self.shading_gaussians.get_depthmap_from_ids(
            camera.full_proj_transform, mainIds))
        depth_np = inv_depth.cpu().numpy()

        self.complex_edit.step(camera=self.camera,
                               gaussians_renderer_depth=depth_np,
                               rel_mouse_pos=self.mouse_pos)

        out = self.complex_edit.output_image
        if out is None:
            return net_image

        rgb = torch.from_numpy(out.copy()[:, :, :3]).permute(2, 0, 1).to("cuda")
        alpha = torch.from_numpy(out.copy()[:, :, 3]).to("cuda")
        return rgb * alpha + (1.0 - alpha) * net_image

    def _render_splats(self, camera):
        """Shared splat-rendering dispatch keyed off the three Gaussian sets."""
        if self.gaussian_color == "ALBEDO" and self._ready(self.albedo_gaussians):
            albedo_pkg = render(camera, self.albedo_gaussians, self.pipe, self.background,
                                scaling_modifier=self.scaling_modifier,
                                override_color=self.colors_precomp,
                                opacity_modifier=self.opacity_modifier,
                                colors_activation=True)
            net_image = albedo_pkg["render"]
            if self._mouse_pos_in_bounds((net_image.shape[2], net_image.shape[1])):
                self.albedo_pixel_col = net_image[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]

            # In editing mode, drive ComplexEdit and composite its textured plane
            # over the albedo render. Mirrors local_viewer_decomp ALBEDO branch.
            if self.editing_enabled and self._ready(self.shading_gaussians):
                net_image = self._composite_complex_edit(camera, net_image)
            return net_image

        if self.gaussian_color == "SHADING" and self._ready(self.shading_gaussians):
            net_image = render(camera, self.shading_gaussians, self.pipe, self.background,
                               scaling_modifier=self.scaling_modifier, colors_activation=False)["render"]
            if self._mouse_pos_in_bounds((net_image.shape[2], net_image.shape[1])):
                self.shading_pixel_col = net_image[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
            return net_image

        if self.gaussian_color == "RESIDUAL" and self._ready(self.residual_gaussians):
            net_image = render(camera, self.residual_gaussians, self.pipe, self.background,
                               scaling_modifier=self.scaling_modifier, colors_activation=False,
                               enable_shs=self.use_shs)["render"]
            if self._mouse_pos_in_bounds((net_image.shape[2], net_image.shape[1])):
                self.glossy_pixel_col = net_image[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
            return net_image

        if (self.gaussian_color == "GLOSSY" and self._ready(self.residual_gaussians)
                and self._ready(self.albedo_gaussians) and self._ready(self.shading_gaussians)):
            albedo_render = render(camera, self.albedo_gaussians, self.pipe, self.background,
                                   scaling_modifier=self.scaling_modifier, colors_activation=True)["render"]
            shading_render = render(camera, self.shading_gaussians, self.pipe, self.background,
                                    scaling_modifier=self.scaling_modifier, colors_activation=False)["render"]
            glossy_render = render(camera, self.residual_gaussians, self.pipe, self.background,
                                   scaling_modifier=self.scaling_modifier, colors_activation=False,
                                   enable_shs=self.use_shs)["render"]
            if self._mouse_pos_in_bounds((albedo_render.shape[2], albedo_render.shape[1])):
                self.albedo_pixel_col = albedo_render[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
                self.shading_pixel_col = shading_render[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
                self.glossy_pixel_col = glossy_render[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
            return albedo_render * shading_render + glossy_render

        # DIFFUSE (default): albedo * shading if both exist; else whichever is present.
        if self._ready(self.albedo_gaussians) and self._ready(self.shading_gaussians):
            albedo_render = render(camera, self.albedo_gaussians, self.pipe, self.background,
                                   scaling_modifier=self.scaling_modifier, colors_activation=True)["render"]
            shading_render = render(camera, self.shading_gaussians, self.pipe, self.background,
                                    scaling_modifier=self.scaling_modifier, colors_activation=False)["render"]
            if self._mouse_pos_in_bounds((albedo_render.shape[2], albedo_render.shape[1])):
                self.albedo_pixel_col = albedo_render[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
                self.shading_pixel_col = shading_render[:, int(self.mouse_pos.y), int(self.mouse_pos.x)]
            return albedo_render * shading_render

        if self._ready(self.albedo_gaussians):
            return render(camera, self.albedo_gaussians, self.pipe, self.background,
                          scaling_modifier=self.scaling_modifier, colors_activation=True)["render"]

        if self._ready(self.shading_gaussians):
            return render(camera, self.shading_gaussians, self.pipe, self.background,
                          scaling_modifier=self.scaling_modifier, colors_activation=False)["render"]

        if self._ready(self.residual_gaussians):
            return render(camera, self.residual_gaussians, self.pipe, self.background,
                          scaling_modifier=self.scaling_modifier, colors_activation=False,
                          enable_shs=self.use_shs)["render"]

        return torch.zeros((3, camera.image_height, camera.image_width), device="cuda")

    def show_gui(self):
        if self.editing_enabled:
            gizmo.begin_frame()

        with imgui_ctx.begin("Point View Settings"):
            _, self.render_mode = imgui.list_box("Render Mode", self.render_mode, self.render_modes)

            bg_edited, self.background_color = imgui.color_picker3("Background Color", self.background_color)
            if bg_edited:
                self.background = torch.tensor(self.background_color, dtype=torch.float32, device="cuda")

            _, self.enable_pixel_inspector = imgui.checkbox("pixel inspector", self.enable_pixel_inspector)


            if self.albedo_gaussians is not None:
                imgui.text(f"{self.albedo_gaussians.get_xyz.shape[0]} albedo gaussians")
            if self.shading_gaussians is not None:
                imgui.text(f"{self.shading_gaussians.get_xyz.shape[0]} shading gaussians")
            if self.residual_gaussians is not None:
                imgui.text(f"{self.residual_gaussians.get_xyz.shape[0]} residual gaussians")

            imgui.separator_text("Render Settings")
            if imgui.radio_button("Albedo", self.gaussian_color == "ALBEDO"):
                self.gaussian_color = "ALBEDO"
            if imgui.radio_button("Shading", self.gaussian_color == "SHADING"):
                self.gaussian_color = "SHADING"
            if imgui.radio_button("Residual", self.gaussian_color == "RESIDUAL"):
                self.gaussian_color = "RESIDUAL"
            if imgui.radio_button("Diffuse", self.gaussian_color == "DIFFUSE"):
                self.gaussian_color = "DIFFUSE"
            if imgui.radio_button("Glossy", self.gaussian_color == "GLOSSY"):
                self.gaussian_color = "GLOSSY"
            _, self.exposure = imgui.drag_float("Exposure", self.exposure, v_min=-5, v_max=5, v_speed=0.01)
            _, self.gamma = imgui.drag_float("Gamma", self.gamma, v_min=0.01, v_max=5, v_speed=0.01)

            if self.render_mode == 0:
                _, self.scaling_modifier = imgui.drag_float("Scaling Factor", self.scaling_modifier,
                                                            v_min=0, v_max=10, v_speed=0.01)
                if self.residual_gaussians is not None:
                    imgui.separator_text(
                        f"Residual Gaussians Settings {self.residual_gaussians.active_sh_degree} degree shs for now.")
                    _, self.use_shs = imgui.checkbox("Enable SHs", self.use_shs)
            if self.render_mode == 1:
                _, self.ellipsoid_viewer.scaling_modifier = imgui.drag_float(
                    "Scaling Factor", self.ellipsoid_viewer.scaling_modifier,
                    v_min=0, v_max=10, v_speed=0.01)
                _, self.ellipsoid_viewer.render_floaters = imgui.checkbox(
                    "Render Floaters", self.ellipsoid_viewer.render_floaters)
                _, self.ellipsoid_viewer.limit = imgui.drag_float(
                    "Alpha Threshold", self.ellipsoid_viewer.limit,
                    v_min=0, v_max=1, v_speed=0.01)

            imgui.separator_text("Camera Settings")
            self.camera.show_gui()

            if self.editing_enabled:
                self._show_editing_settings_gui()

        with imgui_ctx.begin("Point View"):
            self._image_top_left = imgui.get_cursor_screen_pos()

            if self.render_mode == 0:
                self.point_view.show_gui()
            else:
                self.ellipsoid_viewer.show_gui()

            if self.enable_pixel_inspector:
                self.pixel_inspector.show_gui(self.point_view.texture)

            if imgui.is_item_hovered():
                self.camera.process_mouse_input()

            if imgui.is_item_focused() or imgui.is_item_hovered():
                self.camera.process_keyboard_input()

            self.mouse_pos = imgui.get_mouse_pos() - (imgui.get_window_pos() + imgui.get_style().window_padding)
            self.mouse_pos.y -= imgui.get_text_line_height_with_spacing()

            if self.editing_enabled:
                self._handle_selection_box()
                if self.display_gizmo:
                    self._draw_gizmo()

        with imgui_ctx.begin("Performance"):
            self.monitor.show_gui()

        if self.editing_enabled:
            self._show_editing_panels()

    def _show_editing_settings_gui(self):
        """Editing-only controls inside the Point View Settings window."""
        imgui.separator_text("Edit plane")
        _, self.tool_id = imgui.list_box("Edit tool", self.tool_id, self.tools)
        _, self.display_gizmo = imgui.checkbox("Display gizmo", self.display_gizmo)

        imgui.separator_text("Depth view")
        _, self.show_depth = imgui.checkbox("Show depth", self.show_depth)
        if self.show_depth:
            _, self.display_normalized_depth = imgui.checkbox(
                "Display normalized depth", self.display_normalized_depth)

        imgui.separator_text("Re-optimization")
        _, self.edit_use_shading_depth = imgui.checkbox("Use shading depth", self.edit_use_shading_depth)
        if self.is_optimizing:
            imgui.text(f"Optimizing: {self.complex_edit.current_iteration}/{self.complex_edit.optimization_iterations}")
        if self.prev_albedo_gaussians_params is not None and not self.is_optimizing:
            if imgui.button("Undo last edit"):
                with self.gaussian_lock:
                    self.albedo_gaussians.from_params(self.prev_albedo_gaussians_params)
                    self.prev_albedo_gaussians_params = None
                    self.edition_gaussians = None
                    self.sampled_views = []

        imgui.text("Hold SHIFT and drag to select an edit region.")

    def _show_editing_panels(self):
        """Editing-only top-level windows."""
        with imgui_ctx.begin("Complex Edit"):
            self.complex_edit.show_gui()

    def client_send(self):
        return None, {
            "scaling_modifier": self.scaling_modifier,
            "render_mode": self.render_mode,
        }

    def server_recv(self, _, text):
        self.scaling_modifier = text["scaling_modifier"]
        self.render_mode = text["render_mode"]


def _build_arg_parser():
    parser = ArgumentParser()
    subparsers = parser.add_subparsers(title="mode", dest="mode", required=True)

    local = subparsers.add_parser("local")
    local.add_argument("model_path")
    local.add_argument("iter", type=int, default=7000)


    client = subparsers.add_parser("client")
    client.add_argument("--ip", default="localhost")
    client.add_argument("--port", type=int, default=6009)

    server = subparsers.add_parser("server")
    server.add_argument("model_path")
    server.add_argument("iter", type=int, default=7000)
    server.add_argument("--ip", default="localhost")
    server.add_argument("--port", type=int, default=6009)

    return parser


if __name__ == "__main__":
    args = _build_arg_parser().parse_args()

    match args.mode:
        case "local":
            mode = ViewerMode.LOCAL
        case "client":
            mode = ViewerMode.CLIENT
        case "server":
            mode = ViewerMode.SERVER

    if mode is ViewerMode.CLIENT:
        viewer = GaussianViewer(mode)
    else:  # LOCAL or SERVER
        viewer = GaussianViewer.from_ply(args.model_path, args.iter, mode)

    viewer.run()