visualizer/run_vis.py

#!/usr/bin/env python3
# Copyright © Niantic, Inc. 2023.

import logging
import shutil
from pathlib import Path
import os
import numpy as np
import argparse
from distutils.util import strtobool
import time
import ace_zero_util as zutil
from ace_util import load_npz_file, compute_knn_mask
from joblib import Parallel, delayed
# from ace_trainer import TrainerACE
# import dataset_io
from PIL import Image
from ace_visualizer import ACEVisualizer
import subprocess
import matplotlib
_logger = logging.getLogger(__name__)
import imageio



def colorize(depth: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray:
    if mask is None:
        depth = np.where(depth > 0, depth, np.nan)
    else:
        depth = np.where((depth > 0) & mask, depth, np.nan)
    disp = 1 / depth
    if normalize:
        min_disp, max_disp = np.nanquantile(disp, 0.001), np.nanquantile(disp, 0.999)
        disp = (disp - min_disp) / (max_disp - min_disp)
    colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disp), 0)
    colored = (colored.clip(0, 1) * 255).astype(np.uint8)[:, :, :3]
    return colored

def _strtobool(x):
    return bool(strtobool(x))

def npz2image(npz_file, save_path):
    data = np.load(npz_file)
    os.makedirs(save_path, exist_ok=True)
    image_data = data['images_gt'] if 'images_gt' in data else data['images']
    if image_data.shape[1] == 3:
        image_data = np.transpose(image_data, (0, 2, 3, 1))
    for idx in range(image_data.shape[0]):
        img = image_data[idx, :, :]
        if img.max() <= 1.0 and img.min() >= 0.0:
            img = (img * 255).astype(np.uint8)
        elif img.max() <= 1.0 and img.min() < 0.0:
            img = (img + 1.0) * 255 / 2.0
            img = img.astype(np.uint8)
        else:
            img = img.astype(np.uint8)
        # import pdb; pdb.set_trace()
        image = Image.fromarray(img)
        image.save(f"{save_path}/gt_image_{idx}.png")
        # image.save(os.path.join(save_path, str(idx) + '.png'))


class Options:
    def __init__(self, ** kwargs):
        self.__dict__.update(kwargs)


