utils/video_utils.py

import logging
import os
from typing import Callable, Dict, List, Optional
from plyfile import PlyData, PlyElement

import imageio
import numpy as np
import torch
from utils.loss_utils import l1_loss
from skimage.metrics import structural_similarity as ssim
from lpipsPyTorch import lpips
from torch import Tensor
from tqdm import tqdm, trange
from gaussian_renderer import render

from utils.image_utils import psnr
# from gs_renderer import Renderer,GaussianModel,TrainCam
# from trainer import Trainer
# from utils.misc import get_robust_pca
from utils.visualization_tools import (
    resize_five_views,
    scene_flow_to_rgb,
    to8b,
    visualize_depth,
)

depth_visualizer = lambda frame, opacity: visualize_depth(
    frame,
    opacity,
    lo=4.0,
    hi=120,
    depth_curve_fn=lambda x: -np.log(x + 1e-6),
)
flow_visualizer = (
    lambda frame: scene_flow_to_rgb(
        frame,
        background="bright",
        flow_max_radius=1.0,
    )
    .cpu()
    .numpy()
)
get_numpy: Callable[[Tensor], np.ndarray] = lambda x: x.squeeze().cpu().numpy()
non_zero_mean: Callable[[Tensor], float] = (
    lambda x: sum(x) / len(x) if len(x) > 0 else -1
)

def get_robust_pca(features: torch.Tensor, m: float = 2, remove_first_component=False):
    # features: (N, C)
    # m: a hyperparam controlling how many std dev outside for outliers
    assert len(features.shape) == 2, "features should be (N, C)"
    reduction_mat = torch.pca_lowrank(features, q=3, niter=20)[2]
    colors = features @ reduction_mat
    if remove_first_component:
        colors_min = colors.min(dim=0).values
        colors_max = colors.max(dim=0).values
        tmp_colors = (colors - colors_min) / (colors_max - colors_min)
        fg_mask = tmp_colors[..., 0] < 0.2
        reduction_mat = torch.pca_lowrank(features[fg_mask], q=3, niter=20)[2]
        colors = features @ reduction_mat
    else:
        fg_mask = torch.ones_like(colors[:, 0]).bool()
    d = torch.abs(colors[fg_mask] - torch.median(colors[fg_mask], dim=0).values)
    mdev = torch.median(d, dim=0).values
    s = d / mdev
    rins = colors[fg_mask][s[:, 0] < m, 0]
    gins = colors[fg_mask][s[:, 1] < m, 1]
    bins = colors[fg_mask][s[:, 2] < m, 2]

    rgb_min = torch.tensor([rins.min(), gins.min(), bins.min()])
    rgb_max = torch.tensor([rins.max(), gins.max(), bins.max()])
    return reduction_mat, rgb_min.to(reduction_mat), rgb_max.to(reduction_mat)

def render_pixels(
    viewpoint_stack,
    gaussians,
    bg,
    pipe,
    compute_metrics: bool = True,
    return_decomposition: bool = True,
    debug:bool = False
):
    """
    Render pixel-related outputs from a model.

    Args:
        ....skip obvious args
        compute_metrics (bool, optional): Whether to compute metrics. Defaults to False.
        vis_indices (Optional[List[int]], optional): Indices to visualize. Defaults to None.
        return_decomposition (bool, optional): Whether to visualize the static-dynamic decomposition. Defaults to True.
    """
    # set up render function
    render_results = render_func(
        viewpoint_stack,
        gaussians,
        pipe,
        bg,
        compute_metrics=compute_metrics,
        return_decomposition=return_decomposition,
        debug = debug
    )
    if compute_metrics:
        num_samples = len(viewpoint_stack)
        print(f"Eval over {num_samples} images:")
        print(f"\tPSNR: {render_results['psnr']:.4f}")
        print(f"\tSSIM: {render_results['ssim']:.4f}")
        print(f"\tLPIPS: {render_results['lpips']:.4f}")
        # print(f"\tFeature PSNR: {render_results['feat_psnr']:.4f}")
        print(f"\tMasked PSNR: {render_results['masked_psnr']:.4f}")
        print(f"\tMasked SSIM: {render_results['masked_ssim']:.4f}")
        # print(f"\tMasked Feature PSNR: {render_results['masked_feat_psnr']:.4f}")

    return render_results


