demo.py

from typing import Literal
import tyro
import os
import torch
import cv2
import imageio # To generate gifs
import pycolmap_scene_manager as pycolmap
from gsplat import rasterization
from ultralytics import YOLOWorld
from sam2.build_sam import build_sam2_video_predictor
import numpy as np

def torch_to_cv(tensor):
    img_cv = tensor.detach().cpu().numpy()[..., ::-1]
    img_cv = np.clip(img_cv*255, 0, 255).astype(np.uint8)
    return img_cv


def _detach_tensors_from_dict(d, inplace=True):
    if not inplace:
        d = d.copy()
    for key in d:
        if isinstance(d[key], torch.Tensor):
            d[key] = d[key].detach()
    return d


def load_checkpoint(checkpoint: str, data_dir: str, rasterizer: Literal["original", "gsplat"]="original", data_factor: int = 1):

    colmap_project = pycolmap.SceneManager(f"{data_dir}/sparse/0")
    colmap_project.load_cameras()
    colmap_project.load_images()
    colmap_project.load_points3D()
    model = torch.load(checkpoint) # Make sure it is generated by 3DGS original repo
    if rasterizer == "original":
        model_params, _ = model
        splats = {
            "active_sh_degree": model_params[0],
            "means": model_params[1],
            "features_dc": model_params[2],
            "features_rest": model_params[3],
            "scaling": model_params[4],
            "rotation": model_params[5],
            "opacity": model_params[6].squeeze(1),
        }
    elif rasterizer == "gsplat":
        print(model["splats"].keys())
        model_params = model["splats"]
        splats = {
            "active_sh_degree": 3,
            "means": model_params["means"],
            "features_dc": model_params["sh0"],
            "features_rest": model_params["shN"],
            "scaling": model_params["scales"],
            "rotation": model_params["quats"],
            "opacity": model_params["opacities"],
        }
    else:
        raise ValueError("Invalid rasterizer")

    _detach_tensors_from_dict(splats)

    # Assuming only one camera
    for camera in colmap_project.cameras.values():
        camera_matrix = torch.tensor(
            [
                [camera.fx, 0, camera.cx],
                [0, camera.fy, camera.cy],
                [0, 0, 1],
            ]
        )
        break

    camera_matrix[:2,:3] /= data_factor

    splats["camera_matrix"] = camera_matrix
    splats["colmap_project"] = colmap_project
    splats["colmap_dir"] = data_dir

    return splats

def get_viewmat_from_colmap_image(image):
    viewmat = torch.eye(4).float()#.to(device)
    viewmat[:3, :3] = torch.tensor(image.R()).float()#.to(device)
    viewmat[:3, 3] = torch.tensor(image.t).float()#.to(device)
    return viewmat

def create_checkerboard(width, height, size=64):
    checkerboard = np.zeros((height, width, 3), dtype=np.uint8)
    for y in range(0, height, size):
        for x in range(0, width, size):
            if (x // size + y // size) % 2 == 0:
                checkerboard[y:y + size, x:x + size] = 255
            else:
                checkerboard[y:y + size, x:x + size] = 128
    return checkerboard




def render_to_dir(output_dir: str, splats, feedback: bool = False):
    if feedback:
        cv2.destroyAllWindows()
        cv2.namedWindow("Initial Rendering", cv2.WINDOW_NORMAL)
    os.makedirs(output_dir, exist_ok=True)
    colmap_project = splats["colmap_project"]
    frame_idx = 0
    for image in sorted(colmap_project.images.values(), key=lambda x: x.name):
        image_name = image.name#.split(".")[0] + ".jpg"
        image_path = f"{output_dir}/{image_name}"
        means = splats["means"]
        colors_dc = splats["features_dc"]
        colors_rest = splats["features_rest"]
        colors = torch.cat([colors_dc, colors_rest], dim=1)
        opacities = torch.sigmoid(splats["opacity"])
        scales = torch.exp(splats["scaling"])
        quats = splats["rotation"]
        viewmat = get_viewmat_from_colmap_image(image)
        K = splats["camera_matrix"]
        output, _, info = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmats=viewmat[None],
            Ks=K[None],
            sh_degree=3,
            width=K[0, 2] * 2,
            height=K[1, 2] * 2,
        )
        output_cv = torch_to_cv(output[0])
        imageio.imsave(image_path, output_cv[:, :, ::-1])
        if feedback:
            cv2.imshow("Initial Rendering", output_cv)
            cv2.waitKey(1)
        frame_idx += 1

