dva/visualize.py

import cv2
import os
import numpy as np
import torch
import imageio
from torchvision.utils import make_grid, save_image
from .ray_marcher import RayMarcher, generate_colored_boxes

def get_pose_on_orbit(radius, height, angles, world_up=torch.Tensor([0, 1, 0])):
    num_points = angles.shape[0]
    x = radius * torch.cos(angles)
    h = torch.ones((num_points,)) * height
    z = radius * torch.sin(angles)
    position = torch.stack([x, h, z], dim=-1)
    forward = position / torch.norm(position, p=2, dim=-1, keepdim=True)
    right = -torch.cross(world_up[None, ...], forward)
    right /= torch.norm(right, dim=-1, keepdim=True)
    up = torch.cross(forward, right)
    up /= torch.norm(up, p=2, dim=-1, keepdim=True)
    rotation = torch.stack([right, up, forward], dim=1)
    translation = torch.Tensor([0, 0, radius])[None, :, None].repeat(num_points, 1, 1)
    return torch.concat([rotation, translation], dim=2)

def render_mvp_boxes(rm, batch, preds):
    with torch.no_grad():
        boxes_rgba = generate_colored_boxes(
            preds["prim_rgba"],
            preds["prim_rot"],
        )
        preds_boxes = rm(
            prim_rgba=boxes_rgba,
            prim_pos=preds["prim_pos"],
            prim_scale=preds["prim_scale"],
            prim_rot=preds["prim_rot"],
            RT=batch["Rt"],
            K=batch["K"],
        )

    return preds_boxes["rgba_image"][:, :3].permute(0, 2, 3, 1)


def save_image_summary(path, batch, preds):
    rgb = preds["rgb"].detach().permute(0, 3, 1, 2)
    # rgb_gt = batch["image"]
    rgb_boxes = preds["rgb_boxes"].detach().permute(0, 3, 1, 2)
    bs = rgb_boxes.shape[0]
    if "folder" in batch and "key" in batch:
        obj_list = []
        for bs_idx in range(bs):
            tmp_img = rgb_boxes[bs_idx].permute(1, 2, 0).to(torch.uint8).cpu().numpy()
            tmp_img = np.ascontiguousarray(tmp_img)
            folder = batch['folder'][bs_idx]
            key = batch['key'][bs_idx]
            obj_list.append("{}/{}\n".format(folder, key))
            cv2.putText(tmp_img, "{}".format(folder), (200, 200), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 2)
            cv2.putText(tmp_img, "{}".format(key), (200, 400), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 2)
            tmp_img_torch = torch.as_tensor(tmp_img).permute(2, 0, 1).float()
            rgb_boxes[bs_idx] = tmp_img_torch
        with open(os.path.splitext(path)[0]+".txt", "w") as f:
            f.writelines(obj_list)
    img = make_grid(torch.cat([rgb, rgb_boxes], dim=2) / 255.0).clip(0.0, 1.0)
    save_image(img, path)


@torch.no_grad()
def visualize_primsdf_box(image_save_path, model, rm: RayMarcher, device):
    # prim_rgba: primitive payload [B, K, 4, S, S, S],
    # K - # of primitives, S - primitive size
    # prim_pos: locations [B, K, 3]
    # prim_rot: rotations [B, K, 3, 3]
    # prim_scale: scales [B, K, 3]
    # K: intrinsics [B, 3, 3]
    # RT: extrinsics [B, 3, 4]
    preds = {}
    batch = {}
    prim_alpha = model.sdf2alpha(model.feat_geo).reshape(1, model.num_prims, 1, model.prim_shape, model.prim_shape, model.prim_shape) * 255
    prim_rgb = model.feat_tex.reshape(1, model.num_prims, 3, model.prim_shape, model.prim_shape, model.prim_shape) * 255
    preds['prim_rgba'] = torch.concat([prim_rgb, prim_alpha], dim=2)
    preds['prim_pos'] = model.pos.reshape(1, model.num_prims, 3) * rm.volradius
    preds['prim_rot'] = torch.eye(3).to(preds['prim_pos'])[None, None, ...].repeat(1, model.num_prims, 1, 1)
    preds['prim_scale'] = (1 / model.scale.reshape(1, model.num_prims, 1).repeat(1, 1, 3))
    batch['Rt'] = torch.Tensor([
            [
                1.0,
                0.0,
                0.0,
                0.0 * rm.volradius
            ],
            [
                0.0,
                -1.0,
                0.0,
                0.0 * rm.volradius
            ],
            [
                0.0,
                0.0,
                -1.0,
                5 * rm.volradius
            ]
        ]).to(device)[None, ...]
    batch['K'] = torch.Tensor([
            [
                2084.9526697685183,
                0.0,
                512.0
            ],
            [
                0.0,
                2084.9526697685183,
                512.0
            ],
            [
                0.0,
                0.0,
                1.0
            ]]).to(device)[None, ...]
    ratio_h = rm.image_height / 1024.
    ratio_w = rm.image_width / 1024.
    batch['K'][:, 0:1, :] *= ratio_h
    batch['K'][:, 1:2, :] *= ratio_w
    # raymarcher is in mm
    rm_preds = rm(
        prim_rgba=preds["prim_rgba"],
        prim_pos=preds["prim_pos"],
        prim_scale=preds["prim_scale"],
        prim_rot=preds["prim_rot"],
        RT=batch["Rt"],
        K=batch["K"],
    )
    rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)
    preds.update(alpha=rgba[..., -1].contiguous(), rgb=rgba[..., :3].contiguous())
    with torch.no_grad():
        preds["rgb_boxes"] = render_mvp_boxes(rm, batch, preds)
    save_image_summary(image_save_path, batch, preds)