def render_func(
    viewpoint_stack,
    gaussians,
    pipe,
    bg,
    compute_metrics: bool = False,
    return_decomposition:bool = False,
    num_cams: int = 3,
    debug: bool = False,
    save_seperate_pcd = False
):
    """
    Renders a dataset utilizing a specified render function.
    For efficiency and space-saving reasons, this function doesn't store the original features; instead, it keeps
    the colors reduced via PCA.
    TODO: clean up this function

    Parameters:
        dataset: Dataset to render.
        render_func: Callable function used for rendering the dataset.
        compute_metrics: Optional; if True, the function will compute and return metrics. Default is False.
    """
    # rgbs
    rgbs, gt_rgbs = [], []
    static_rgbs, dynamic_rgbs = [], []
    shadow_reduced_static_rgbs, shadow_only_static_rgbs = [], []

    # depths
    depths, median_depths = [], []
    static_depths, static_opacities = [], []
    dynamic_depths, dynamic_opacities = [], []

    # sky
    opacities, sky_masks = [], []

    # features
    pred_dinos, gt_dinos = [], []
    pred_dinos_pe_free, pred_dino_pe = [], []
    static_dinos, dynamic_dinos = [], []  # should we also render this?

    # cross-rendering results
    dynamic_dino_on_static_rgbs, dynamic_rgb_on_static_dinos = [], []

    # flows
    forward_flows, backward_flows = [], []
    dx_list = []

    if compute_metrics:
        psnrs, ssim_scores, feat_psnrs = [], [], []
        masked_psnrs, masked_ssims = [], []
        masked_feat_psnrs = [],
        lpipss = []

    with torch.no_grad():
        for i in tqdm(range(len(viewpoint_stack)), desc=f"rendering full data", dynamic_ncols=True):
            viewpoint_cam = viewpoint_stack[i]

            render_pkg = render(viewpoint_cam, gaussians, pipe, bg,return_decomposition = return_decomposition,return_dx=True)
            image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]
            
            # ------------- rgb ------------- #
            rgb = image
            gt_rgb = viewpoint_cam.original_image.cuda()

            rgbs.append(get_numpy(rgb.permute(1, 2, 0)))
            gt_rgbs.append(get_numpy(gt_rgb.permute(1, 2, 0)))
            
            if "render_s" in render_pkg:
                static_rgbs.append(get_numpy(render_pkg["render_s"].permute(1, 2, 0)))
                visibility_filter_s = render_pkg['visibility_filter_s']
            if "render_d" in render_pkg:
                # green screen blending for better visualization
                # green_background = torch.tensor([0.0, 177, 64]) / 255.0
                # green_background = green_background.to(render_pkg["render_d"].device)
                dy_rgb = render_pkg["render_d"].permute(1, 2, 0)
                # dy_rgb = dy_rgb * 0.8 + green_background * 0.2
                dynamic_rgbs.append(get_numpy(dy_rgb))
                visibility_filter_d = render_pkg['visibility_filter_d']
            
            # ------------- depth ------------- #
            depth = render_pkg["depth"]
            depth_np = depth.permute(1, 2, 0).cpu().numpy()
            depth_np /= depth_np.max()
            # depth_np = np.repeat(depth_np, 3, axis=2)

            depths.append(depth_np)

            # ------------- flow ------------- #
            if "dx" in render_pkg and render_pkg['dx'] is not None:
                dx = render_pkg['dx']
                dx = torch.tensor(dx)
                dx_max = torch.max(dx)
                dx_min = torch.min(dx)
                dx_list.append(dx)     
            if compute_metrics:
                psnrs.append(psnr(rgb, gt_rgb).mean().double().item())
                # ssim_scores.append(ssim(rgb, gt_rgb).mean().item())
                ssim_scores.append(
                    ssim(
                        get_numpy(rgb),
                        get_numpy(gt_rgb),
                        data_range=1.0,
                        channel_axis=0,
                    )
                )
                lpipss.append(torch.tensor(lpips(rgb, gt_rgb,net_type='alex')).mean().item())
                
                dynamic_mask = get_numpy(viewpoint_cam.dynamic_mask).astype(bool)
                if dynamic_mask.sum() > 0:
                    rgb_d = rgb.permute(1, 2, 0)[dynamic_mask]
                    rgb_d = rgb_d.permute(1, 0)
                    gt_rgb_d = gt_rgb.permute(1, 2, 0)[dynamic_mask]
                    gt_rgb_d = gt_rgb_d.permute(1, 0)

                    masked_psnrs.append(
                    psnr(rgb_d, gt_rgb_d).mean().double().item()
                    )
                    masked_ssims.append(
                        ssim(
                            get_numpy(rgb.permute(1, 2, 0)),
                            get_numpy(gt_rgb.permute(1, 2, 0)),
                            data_range=1.0,
                            channel_axis=-1,
                            full=True,
                        )[1][dynamic_mask].mean()
                    )

        if save_seperate_pcd and len(dx_list)>1:    
            # 首先根据visibility_filter 选出所有的可见范围内的点
            # 然后得到dynamic 和 static 的mask，把点保存
                
            dynamic_pcd_path = os.path.join('test','dynamic.ply')
            static_pcd_path = os.path.join('test','static.ply')

            gaussians.save_ply_split(dynamic_pcd_path, static_pcd_path, dx_list, visibility_filter)

        if len(dx_list)>1:
            # deformation flow -> forward & backward flow
            bf_color_first = []
            ff_color_last = []
            for t in range(len(dx_list)): # 防止越界
                if t < len(dx_list)-num_cams:
                    # forward_flow_t 归一化一下
                    forward_flow_t = dx_list[t + num_cams] - dx_list[t]
                    ff_color = flow_visualizer(forward_flow_t)
                    ff_color = torch.from_numpy(ff_color).to("cuda") 
                    if debug:
                        ff_color = (ff_color - torch.min(ff_color)) / (torch.max(ff_color) - torch.min(ff_color) + 1e-6)  # 归一化，避免除零错误

                    if t == len(dx_list)-num_cams-1 or t == len(dx_list)-num_cams-2 or t == len(dx_list)-num_cams-3: 
                        ff_color_last.append(ff_color)              
                    render_pkg2 = render(viewpoint_stack[t], gaussians, pipe, bg, override_color=ff_color)
                    ff_map = render_pkg2['render'].permute(1, 2, 0).cpu().numpy()

                    # print(ff_map.max())
                    # print(ff_map.min())
                    forward_flows.append(ff_map)
                
                # 同时处理 backward flow，除第一个时刻外
                if t > num_cams-1:
                    backward_flow_t = dx_list[t] - dx_list[t - num_cams]
                    bf_color = flow_visualizer(backward_flow_t)
                    bf_color = torch.from_numpy(bf_color).to("cuda") 
                    if debug:
                        bf_color = (bf_color - torch.min(bf_color)) / (torch.max(bf_color) - torch.min(bf_color) + 1e-6)  # 归一化，避免除零错误
                    if t == num_cams or t == num_cams+1 or t == num_cams+2: 
                        bf_color_first.append(bf_color)                 
                    # viewpoint_cam 要变化
                    render_pkg2 = render(viewpoint_stack[t], gaussians, pipe, bg, override_color=bf_color)
                    bf_map = render_pkg2['render'].permute(1, 2, 0).cpu().numpy()

                    backward_flows.append(bf_map)

            for i, bf_color in enumerate(bf_color_first):
                render_pkg3 = render(viewpoint_stack[i], gaussians, pipe, bg, override_color=bf_color)            
                bf_map_first = render_pkg3['render'].permute(1, 2, 0).cpu().numpy()       
                # 对于 backward flow 的第一个时刻，复制第一个计算的 forward flow
                backward_flows.insert(i, bf_map_first)

            for i, ff_color in enumerate(ff_color_last):
                # 对于 forward flow 的最后一个时刻，复制最后一个计算的 backward flow
                render_pkg4 = render(viewpoint_stack[len(viewpoint_stack)-num_cams+i], gaussians, pipe, bg, override_color=ff_color)            
                ff_map_last = render_pkg4['render'].permute(1, 2, 0).cpu().numpy()       
                forward_flows.append(ff_map_last)           

    # messy aggregation...
    results_dict = {}
    results_dict["psnr"] = non_zero_mean(psnrs) if compute_metrics else -1
    results_dict["ssim"] = non_zero_mean(ssim_scores) if compute_metrics else -1
    results_dict["lpips"] = non_zero_mean(lpipss) if compute_metrics else -1
    results_dict["masked_psnr"] = non_zero_mean(masked_psnrs) if compute_metrics else -1
    results_dict["masked_ssim"] = non_zero_mean(masked_ssims) if compute_metrics else -1

    results_dict["rgbs"] = rgbs
    results_dict["depths"] = depths
    results_dict["opacities"] = opacities

    if len(gt_rgbs) > 0:
        results_dict["gt_rgbs"] = gt_rgbs
    if len(static_rgbs)>0:
        results_dict["static_rgbs"] = static_rgbs
    if len(dynamic_rgbs)>0:
        results_dict["dynamic_rgbs"] = dynamic_rgbs
    if len(sky_masks) > 0:
        results_dict["gt_sky_masks"] = sky_masks
    if len(pred_dinos) > 0:
        results_dict["dino_feats"] = pred_dinos
    if len(gt_dinos) > 0:
        results_dict["gt_dino_feats"] = gt_dinos
    if len(pred_dinos_pe_free) > 0:
        results_dict["dino_feats_pe_free"] = pred_dinos_pe_free
    if len(pred_dino_pe) > 0:
        results_dict["dino_pe"] = pred_dino_pe
    if len(static_dinos) > 0:
        results_dict["static_dino_feats"] = static_dinos
    if len(dynamic_dinos) > 0:
        results_dict["dynamic_dino_feats"] = dynamic_dinos
    if len(dynamic_dino_on_static_rgbs) > 0:
        results_dict["dynamic_dino_on_static_rgbs"] = dynamic_dino_on_static_rgbs
    if len(dynamic_rgb_on_static_dinos) > 0:
        results_dict["dynamic_rgb_on_static_dinos"] = dynamic_rgb_on_static_dinos
    if len(shadow_reduced_static_rgbs) > 0:
        results_dict["shadow_reduced_static_rgbs"] = shadow_reduced_static_rgbs
    if len(shadow_only_static_rgbs) > 0:
        results_dict["shadow_only_static_rgbs"] = shadow_only_static_rgbs
    if len(forward_flows) > 0:
        results_dict["forward_flows"] = forward_flows
    if len(backward_flows) > 0:
        results_dict["backward_flows"] = backward_flows
    if len(median_depths) > 0:
        results_dict["median_depths"] = median_depths
    if len(dx_list) > 0:
        results_dict['dx_list'] = dx_list
    return results_dict