def prune_by_gradients(splats):
    colmap_project = splats["colmap_project"]
    frame_idx = 0
    means = splats["means"]
    colors_dc = splats["features_dc"]
    colors_rest = splats["features_rest"]
    colors = torch.cat([colors_dc, colors_rest], dim=1)
    opacities = torch.sigmoid(splats["opacity"])
    scales = torch.exp(splats["scaling"])
    quats = splats["rotation"]
    K = splats["camera_matrix"]
    colors.requires_grad = True
    gaussian_grads = torch.zeros(colors.shape[0], device=colors.device)
    for image in sorted(colmap_project.images.values(), key=lambda x: x.name):
        viewmat = get_viewmat_from_colmap_image(image)
        output, _, _ = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors[:,0,:],
            viewmats=viewmat[None],
            Ks=K[None],
            # sh_degree=3,
            width=K[0, 2] * 2,
            height=K[1, 2] * 2,
        )
        frame_idx += 1
        # pseudo_loss = ((output.detach() + 1 - output)**2).mean()
        pseudo_loss = output.mean()
        pseudo_loss.backward()
        # print(colors.grad.shape)
        gaussian_grads += (colors.grad[:,0]).norm(dim=[1])
        colors.grad.zero_()

    mask = gaussian_grads > 0
    print("Total splats", len(gaussian_grads))
    print("Pruned", (~mask).sum(), "splats")
    print("Remaining", mask.sum(), "splats")
    splats = splats.copy()
    splats["means"] = splats["means"][mask]
    splats["features_dc"] = splats["features_dc"][mask]
    splats["features_rest"] = splats["features_rest"][mask]
    splats["scaling"] = splats["scaling"][mask]
    splats["rotation"] = splats["rotation"][mask]
    splats["opacity"] = splats["opacity"][mask]
    return splats, mask

def test_proper_pruning(splats, splats_after_pruning):
    colmap_project = splats["colmap_project"]
    frame_idx = 0
    means = splats["means"]
    colors_dc = splats["features_dc"]
    colors_rest = splats["features_rest"]
    colors = torch.cat([colors_dc, colors_rest], dim=1)
    opacities = torch.sigmoid(splats["opacity"])
    scales = torch.exp(splats["scaling"])
    quats = splats["rotation"]

    means_pruned = splats_after_pruning["means"]
    colors_dc_pruned = splats_after_pruning["features_dc"]
    colors_rest_pruned = splats_after_pruning["features_rest"]
    colors_pruned = torch.cat([colors_dc_pruned, colors_rest_pruned], dim=1)
    opacities_pruned = torch.sigmoid(splats_after_pruning["opacity"])
    scales_pruned = torch.exp(splats_after_pruning["scaling"])
    quats_pruned = splats_after_pruning["rotation"]



    K = splats["camera_matrix"]
    total_error = 0
    max_pixel_error = 0
    for image in sorted(colmap_project.images.values(), key=lambda x: x.name):
        viewmat = get_viewmat_from_colmap_image(image)
        output, _, _ = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmats=viewmat[None],
            Ks=K[None],
            sh_degree=3,
            width=K[0, 2] * 2,
            height=K[1, 2] * 2,
        )

        output_pruned, _, _ = rasterization(
            means_pruned,
            quats_pruned,
            scales_pruned,
            opacities_pruned,
            colors_pruned,
            viewmats=viewmat[None],
            Ks=K[None],
            sh_degree=3,
            width=K[0, 2] * 2,
            height=K[1, 2] * 2,
        )

        total_error += torch.abs((output - output_pruned)).sum()
        max_pixel_error = max(max_pixel_error, torch.abs((output - output_pruned)).max())

    percentage_pruned = (len(splats["means"]) - len(splats_after_pruning["means"])) / len(splats["means"]) * 100

    assert max_pixel_error < 1 / (255*2), "Max pixel error should be less than 1/(255*2), safety margin"
    print("Report {}% pruned, max pixel error = {}, total pixel error = {}".format(percentage_pruned, max_pixel_error, total_error))