@torch.no_grad()
def render_primsdf(image_save_path, model, rm, device):
    preds = {}
    batch = {}
    preds['prim_pos'] = model.pos.reshape(1, model.num_prims, 3) * rm.volradius
    preds['prim_rot'] = torch.eye(3).to(preds['prim_pos'])[None, None, ...].repeat(1, model.num_prims, 1, 1)
    preds['prim_scale'] = (1 / model.scale.reshape(1, model.num_prims, 1).repeat(1, 1, 3))
    batch['Rt'] = torch.Tensor([
            [
                1.0,
                0.0,
                0.0,
                0.0 * rm.volradius
            ],
            [
                0.0,
                -1.0,
                0.0,
                0.0 * rm.volradius
            ],
            [
                0.0,
                0.0,
                -1.0,
                5 * rm.volradius
            ]
        ]).to(device)[None, ...]
    batch['K'] = torch.Tensor([
            [
                2084.9526697685183,
                0.0,
                512.0
            ],
            [
                0.0,
                2084.9526697685183,
                512.0
            ],
            [
                0.0,
                0.0,
                1.0
            ]]).to(device)[None, ...]
    ratio_h = rm.image_height / 1024.
    ratio_w = rm.image_width / 1024.
    batch['K'][:, 0:1, :] *= ratio_h
    batch['K'][:, 1:2, :] *= ratio_w
    # test rendering
    all_sampled_sdf = []
    all_sampled_tex = []
    for i in range(model.prim_shape ** 3):
        with torch.no_grad():
            model_prediction = model(model.sdf_sampled_point[:, i, :].to(device))
            sampled_sdf = model_prediction['sdf']
            sampled_rgb = model_prediction['tex']
            all_sampled_sdf.append(sampled_sdf)
            all_sampled_tex.append(sampled_rgb)
    sampled_sdf = torch.stack(all_sampled_sdf, dim=1)
    sampled_tex = torch.stack(all_sampled_tex, dim=1).permute(0, 2, 1).reshape(1, model.num_prims, 3, model.prim_shape, model.prim_shape, model.prim_shape) * 255
    prim_rgb = sampled_tex
    prim_alpha = model.sdf2alpha(sampled_sdf).reshape(1, model.num_prims, 1, model.prim_shape, model.prim_shape, model.prim_shape) * 255
    preds['prim_rgba'] = torch.concat([prim_rgb, prim_alpha], dim=2)
    rm_preds = rm(
        prim_rgba=preds["prim_rgba"],
        prim_pos=preds["prim_pos"],
        prim_scale=preds["prim_scale"],
        prim_rot=preds["prim_rot"],
        RT=batch["Rt"],
        K=batch["K"],
    )

    rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)
    preds.update(alpha=rgba[..., -1].contiguous(), rgb=rgba[..., :3].contiguous())
    with torch.no_grad():
        preds["rgb_boxes"] = render_mvp_boxes(rm, batch, preds)
    save_image_summary(image_save_path, batch, preds)