def save_videos(
    render_results: Dict[str, List[Tensor]],
    save_pth: str,
    num_timestamps: int,
    keys: List[str] = ["gt_rgbs", "rgbs", "depths"],
    num_cams: int = 3,
    save_seperate_video: bool = False,
    save_images: bool = False,
    fps: int = 10,
    verbose: bool = True,
):
    if save_seperate_video:
        return_frame = save_seperate_videos(
            render_results,
            save_pth,
            num_timestamps=num_timestamps,
            keys=keys,
            num_cams=num_cams,
            save_images=save_images,
            fps=fps,
            verbose=verbose,
        )
    else:
        return_frame = save_concatenated_videos(
            render_results,
            save_pth,
            num_timestamps=num_timestamps,
            keys=keys,
            num_cams=num_cams,
            save_images=save_images,
            fps=fps,
            verbose=verbose,
        )
    return return_frame


def save_concatenated_videos(
    render_results: Dict[str, List[Tensor]],
    save_pth: str,
    num_timestamps: int,
    keys: List[str] = ["gt_rgbs", "rgbs", "depths"],
    num_cams: int = 3,
    save_images: bool = False,
    fps: int = 10,
    verbose: bool = True,
):
    if num_timestamps == 1:  # it's an image
        writer = imageio.get_writer(save_pth, mode="I")
        return_frame_id = 0
    else:
        return_frame_id = num_timestamps // 2
        writer = imageio.get_writer(save_pth, mode="I", fps=fps)
    for i in trange(num_timestamps, desc="saving video", dynamic_ncols=True):
        merged_list = []
        for key in keys:
            if key == "sky_masks":
                frames = render_results["opacities"][i * num_cams : (i + 1) * num_cams]
            else:
                if key not in render_results or len(render_results[key]) == 0:
                    continue
                frames = render_results[key][i * num_cams : (i + 1) * num_cams]
            if key == "gt_sky_masks":
                frames = [np.stack([frame, frame, frame], axis=-1) for frame in frames]
            elif key == "sky_masks":
                frames = [
                    1 - np.stack([frame, frame, frame], axis=-1) for frame in frames
                ]
            # elif "depth" in key:

            #     frames = [
            #         depth_visualizer(frame, opacity)
            #         for frame, opacity in zip(frames, opacities)
            #     ]
            frames = resize_five_views(frames)
            frames = np.concatenate(frames, axis=1)
            merged_list.append(frames)
        merged_frame = to8b(np.concatenate(merged_list, axis=0))
        if i == return_frame_id:
            return_frame = merged_frame
        writer.append_data(merged_frame)
    writer.close()
    if verbose:
        print(f"saved video to {save_pth}")
    del render_results
    return {"concatenated_frame": return_frame}