def get_mask3d(splats, prompt, data_dir: str, results_dir: str, show_visual_feedback: bool = False, mask_interval: int = 1, voting_method: Literal["gradient", "binary", "projection"] = "gradient", mask_dir=None):
    checkpoint = "./checkpoints/sam2_hiera_large.pt"
    if not os.path.exists(checkpoint):
        raise RuntimeError("Please download the checkpoint sam2_hiera_large.pt to checkpoints folder")
    
    if show_visual_feedback:
        cv2.destroyAllWindows()
        cv2.namedWindow("2D Mask", cv2.WINDOW_NORMAL)
    model_cfg = "sam2_hiera_l.yaml"
    mask_predictor = build_sam2_video_predictor(model_cfg, checkpoint)
    yolo_world = YOLOWorld("yolov8s-worldv2.pt")
    yolo_world.set_classes([prompt])
    colmap_project = splats["colmap_project"]
    first_image_name = sorted(colmap_project.images.values(), key=lambda x: x.name)[0].name
    first_image_path = f"{results_dir}/images/{first_image_name}"
    frame_idx = 0
    with torch.autocast("cuda", dtype=torch.bfloat16):
        state = mask_predictor.init_state(f"{results_dir}/images/")

        result = yolo_world(first_image_path)[0]

        box = result.boxes[0].xyxy[0].tolist()


        # add new prompts and instantly get the output on the same frame
        _, object_ids, masks = mask_predictor.add_new_points_or_box(
            state, box=box, frame_idx=0, obj_id=0
        )

        means = splats["means"]
        colors_dc = splats["features_dc"]
        colors_rest = splats["features_rest"]

        colors = colors_dc[:,0,:] * 0 # Just to show that gradient (opacity * transmittance) is independent of color. Any value will work.


        opacities = torch.sigmoid(splats["opacity"])
        scales = torch.exp(splats["scaling"])
        quats = splats["rotation"]
        K = splats["camera_matrix"]
        colors.requires_grad = True

        gaussian_grads = torch.zeros(colors.shape[0], device=colors.device)
        mask_dir = f"{results_dir}/masks_with_images"
        mask_bin_dir = f"{results_dir}/masks_bin"
        os.makedirs(mask_dir, exist_ok=True)
        os.makedirs(mask_bin_dir, exist_ok=True)
        
        # propagate the prompts to get masklets throughout the video
        frame_idx = 0
        for image, (frame_idx, object_ids, masks) in zip(
            sorted(colmap_project.images.values(), key=lambda x: x.name),
            mask_predictor.propagate_in_video(state),
        ):

            image_name = image.name#.split(".")[0] + ".jpg"
            # image_name = f"frame_"
            image_path = f"{results_dir}/images/{image_name}"
            mask_path = f"{mask_dir}/{image_name}"
            mask_bin_path = f"{mask_bin_dir}/{image.name}"

            frame = cv2.imread(image_path)
            mask = masks[0, 0].cpu().numpy() >= 0
            mask = mask.astype(float)
            mask = cv2.blur(mask, (7, 7))
            mask = mask > 0
            mask = mask.astype(bool)

            mask_red = np.zeros_like(frame)

            mask_red[:, :, -1][mask] = 255
            mask_bin = mask.astype(np.uint8) * 255
            cv2.imwrite(mask_bin_path, mask_bin)
            output = cv2.addWeighted(frame, 1, mask_red, 0.5, 0)
            cv2.imwrite(mask_path, output)


            frame_idx += 1
            if (frame_idx % mask_interval != 1) and (mask_interval != 1):
                continue
            if show_visual_feedback:
                cv2.imshow("2D Mask", output)
                cv2.waitKey(1)
            
            viewmat = get_viewmat_from_colmap_image(image)

            
            width = frame.shape[1]
            height = frame.shape[0]

            output_for_grad, _, meta = rasterization(
                means,
                quats,
                scales,
                opacities,
                colors,
                viewmat[None],
                K[None],
                width=width,
                height=height,
                # sh_degree=3,
            )

            target = output_for_grad[0] * torch.from_numpy(mask)[..., None].cuda().float()
            loss = 1 * target.mean()

            loss.backward(retain_graph=True)
            if voting_method == "gradient":
                mins = torch.min(colors.grad, dim=-1).values
                maxes = torch.max(colors.grad, dim=-1).values
                assert torch.allclose(mins , maxes), "Something is wrong with gradient calculation"
                gaussian_grads += (colors.grad).norm(dim=[1])
            elif voting_method == "binary":
                gaussian_grads += 1 * (colors.grad.norm(dim=[1]) > 0)
            elif voting_method == "projection":

                means2d = np.round(meta["means2d"].detach().cpu().numpy()).astype(int)
                means2d_mask = (means2d[:, 0] >= 0) & (means2d[:, 0] < width) & (means2d[:, 1] >= 0) & (means2d[:, 1] < height)
                means2d = means2d[means2d_mask]
                gaussian_ids = meta["gaussian_ids"].detach().cpu().numpy()
                gaussian_ids = gaussian_ids[means2d_mask]

                means2d_mask = mask[means2d[:, 1], means2d[:, 0]] # Check if the splat is in the mask
                gaussian_grads[torch.from_numpy(gaussian_ids[~means2d_mask]).long()] -= 1
                gaussian_grads[torch.from_numpy(gaussian_ids[means2d_mask]).long()] += 1
            else:
                raise ValueError("Invalid voting method")

            colors.grad.zero_()


            mask_inverted = ~mask
            target = output_for_grad[0] * torch.from_numpy(mask_inverted).cuda()[
                ..., None
            ]
            loss = 1 * target.mean()
            loss.backward(retain_graph=False)

            if voting_method == "gradient":
                gaussian_grads -= (colors.grad).norm(dim=[1])
            elif voting_method == "binary":
                gaussian_grads -= 1 * ((colors.grad).norm(dim=[1]) > 0)
            elif voting_method == "projection":
                pass
            else:
                raise ValueError("Invalid voting method")
            colors.grad.zero_()

        mask_3d = gaussian_grads > 0
        mask_3d_inverted = gaussian_grads < 0 # We don't need Gaussians without any influence ie gaussian_grads == 0
        return mask_3d, mask_3d_inverted