@torch.no_grad()
def visualize_primvolume(image_save_path, batch, prim_volume, rm: RayMarcher, device):
    # prim_volume - [B, nprims, 4+6*8^3]
    def sdf2alpha(sdf):
        return torch.exp(-(sdf / 0.005) ** 2)
    preds = {}
    prim_shape = int(np.round(((prim_volume.shape[2] - 4) / 6) ** (1/3)))
    num_prims = prim_volume.shape[1]
    bs = prim_volume.shape[0]
    geo_start_index = 4
    geo_end_index = geo_start_index + prim_shape ** 3 # non-inclusive
    tex_start_index = geo_end_index
    tex_end_index = tex_start_index + prim_shape ** 3 * 3 # non-inclusive
    mat_start_index = tex_end_index
    mat_end_index = mat_start_index + prim_shape ** 3 * 2

    feat_geo = prim_volume[:, :, geo_start_index: geo_end_index]
    feat_tex = prim_volume[:, :, tex_start_index: tex_end_index]
    prim_alpha = sdf2alpha(feat_geo).reshape(bs, num_prims, 1, prim_shape, prim_shape, prim_shape) * 255
    prim_rgb = feat_tex.reshape(bs, num_prims, 3, prim_shape, prim_shape, prim_shape) * 255
    preds['prim_rgba'] = torch.concat([prim_rgb, prim_alpha], dim=2)
    pos = prim_volume[:, :, 1:4]
    scale = prim_volume[:, :, 0:1]
    preds['prim_pos'] = pos.reshape(bs, num_prims, 3) * rm.volradius
    preds['prim_rot'] = torch.eye(3).to(preds['prim_pos'])[None, None, ...].repeat(bs, num_prims, 1, 1)
    preds['prim_scale'] = (1 / scale.reshape(bs, num_prims, 1).repeat(1, 1, 3))
    batch['Rt'] = torch.Tensor([
            [
                1.0,
                0.0,
                0.0,
                0.0 * rm.volradius
            ],
            [
                0.0,
                -1.0,
                0.0,
                0.0 * rm.volradius
            ],
            [
                0.0,
                0.0,
                -1.0,
                5 * rm.volradius
            ]
        ]).to(device)[None, ...].repeat(bs, 1, 1)
    batch['K'] = torch.Tensor([
            [
                2084.9526697685183,
                0.0,
                512.0
            ],
            [
                0.0,
                2084.9526697685183,
                512.0
            ],
            [
                0.0,
                0.0,
                1.0
            ]]).to(device)[None, ...].repeat(bs, 1, 1)
    ratio_h = rm.image_height / 1024.
    ratio_w = rm.image_width / 1024.
    batch['K'][:, 0:1, :] *= ratio_h
    batch['K'][:, 1:2, :] *= ratio_w
    # raymarcher is in mm
    rm_preds = rm(
        prim_rgba=preds["prim_rgba"],
        prim_pos=preds["prim_pos"],
        prim_scale=preds["prim_scale"],
        prim_rot=preds["prim_rot"],
        RT=batch["Rt"],
        K=batch["K"],
    )
    rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)
    preds.update(alpha=rgba[..., -1].contiguous(), rgb=rgba[..., :3].contiguous())
    with torch.no_grad():
        preds["rgb_boxes"] = render_mvp_boxes(rm, batch, preds)
    save_image_summary(image_save_path, batch, preds)