def save_seperate_videos(
    render_results: Dict[str, List[Tensor]],
    save_pth: str,
    num_timestamps: int,
    keys: List[str] = ["gt_rgbs", "rgbs", "depths"],
    num_cams: int = 3,
    fps: int = 10,
    verbose: bool = False,
    save_images: bool = False,
):
    return_frame_id = num_timestamps // 2
    return_frame_dict = {}
    for key in keys:
        tmp_save_pth = save_pth.replace(".mp4", f"_{key}.mp4")
        tmp_save_pth = tmp_save_pth.replace(".png", f"_{key}.png")
        if num_timestamps == 1:  # it's an image
            writer = imageio.get_writer(tmp_save_pth, mode="I")
        else:
            writer = imageio.get_writer(tmp_save_pth, mode="I", fps=fps)
        if key not in render_results or len(render_results[key]) == 0:
            continue
        for i in range(num_timestamps):
            if key == "sky_masks":
                frames = render_results["opacities"][i * num_cams : (i + 1) * num_cams]
            else:
                # 这里取3个，得到3个视角
                frames = render_results[key][i * num_cams : (i + 1) * num_cams]
            if key == "gt_sky_masks":
                frames = [np.stack([frame, frame, frame], axis=-1) for frame in frames]
            elif key == "sky_masks":
                frames = [
                    1 - np.stack([frame, frame, frame], axis=-1) for frame in frames
                ]
            # elif "depth" in key:
            #     opacities = render_results[key.replace("depths", "opacities")][
            #         i * num_cams : (i + 1) * num_cams
            #     ]
            #     frames = [
            #         depth_visualizer(frame, opacity)
            #         for frame, opacity in zip(frames, opacities)
            #     ]
            frames = resize_five_views(frames)
            if save_images:
                image_dir = tmp_save_pth.replace(".mp4", "")
                if i == 0:
                    os.makedirs(image_dir, exist_ok=True)
                for j, frame in enumerate(frames):
                    image = to8b(frame)
                    if image.ndim == 2:
                        image = np.stack([image, image, image], axis=-1)
                    elif image.ndim == 3 and image.shape[-1] == 1:
                        image = np.repeat(image, 3, axis=-1)
                    imageio.imwrite(
                        os.path.join(image_dir, f"{i * num_cams + j:03d}.png"),
                        image,
                    )
            frames = to8b(np.concatenate(frames, axis=1))
            writer.append_data(frames) # [H,W,3]
            if i == return_frame_id:
                return_frame_dict[key] = frames
        # close the writer
        writer.close()
        del writer
        if verbose:
            print(f"saved video to {tmp_save_pth}")
    del render_results
    return return_frame_dict