def apply_mask3d(splats, mask3d,mask3d_inverted, results_dir: str):
    if mask3d_inverted == None:
        mask3d_inverted = ~mask3d
    extracted = splats.copy()
    deleted = splats.copy()
    masked = splats.copy()
    extracted["means"] = extracted["means"][mask3d]
    extracted["features_dc"] = extracted["features_dc"][mask3d]
    extracted["features_rest"] = extracted["features_rest"][mask3d]
    extracted["scaling"] = extracted["scaling"][mask3d]
    extracted["rotation"] = extracted["rotation"][mask3d]
    extracted["opacity"] = extracted["opacity"][mask3d]

    deleted["means"] = deleted["means"][mask3d_inverted]
    deleted["features_dc"] = deleted["features_dc"][mask3d_inverted]
    deleted["features_rest"] = deleted["features_rest"][mask3d_inverted]
    deleted["scaling"] = deleted["scaling"][mask3d_inverted]
    deleted["rotation"] = deleted["rotation"][mask3d_inverted]
    deleted["opacity"] = deleted["opacity"][mask3d_inverted]

    masked["features_dc"][mask3d] = 1#(1 - 0.5) / 0.2820947917738781
    masked["features_dc"][~mask3d] = 0#(0 - 0.5) / 0.2820947917738781
    masked["features_rest"][~mask3d] = 0

    return extracted, deleted, masked

    