@torch.no_grad()
def visualize_multiview_primvolume(image_save_path, batch, prim_volume, view_counts, rm: RayMarcher, device):
    # prim_volume - [B, nprims, 4+6*8^3]
    view_angles = torch.linspace(0.5, 2.5, view_counts + 1) * torch.pi
    view_angles = view_angles[:-1]
    def sdf2alpha(sdf):
        return torch.exp(-(sdf / 0.005) ** 2)
    preds = {}
    prim_shape = int(np.round(((prim_volume.shape[2] - 4) / 6) ** (1/3)))
    num_prims = prim_volume.shape[1]
    bs = prim_volume.shape[0]
    geo_start_index = 4
    geo_end_index = geo_start_index + prim_shape ** 3 # non-inclusive
    tex_start_index = geo_end_index
    tex_end_index = tex_start_index + prim_shape ** 3 * 3 # non-inclusive
    mat_start_index = tex_end_index
    mat_end_index = mat_start_index + prim_shape ** 3 * 2

    feat_geo = prim_volume[:, :, geo_start_index: geo_end_index]
    feat_tex = prim_volume[:, :, tex_start_index: tex_end_index]
    prim_alpha = sdf2alpha(feat_geo).reshape(bs, num_prims, 1, prim_shape, prim_shape, prim_shape) * 255
    prim_rgb = feat_tex.reshape(bs, num_prims, 3, prim_shape, prim_shape, prim_shape) * 255
    preds['prim_rgba'] = torch.concat([prim_rgb, prim_alpha], dim=2)
    pos = prim_volume[:, :, 1:4]
    scale = prim_volume[:, :, 0:1]
    preds['prim_pos'] = pos.reshape(bs, num_prims, 3) * rm.volradius
    preds['prim_rot'] = torch.eye(3).to(preds['prim_pos'])[None, None, ...].repeat(bs, num_prims, 1, 1)
    preds['prim_scale'] = (1 / scale.reshape(bs, num_prims, 1).repeat(1, 1, 3))
    batch['K'] = torch.Tensor([
            [
                2084.9526697685183,
                0.0,
                512.0
            ],
            [
                0.0,
                2084.9526697685183,
                512.0
            ],
            [
                0.0,
                0.0,
                1.0
            ]]).to(device)[None, ...].repeat(bs, 1, 1)
    ratio_h = rm.image_height / 1024.
    ratio_w = rm.image_width / 1024.
    batch['K'][:, 0:1, :] *= ratio_h
    batch['K'][:, 1:2, :] *= ratio_w

    final_preds = {}
    final_preds['rgb'] = []
    final_preds['rgb_boxes'] = []
    for view_ang in view_angles:
        bs_view_ang = view_ang.repeat(bs,)
        batch['Rt'] = get_pose_on_orbit(radius=5*rm.volradius, height=0, angles=bs_view_ang).to(prim_volume)
        # raymarcher is in mm
        rm_preds = rm(
            prim_rgba=preds["prim_rgba"],
            prim_pos=preds["prim_pos"],
            prim_scale=preds["prim_scale"],
            prim_rot=preds["prim_rot"],
            RT=batch["Rt"],
            K=batch["K"],
        )
        rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)
        preds.update(alpha=rgba[..., -1].contiguous(), rgb=rgba[..., :3].contiguous())
        with torch.no_grad():
            preds["rgb_boxes"] = render_mvp_boxes(rm, batch, preds)
        final_preds['rgb'].append(preds['rgb'])
        final_preds['rgb_boxes'].append(preds['rgb_boxes'])
    final_preds['rgb'] = torch.concat(final_preds['rgb'], dim=0)
    final_preds['rgb_boxes'] = torch.concat(final_preds['rgb_boxes'], dim=0)
    save_image_summary(image_save_path, batch, final_preds)


