Dororo99/Ours_S3GS_Waymo

1import logging2import os3from collections import namedtuple4from itertools import accumulate5from typing import List, Optional, Union6 7import matplotlib.cm as cm8import numpy as np9import plotly.graph_objects as go10import torch11# from omegaconf import OmegaConf12from scipy import ndimage13from tqdm import tqdm14import json15import cv216 17# from datasets import SceneDataset18# from datasets.utils import voxel_coords_to_world_coords, world_coords_to_voxel_coords19# from radiance_fields import DensityField, RadianceField20# from radiance_fields.render_utils import render_rays21# from third_party.nerfacc_prop_net import PropNetEstimator22# from utils.misc import get_robust_pca23# from utils.misc import NumpyEncoder24 25DEFAULT_TRANSITIONS = (15, 6, 4, 11, 13, 6)26 27logger = logging.getLogger()28turbo_cmap = cm.get_cmap("turbo")29 30# 定义函数用于将光流可视化为RGB颜色31def flow_to_color(flow):32    hsv = np.zeros((flow.shape[0], flow.shape[1], 3), dtype=np.uint8)33    hsv[..., 1] = 25534 35    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])36    hsv[..., 0] = ang * 180 / np.pi / 237    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)38 39    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)40 41# 定义函数用于计算光流并保存可视化结果42def compute_optical_flow_and_save(frames, output_path):43    # 初始化光流计算器44    prev_frame = frames[0]45    hsv = np.zeros_like(prev_frame)46    hsv[..., 1] = 25547 48    # 逐帧计算光流并保存可视化结果49    for i in range(1, len(frames)):50        next_frame = frames[i]51 52        # 计算光流53        flow = cv2.calcOpticalFlowFarneback(54            cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY),55            cv2.cvtColor(next_frame, cv2.COLOR_BGR2GRAY),56            None, 0.5, 5, 15, 5, 7, 1.5, 0)57 58        # 将光流转换为RGB颜色59        flow_rgb = flow_to_color(flow)60 61        # 保存可视化结果62        cv2.imwrite(f"{output_path}/optical_flow_{i}.jpg", flow_rgb)63 64        # 更新前一帧65        prev_frame = next_frame66 67 68def to8b(x):69    if isinstance(x, torch.Tensor):70        x = x.detach().cpu().numpy()71    return (255 * np.clip(x, 0, 1)).astype(np.uint8)72 73 74def resize_five_views(imgs: np.array):75    if len(imgs) != 5:76        return imgs77    for idx in [0, -1]:78        img = imgs[idx]79        new_shape = [int(img.shape[1] * 0.46), img.shape[1], 3]80        new_img = np.zeros_like(img)81        new_img[-new_shape[0] :, : new_shape[1], :] = ndimage.zoom(82            img, [new_shape[0] / img.shape[0], new_shape[1] / img.shape[1], 1]83        )84        # clip the image to 0-185        new_img = np.clip(new_img, 0, 1)86        imgs[idx] = new_img87    return imgs88 89 90def sinebow(h):91    """A cyclic and uniform colormap, see http://basecase.org/env/on-rainbows."""92    f = lambda x: np.sin(np.pi * x) ** 293    return np.stack([f(3 / 6 - h), f(5 / 6 - h), f(7 / 6 - h)], -1)94 95 96def matte(vis, acc, dark=0.8, light=1.0, width=8):97    """Set non-accumulated pixels to a Photoshop-esque checker pattern."""98    bg_mask = np.logical_xor(99        (np.arange(acc.shape[0]) % (2 * width) // width)[:, None],100        (np.arange(acc.shape[1]) % (2 * width) // width)[None, :],101    )102    bg = np.where(bg_mask, light, dark)103    return vis * acc[:, :, None] + (bg * (1 - acc))[:, :, None]104 105 106def weighted_percentile(x, w, ps, assume_sorted=False):107    """Compute the weighted percentile(s) of a single vector."""108    x = x.reshape([-1])109    w = w.reshape([-1])110    if not assume_sorted:111        sortidx = np.argsort(x)112        x, w = x[sortidx], w[sortidx]113    acc_w = np.cumsum(w)114    return np.interp(np.array(ps) * (acc_w[-1] / 100), acc_w, x)115 116 117def visualize_cmap(118    value,119    weight,120    colormap,121    lo=None,122    hi=None,123    percentile=99.0,124    curve_fn=lambda x: x,125    modulus=None,126    matte_background=True,127):128    """Visualize a 1D image and a 1D weighting according to some colormap.129    from mipnerf130 131    Args:132      value: A 1D image.133      weight: A weight map, in [0, 1].134      colormap: A colormap function.135      lo: The lower bound to use when rendering, if None then use a percentile.136      hi: The upper bound to use when rendering, if None then use a percentile.137      percentile: What percentile of the value map to crop to when automatically138        generating `lo` and `hi`. Depends on `weight` as well as `value'.139      curve_fn: A curve function that gets applied to `value`, `lo`, and `hi`140        before the rest of visualization. Good choices: x, 1/(x+eps), log(x+eps).141      modulus: If not None, mod the normalized value by `modulus`. Use (0, 1]. If142        `modulus` is not None, `lo`, `hi` and `percentile` will have no effect.143      matte_background: If True, matte the image over a checkerboard.144 145    Returns:146      A colormap rendering.147    """148    # Identify the values that bound the middle of `value' according to `weight`.149    if lo is None or hi is None:150        lo_auto, hi_auto = weighted_percentile(151            value, weight, [50 - percentile / 2, 50 + percentile / 2]152        )153        # If `lo` or `hi` are None, use the automatically-computed bounds above.154        eps = np.finfo(np.float32).eps155        lo = lo or (lo_auto - eps)156        hi = hi or (hi_auto + eps)157 158    # Curve all values.159    value, lo, hi = [curve_fn(x) for x in [value, lo, hi]]160 161    # Wrap the values around if requested.162    if modulus:163        value = np.mod(value, modulus) / modulus164    else:165        # Otherwise, just scale to [0, 1].166        value = np.nan_to_num(167            np.clip((value - np.minimum(lo, hi)) / np.abs(hi - lo), 0, 1)168        )169    if weight is not None:170        value *= weight171    else:172        weight = np.ones_like(value)173    if colormap:174        colorized = colormap(value)[..., :3]175    else:176        assert len(value.shape) == 3 and value.shape[-1] == 3177        colorized = value178 179    return matte(colorized, weight) if matte_background else colorized180 181 182def visualize_depth(183    x, acc=None, lo=None, hi=None, depth_curve_fn=lambda x: -np.log(x + 1e-6)184):185    """Visualizes depth maps."""186    return visualize_cmap(187        x,188        acc,189        cm.get_cmap("turbo"),190        curve_fn=depth_curve_fn,191        lo=lo,192        hi=hi,193        matte_background=False,194    )195 196 197def _make_colorwheel(transitions: tuple = DEFAULT_TRANSITIONS) -> torch.Tensor:198    """Creates a colorwheel (borrowed/modified from flowpy).199    A colorwheel defines the transitions between the six primary hues:200    Red(255, 0, 0), Yellow(255, 255, 0), Green(0, 255, 0), Cyan(0, 255, 255), Blue(0, 0, 255) and Magenta(255, 0, 255).201    Args:202        transitions: Contains the length of the six transitions, based on human color perception.203    Returns:204        colorwheel: The RGB values of the transitions in the color space.205    Notes:206        For more information, see:207        https://web.archive.org/web/20051107102013/http://members.shaw.ca/quadibloc/other/colint.htm208        http://vision.middlebury.edu/flow/flowEval-iccv07.pdf209    """210    colorwheel_length = sum(transitions)211    # The red hue is repeated to make the colorwheel cyclic212    base_hues = map(213        np.array,214        (215            [255, 0, 0],216            [255, 255, 0],217            [0, 255, 0],218            [0, 255, 255],219            [0, 0, 255],220            [255, 0, 255],221            [255, 0, 0],222        ),223    )224    colorwheel = np.zeros((colorwheel_length, 3), dtype="uint8")225    hue_from = next(base_hues)226    start_index = 0227    for hue_to, end_index in zip(base_hues, accumulate(transitions)):228        transition_length = end_index - start_index229        colorwheel[start_index:end_index] = np.linspace(230            hue_from, hue_to, transition_length, endpoint=False231        )232        hue_from = hue_to233        start_index = end_index234    return torch.FloatTensor(colorwheel)235 236 237WHEEL = _make_colorwheel()238N_COLS = len(WHEEL)239WHEEL = torch.vstack((WHEEL, WHEEL[0]))  # Make the wheel cyclic for interpolation240 241 242def scene_flow_to_rgb(243    flow: torch.Tensor,244    flow_max_radius: Optional[float] = None,245    background: Optional[str] = "dark",246) -> torch.Tensor:247    """Creates a RGB representation of an optical flow (borrowed/modified from flowpy).248    Adapted from https://github.com/Lilac-Lee/Neural_Scene_Flow_Prior/blob/main/visualize.py249    Args:250        flow: scene flow.251            flow[..., 0] should be the x-displacement252            flow[..., 1] should be the y-displacement253            flow[..., 2] should be the z-displacement254        flow_max_radius: Set the radius that gives the maximum color intensity, useful for comparing different flows.255            Default: The normalization is based on the input flow maximum radius.256        background: States if zero-valued flow should look 'bright' or 'dark'.257    Returns: An array of RGB colors.258    """259    flow_min = flow.min()  # 找到最小值260    flow_max = flow.max()  # 找到最大值261    eps = 1e-6  # 一个小常数，防止除以零262    flow = (flow - flow_min) / (flow_max - flow_min + eps)  # 归一化，避免除零错误263    # flow = flow * 100264 265    valid_backgrounds = ("bright", "dark")266    if background not in valid_backgrounds:267        raise ValueError(268            f"background should be one the following: {valid_backgrounds}, not {background}."269        )270 271    # For scene flow, it's reasonable to assume displacements in x and y directions only for visualization pursposes.272    complex_flow = flow[..., 0] + 1j * flow[..., 1]273    radius, angle = torch.abs(complex_flow), torch.angle(complex_flow)274    if flow_max_radius is None:275        # flow_max_radius = torch.max(radius)276        flow_max_radius = torch.quantile(radius, 0.99)277    if flow_max_radius > 0:278        radius /= flow_max_radius279    # Map the angles from (-pi, pi] to [0, 2pi) to [0, ncols - 1)280    angle[angle < 0] += 2 * np.pi281    angle = angle * ((N_COLS - 1) / (2 * np.pi))282 283    # Interpolate the hues284    angle_fractional, angle_floor, angle_ceil = (285        torch.fmod(angle, 1),286        angle.trunc(),287        torch.ceil(angle),288    )289    angle_fractional = angle_fractional.unsqueeze(-1)290    wheel = WHEEL.to(angle_floor.device)291    float_hue = (292        wheel[angle_floor.long()] * (1 - angle_fractional)293        + wheel[angle_ceil.long()] * angle_fractional294    )295    ColorizationArgs = namedtuple(296        "ColorizationArgs",297        ["move_hue_valid_radius", "move_hue_oversized_radius", "invalid_color"],298    )299 300    def move_hue_on_V_axis(hues, factors):301        return hues * factors.unsqueeze(-1)302 303    def move_hue_on_S_axis(hues, factors):304        return 255.0 - factors.unsqueeze(-1) * (255.0 - hues)305 306    if background == "dark":307        parameters = ColorizationArgs(308            move_hue_on_V_axis, move_hue_on_S_axis, torch.FloatTensor([255, 255, 255])309        )310    else:311        parameters = ColorizationArgs(312            move_hue_on_S_axis, move_hue_on_V_axis, torch.zeros(3)313        )314    colors = parameters.move_hue_valid_radius(float_hue, radius)315    oversized_radius_mask = radius > 1316    colors[oversized_radius_mask] = parameters.move_hue_oversized_radius(317        float_hue[oversized_radius_mask], 1 / radius[oversized_radius_mask]318    )319    # print(colors.max())320    # print(colors.min())321 322    return colors / 255.0323 324 325def vis_occ_plotly(326    vis_aabb: List[Union[int, float]],327    coords: np.array = None,328    colors: np.array = None,329    dynamic_coords: List[np.array] = None,330    dynamic_colors: List[np.array] = None,331    x_ratio: float = 1.0,332    y_ratio: float = 1.0,333    z_ratio: float = 0.125,334    size: int = 5,335    black_bg: bool = False,336    title: str = None,337) -> go.Figure:  # type: ignore338    fig = go.Figure()  # start with an empty figure339 340    if coords is not None:341        # Add static trace342        static_trace = go.Scatter3d(343            x=coords[:, 0],344            y=coords[:, 1],345            z=coords[:, 2],346            mode="markers",347            marker=dict(348                size=size,349                color=colors,350                symbol="square",351            ),352        )353        fig.add_trace(static_trace)354 355    # Add temporal traces356    if dynamic_coords is not None:357        for i in range(len(dynamic_coords)):358            fig.add_trace(359                go.Scatter3d(360                    x=dynamic_coords[i][:, 0],361                    y=dynamic_coords[i][:, 1],362                    z=dynamic_coords[i][:, 2],363                    mode="markers",364                    marker=dict(365                        size=size,366                        color=dynamic_colors[i],367                        symbol="diamond",368                    ),369                )370            )371        steps = []372        if coords is not None:373            for i in range(len(dynamic_coords)):374                step = dict(375                    method="restyle",376                    args=[377                        "visible",378                        [False] * (len(dynamic_coords) + 1),379                    ],  # Include the static trace380                    label=f"Second {i}",381                )382                step["args"][1][0] = True  # Make the static trace always visible383                step["args"][1][i + 1] = True  # Toggle i'th temporal trace to "visible"384                steps.append(step)385        else:386            for i in range(len(dynamic_coords)):387                step = dict(388                    method="restyle",389                    args=[390                        "visible",391                        [False] * (len(dynamic_coords)),392                    ],393                    label=f"Second {i}",394                )395                step["args"][1][i] = True  # Toggle i'th temporal trace to "visible"396                steps.append(step)397 398        sliders = [399            dict(400                active=0,401                pad={"t": 1},402                steps=steps,403                font=dict(color="white") if black_bg else {},  # Update for font color404            )405        ]406        fig.update_layout(sliders=sliders)407    title_font_color = "white" if black_bg else "black"408    if not black_bg:409        fig.update_layout(410            scene=dict(411                xaxis=dict(412                    title="x",413                    showspikes=False,414                    range=[vis_aabb[0], vis_aabb[3]],415                ),416                yaxis=dict(417                    title="y",418                    showspikes=False,419                    range=[vis_aabb[1], vis_aabb[4]],420                ),421                zaxis=dict(422                    title="z",423                    showspikes=False,424                    range=[vis_aabb[2], vis_aabb[5]],425                ),426                aspectmode="manual",427                aspectratio=dict(x=x_ratio, y=y_ratio, z=z_ratio),428            ),429            margin=dict(r=0, b=10, l=0, t=10),430            hovermode=False,431            title=dict(432                text=title,433                font=dict(color=title_font_color),434                x=0.5,435                y=0.95,436                xanchor="center",437                yanchor="top",438            )439            if title440            else None,  # Title addition441        )442    else:443        fig.update_layout(444            scene=dict(445                xaxis=dict(446                    title="x",447                    showspikes=False,448                    range=[vis_aabb[0], vis_aabb[3]],449                    backgroundcolor="rgb(0, 0, 0)",450                    gridcolor="gray",451                    showbackground=True,452                    zerolinecolor="gray",453                    tickfont=dict(color="gray"),454                ),455                yaxis=dict(456                    title="y",457                    showspikes=False,458                    range=[vis_aabb[1], vis_aabb[4]],459                    backgroundcolor="rgb(0, 0, 0)",460                    gridcolor="gray",461                    showbackground=True,462                    zerolinecolor="gray",463                    tickfont=dict(color="gray"),464                ),465                zaxis=dict(466                    title="z",467                    showspikes=False,468                    range=[vis_aabb[2], vis_aabb[5]],469                    backgroundcolor="rgb(0, 0, 0)",470                    gridcolor="gray",471                    showbackground=True,472                    zerolinecolor="gray",473                    tickfont=dict(color="gray"),474                ),475                aspectmode="manual",476                aspectratio=dict(x=x_ratio, y=y_ratio, z=z_ratio),477            ),478            margin=dict(r=0, b=10, l=0, t=10),479            hovermode=False,480            paper_bgcolor="black",481            plot_bgcolor="rgba(0,0,0,0)",482            title=dict(483                text=title,484                font=dict(color=title_font_color),485                x=0.5,486                y=0.95,487                xanchor="center",488                yanchor="top",489            )490            if title491            else None,  # Title addition492        )493    eye = np.array([-1, 0, 0.5])494    eye = eye.tolist()495    fig.update_layout(496        scene_camera=dict(497            eye=dict(x=eye[0], y=eye[1], z=eye[2]),498        ),499    )500    return fig