def render_to_gif(output_path: str, splats, feedback: bool = False, use_checkerboard_background: bool = False, no_sh: bool=False):
    if feedback:
        cv2.destroyAllWindows()
        cv2.namedWindow("Rendering", cv2.WINDOW_NORMAL)
    frames = []
    means = splats["means"]
    colors_dc = splats["features_dc"]
    colors_rest = splats["features_rest"]
    colors = torch.cat([colors_dc, colors_rest], dim=1)
    if no_sh == True:
        colors = colors_dc[:,0,:]
    opacities = torch.sigmoid(splats["opacity"])
    scales = torch.exp(splats["scaling"])
    quats = splats["rotation"]
    K = splats["camera_matrix"]
    aux_dir = output_path + ".images"
    os.makedirs(aux_dir, exist_ok=True)
    for image in sorted(splats["colmap_project"].images.values(), key=lambda x: x.name):
        viewmat = get_viewmat_from_colmap_image(image)
        output, alphas, meta = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmat[None],
            K[None],
            width=K[0, 2]*2,
            height=K[1, 2]*2,
            sh_degree=3 if not no_sh else None,
        )
        frame = np.clip(output[0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)
        if use_checkerboard_background:
            checkerboard = create_checkerboard(frame.shape[1], frame.shape[0])
            alphas = alphas[0].detach().cpu().numpy()
            frame = frame * alphas + checkerboard * (1 - alphas)
            frame = np.clip(frame, 0, 255).astype(np.uint8)
        frames.append(frame)
        if feedback:
            cv2.imshow("Rendering", frame[...,::-1])
            cv2.imwrite(f"{aux_dir}/{image.name}", frame[...,::-1])
            cv2.waitKey(1)
    imageio.mimsave(output_path, frames, fps=10)
    if feedback:
        cv2.destroyAllWindows()

def render_mask_pred(output_dir: str, splats, feedback: bool = False):
    if feedback:
        cv2.destroyAllWindows()
        cv2.namedWindow("Rendering", cv2.WINDOW_NORMAL)
    frames = []
    means = splats["means"]
    colors_dc = splats["features_dc"]
    colors_rest = splats["features_rest"]

    colors = colors_dc[:,0,:]
    opacities = torch.sigmoid(splats["opacity"])
    scales = torch.exp(splats["scaling"])
    quats = splats["rotation"]
    K = splats["camera_matrix"]
    aux_dir = output_dir
    os.makedirs(aux_dir, exist_ok=True)
    for image in sorted(splats["colmap_project"].images.values(), key=lambda x: x.name):
        viewmat = get_viewmat_from_colmap_image(image)
        output, alphas, meta = rasterization(
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmat[None],
            K[None],
            width=K[0, 2]*2,
            height=K[1, 2]*2,
            sh_degree=None,
        )
        frame = np.clip(output[0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)
        frame = frame > 128
        frame = frame.astype(np.uint8) * 255
            # cv2.imshow("Checkerboard", checkerboard)
        # frames.append(frame)
        if feedback:
            cv2.imshow("Rendering", frame[...,::-1])
            cv2.imwrite(f"{aux_dir}/{image.name}", frame)
            cv2.waitKey(1)
    # imageio.mimsave(output_path, frames, fps=10)
    if feedback:
        cv2.destroyAllWindows()

def export_mask(mask3d, prune_mask, results_dir: str):
    if prune_mask is not None:
        mask3d_export = torch.zeros_like(prune_mask).bool()
        mask3d_export[prune_mask] = mask3d
        torch.save(mask3d_export, f"{results_dir}/mask3d.pth")
    else:
        torch.save(mask3d, f"{results_dir}/mask3d.pth")

def main(
        data_dir: str = "./data/chair", # colmap path
        checkpoint: str = "./data/chair/checkpoint.pth", # checkpoint path, can generate from original 3DGS repo
        prompt: str = "chair", # prompt
        results_dir: str = "./results/chair", # output path
        show_visual_feedback: bool = True, # Will show opencv window,
        rasterizer: Literal["original", "gsplat"] = "original", # Original or GSplat for checkpoints
        data_factor: int = 1,
        mask_interval: int = 1,
        voting_method: Literal["gradient", "binary", "projection"] = "gradient",

):
    """
    Demo program.


    Args:
        data_dir: Path to the colmap project directory
        checkpoint: Path to the checkpoint file
        prompt: Prompt for the mask prediction
        results_dir: Path to the output directory
        show_visual_feedback: Show visual feedback
        rasterizer: Rasterizer used to create the checkpoint
        data_factor: Factor to scale down the resolution of the images
        mask_interval: The interval between images for taking masked gradients
        voting_method: Voting method to generate 3D mask
    """
    
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this demo")
    
    torch.set_default_device('cuda')

    os.makedirs(results_dir, exist_ok=True)
    splats = load_checkpoint(checkpoint, data_dir, rasterizer=rasterizer, data_factor=data_factor)
    splats_optimized, prune_mask = prune_by_gradients(splats)
    test_proper_pruning(splats, splats_optimized)

    del splats
    splats = splats_optimized

    render_to_dir(f"{results_dir}/images", splats, show_visual_feedback)
    mask3d, mask3d_inverted = get_mask3d(splats, prompt, data_dir, results_dir, show_visual_feedback, mask_interval=mask_interval, voting_method=voting_method)
    
    export_mask(mask3d, prune_mask, results_dir)
    
    extracted, deleted, masked = apply_mask3d(splats, mask3d, mask3d_inverted, results_dir)
    # render_mask_pred(f"{results_dir}/mask_bin_pred", masked, show_visual_feedback)
    render_to_gif(f"{results_dir}/extracted.gif", extracted, show_visual_feedback, use_checkerboard_background=True)
    render_to_gif(f"{results_dir}/deleted.gif", deleted, show_visual_feedback)
    render_to_gif(f"{results_dir}/masked.gif", masked, show_visual_feedback, no_sh=True)


if __name__ == "__main__":
    tyro.cli(main)