@torch.no_grad()
def visualize_video_primvolume(video_save_folder, batch, prim_volume, view_counts, rm: RayMarcher, device):
    # prim_volume - [B, nprims, 4+6*8^3]
    view_angles = torch.linspace(1.5, 3.5, view_counts + 1) * torch.pi
    def sdf2alpha(sdf):
        return torch.exp(-(sdf / 0.005) ** 2)
    preds = {}
    prim_shape = int(np.round(((prim_volume.shape[2] - 4) / 6) ** (1/3)))
    num_prims = prim_volume.shape[1]
    bs = prim_volume.shape[0]
    geo_start_index = 4
    geo_end_index = geo_start_index + prim_shape ** 3 # non-inclusive
    tex_start_index = geo_end_index
    tex_end_index = tex_start_index + prim_shape ** 3 * 3 # non-inclusive
    mat_start_index = tex_end_index
    mat_end_index = mat_start_index + prim_shape ** 3 * 2

    feat_geo = prim_volume[:, :, geo_start_index: geo_end_index]
    feat_tex = prim_volume[:, :, tex_start_index: tex_end_index]
    feat_mat = prim_volume[:, :, mat_start_index: mat_end_index]
    prim_alpha = sdf2alpha(feat_geo).reshape(bs, num_prims, 1, prim_shape, prim_shape, prim_shape) * 255
    prim_rgb = feat_tex.reshape(bs, num_prims, 3, prim_shape, prim_shape, prim_shape) * 255
    prim_mat = feat_mat.reshape(bs, num_prims, 2, prim_shape, prim_shape, prim_shape) * 255
    dummy_prim = torch.zeros_like(prim_mat[:, :, 0:1, ...])
    prim_mat = torch.concat([dummy_prim, prim_mat], dim=2)
    preds['prim_rgba'] = torch.concat([prim_rgb, prim_alpha], dim=2)
    preds['prim_mata'] = torch.concat([prim_mat, prim_alpha], dim=2)
    pos = prim_volume[:, :, 1:4]
    scale = prim_volume[:, :, 0:1]
    preds['prim_pos'] = pos.reshape(bs, num_prims, 3) * rm.volradius
    preds['prim_rot'] = torch.eye(3).to(preds['prim_pos'])[None, None, ...].repeat(bs, num_prims, 1, 1)
    preds['prim_scale'] = (1 / scale.reshape(bs, num_prims, 1).repeat(1, 1, 3))
    batch['K'] = torch.Tensor([
            [
                2084.9526697685183,
                0.0,
                512.0
            ],
            [
                0.0,
                2084.9526697685183,
                512.0
            ],
            [
                0.0,
                0.0,
                1.0
            ]]).to(device)[None, ...].repeat(bs, 1, 1)
    ratio_h = rm.image_height / 1024.
    ratio_w = rm.image_width / 1024.
    batch['K'][:, 0:1, :] *= ratio_h
    batch['K'][:, 1:2, :] *= ratio_w

    final_preds = {}
    final_preds['rgb'] = []
    final_preds['rgb_boxes'] = []
    final_preds['mat_rgb'] = []
    for view_ang in view_angles:
        bs_view_ang = view_ang.repeat(bs,)
        batch['Rt'] = get_pose_on_orbit(radius=5*rm.volradius, height=0, angles=bs_view_ang).to(prim_volume)
        # raymarcher is in mm
        rm_preds = rm(
            prim_rgba=preds["prim_rgba"],
            prim_pos=preds["prim_pos"],
            prim_scale=preds["prim_scale"],
            prim_rot=preds["prim_rot"],
            RT=batch["Rt"],
            K=batch["K"],
        )
        rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)
        preds.update(alpha=rgba[..., -1].contiguous(), rgb=rgba[..., :3].contiguous())
        with torch.no_grad():
            preds["rgb_boxes"] = render_mvp_boxes(rm, batch, preds)
        rm_preds = rm(
            prim_rgba=preds["prim_mata"],
            prim_pos=preds["prim_pos"],
            prim_scale=preds["prim_scale"],
            prim_rot=preds["prim_rot"],
            RT=batch["Rt"],
            K=batch["K"],
        )
        mat_rgba = rm_preds["rgba_image"].permute(0, 2, 3, 1)
        preds.update(mat_rgb=mat_rgba[..., :3].contiguous())
        final_preds['rgb'].append(preds['rgb'])
        final_preds['rgb_boxes'].append(preds['rgb_boxes'])
        final_preds['mat_rgb'].append(preds['mat_rgb'])

    assert len(final_preds['rgb']) == len(final_preds['rgb_boxes'])
    final_preds['rgb'] = torch.concat(final_preds['rgb'], dim=0)
    final_preds['rgb_boxes'] = torch.concat(final_preds['rgb_boxes'], dim=0)
    final_preds['mat_rgb'] = torch.concat(final_preds['mat_rgb'], dim=0)
    total_num_frames = final_preds['rgb'].shape[0]
    rgb_video = os.path.join(video_save_folder, 'rgb.mp4')
    rgb_video_out = imageio.get_writer(rgb_video, fps=20)
    prim_video = os.path.join(video_save_folder, 'prim.mp4')
    prim_video_out = imageio.get_writer(prim_video, fps=20)
    mat_video = os.path.join(video_save_folder, 'mat.mp4')
    mat_video_out = imageio.get_writer(mat_video, fps=20)

    rgb_np = np.clip(final_preds['rgb'].detach().cpu().numpy(), 0, 255).astype(np.uint8)
    prim_np = np.clip(final_preds['rgb_boxes'].detach().cpu().numpy(), 0, 255).astype(np.uint8)
    mat_np = np.clip(final_preds['mat_rgb'].detach().cpu().numpy(), 0, 255).astype(np.uint8)
    for fidx in range(total_num_frames):
        rgb_video_out.append_data(rgb_np[fidx])
        prim_video_out.append_data(prim_np[fidx])
        mat_video_out.append_data(mat_np[fidx])
    rgb_video_out.close()
    prim_video_out.close()
    mat_video_out.close()