if __name__ == '__main__':

    # Setup logging levels.
    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(
        description='Run ACE0 for a dataset or a scene.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--export_point_cloud', type=_strtobool, default=True,
                        help="Export the ACE0 point cloud after reconstruction, "
                             "for visualisation or to initialise splats")

    parser.add_argument('--dense_point_cloud', type=_strtobool, default=True,
                        help='when exporting a point cloud, do not filter points based on reprojection error, '
                             'bad for visualisation but good to initialise splats')

    parser.add_argument('--num_data_workers', type=int, default=12,
                        help='number of data loading workers, set according to the number of available CPU cores')

    #Registration parameters 
    parser.add_argument('--ransac_iterations', type=int, default=32,
                        help="Number of RANSAC hypothesis when registering mapping frames.")
    parser.add_argument('--ransac_threshold', type=float, default=10,
                        help='RANSAC inlier threshold in pixels')

    # Visualisation parameters
    parser.add_argument('--render_visualization', type=_strtobool, default=True,
                        help="Render visualisation frames of the whole reconstruction process.")
    parser.add_argument('--render_flipped_portrait', type=_strtobool, default=False,
                        help="Dataset images are 90deg flipped (like Wayspots).")
    parser.add_argument('--render_marker_size', type=float, default=0.03,
                        help="Size of the camera marker when rendering scenes.")
    # parser.add_argument('--iterations_output', type=int, default=500,
    #                     help='how often to print the loss and render a frame')

    parser.add_argument('--result_npz', type=str, help='Path to the .npz file containing the raw data (e.g., intrinsic matrices, 3D points, etc.) for visualization')
    parser.add_argument('--results_folder', type=Path, default="1217_new", help='path to output folder for result files')
    parser.add_argument('--knn_percent', type=float, default=98.0, help='percentage of points to keep in the knn mask')
    parser.add_argument('--K_neighbors', type=int, default=10, help='number of neighbors to use for knn mask')
    parser.add_argument('--pan_radius_scale', type=float, default=-1, help='factor to control the size of the pan camera')
    parser.add_argument('--pan_start_angle', type=int, default=-90, help='start angle of the pan camera')
    opt = parser.parse_args()
    opt.results_folder = '/'.join(opt.result_npz.split('/')[:-1])
    os.makedirs(opt.results_folder, exist_ok=True)
    opt.results_folder = Path(opt.results_folder)
    image_folder = f'{opt.results_folder}/gt_images'
    npz2image(opt.result_npz, image_folder)
    opt.rgb_files = f'{image_folder}/*.png'

    input_data = load_npz_file(opt.result_npz)
    
    if 'pts_mask' not in input_data.keys():
        start_time = time.time()
        input_data['pts_mask'] = compute_knn_mask(input_data['pts3d'], opt.K_neighbors, opt.knn_percent)
        end_time = time.time()
        _logger.info(f"Computing knn mask took {end_time - start_time:.2f} seconds.")
        npz_with_mask = f"{opt.results_folder}/points_with_mask.npz"
        np.savez_compressed(npz_with_mask, **input_data)
        opt.result_npz = npz_with_mask
    best_seed = 4
    iteration_id = zutil.get_seed_id(best_seed)
    # depth = input_data['pts3d'][0,...,-1]
    poses = input_data['cam_poses']
    H, W = input_data['pts3d'].shape[1:3]
    for idx, (pts3d, pose) in enumerate(zip(input_data['pts3d'], poses)):
        w2c = np.linalg.inv(pose)
        R = w2c[:3,:3]
        T = w2c[:3,-1] 
        pts3d_c2w = np.einsum('kl, Nl -> Nk', R, pts3d.reshape(-1,3)) + T[None]
        disp = pts3d_c2w[...,-1]
        disp = disp.reshape(H, W)
        disp_vis = colorize(disp)
        os.makedirs(f"{opt.results_folder}/depth", exist_ok=True)
        imageio.imwrite(f"{opt.results_folder}/depth/depth_vis_{idx}.png", disp_vis)
        print(f"Visualized depth map is saved to {opt.results_folder}/depth/depth_vis_{idx}.png.")
    register_opt = zutil.get_register_opt(rgb_files=opt.rgb_files, result_npz=opt.result_npz)
    modify_options = {
        'rgb_files': opt.rgb_files,
        'render_visualization': opt.render_visualization,
        'render_target_path': zutil.get_render_path(opt.results_folder),
        'render_marker_size': opt.render_marker_size,
        'render_flipped_portrait': opt.render_flipped_portrait,
        'session': f"{iteration_id}",
        'hypotheses': opt.ransac_iterations,
        'threshold': opt.ransac_threshold,
        'hypotheses_max_tries': 16,
        'result_npz': opt.result_npz,
        'only_frustum': True,
        'pan_radius_scale': opt.pan_radius_scale,
        'pan_start_angle': opt.pan_start_angle
    }
    for k, v in modify_options.items():
        setattr(register_opt, k, v)
    
    reg_state_dict_1 = zutil.regitser_visulization(register_opt)

    scheduled_to_stop_early = False
    prev_iteration_id = iteration_id

    register_opt_iter = zutil.get_register_opt(rgb_files=opt.rgb_files, result_npz=opt.result_npz)
    modify_options = {
        'rgb_files': opt.rgb_files,
        'render_visualization': opt.render_visualization,
        'render_target_path': zutil.get_render_path(opt.results_folder),
        'render_marker_size': opt.render_marker_size,
        'render_flipped_portrait': opt.render_flipped_portrait,
        'hypotheses': opt.ransac_iterations,
        'threshold': opt.ransac_threshold,
        'hypotheses_max_tries': 16,
        'result_npz': opt.result_npz,
        'state_dict': reg_state_dict_1,
        'only_frustum': False,
        'pan_radius_scale': opt.pan_radius_scale,
        'pan_start_angle': opt.pan_start_angle
    }
    for k, v in modify_options.items():
        setattr(register_opt_iter, k, v)

    reg_state_dict_2 = zutil.regitser_visulization(register_opt_iter)

    _logger.info("Rendering final sweep.")

    final_sweep_visualizer = ACEVisualizer(
        zutil.get_render_path(opt.results_folder), 
        flipped_portait=False, 
        map_depth_filter=100,
        marker_size=opt.render_marker_size,
        result_npz=opt.result_npz, 
        pan_start_angle=opt.pan_start_angle, 
        pan_radius_scale=opt.pan_radius_scale
    )
    poses = [final_sweep_visualizer.pose_align(pose) for pose in final_sweep_visualizer.cam_pose]
    rgb_path = f'{opt.results_folder}/gt_images'
    rgb_files = os.listdir(rgb_path)
    pose_dict = {str(rgb_file): 0 for rgb_file in rgb_files}
    pose_iterations = [0 for _ in range(len(rgb_files))]
    # import pdb; pdb.set_trace()
    final_sweep_visualizer.render_final_sweep(
        frame_count=150,
        camera_z_offset=4,
        poses=poses,
        pose_iterations=pose_iterations,
        total_poses=len(pose_dict),
        state_dict=reg_state_dict_2,)

    _logger.info("Converting to video.")

    # get ffmpeg path
    ffmpeg_path = shutil.which("ffmpeg")

    # run ffmpeg to convert the rendered images to a video
    ffmpeg_save_cmd = [ffmpeg_path,
                    "-y",
                    "-framerate", "30",
                    "-pattern_type", "glob",
                    "-i", f"{zutil.get_render_path(opt.results_folder)}/*.png",
                    "-c:v", "libx264",
                    "-pix_fmt", "yuv420p",
                    str(opt.results_folder / "reconstruction.mp4")
                    ]
    
    subprocess.run(ffmpeg_save_cmd, check=True)
    
    print(f'The render result video is saved to {opt.results_folder}/reconstruction.mp4.')
    if opt.export_point_cloud:
        # import pdb; pdb.set_trace()
        import trimesh
        pts = final_sweep_visualizer.pts3d.reshape(-1, 3) 
        image = final_sweep_visualizer.image_gt.transpose(0, 2, 3, 1).reshape(-1, 3)
        clr = (( image + 1.0 ) / 2.0 * 255.0).astype('float64')
        cloud = trimesh.PointCloud(pts, colors=clr)
        cloud.export(f"{opt.results_folder}/point_cloud.ply")
        print(f'The point cloud is saved to {opt.results_folder}/point_cloud.ply.')