import os import glob import imageio import numpy as np import torch from torch.utils.data import DataLoader from body_model import MANO from omegaconf import DictConfig, OmegaConf from data import get_dataset_from_cfg, expand_source_paths from optim.output import ( get_results_paths, load_result, save_input_frames, save_input_poses, ) from util.loaders import load_config_from_log, resolve_cfg_paths from util.tensor import get_device, move_to, detach_all, to_torch from vis.output import prep_result_vis, animate_scene, make_video_grid_2x2 from vis.tools import vis_keypoints from vis.viewer import init_viewer import time from geometry.mesh import save_mesh_scenes, vertices_to_trimesh LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353) def save_meshes_all(cfg, dataset, res_dicts, dev_id, mesh_dirs, num_steps=-1): B = len(dataset) T = dataset.seq_len loader = DataLoader(dataset, batch_size=B, shuffle=False) device = get_device(dev_id) obs_data = move_to(next(iter(loader)), device) # load models cfg = resolve_cfg_paths(cfg) # Instantiate MANO model mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()} print('initializing MANO model with cfgs:', mano_cfg, 'B*T', B*T) hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device) for res_dict, mesh_dir in zip(res_dicts, mesh_dirs): res_dict = move_to(res_dict, device) scene_dict = move_to( prep_result_vis( res_dict, obs_data["vis_mask"], obs_data["track_id"], hand_model, temporal_smooth=cfg.temporal_smooth, smooth_trans=True # For mesh export, smooth everything ), "cpu", ) scene_dir = mesh_dir verts, joints, colors, l_faces, r_faces, is_right, bounds = scene_dict["geometry"] T = len(verts) print(f"{T} mesh frames for ", mesh_dir) times = list(range(0, T, 1)) flag = False for t in times: if len(is_right[t]) > 1: flag = True vv = t if flag: init_trans = (joints[vv][0][9].clone() + joints[vv][1][9].clone()) / 2 else: init_trans = joints[0][0][9].clone() for t in times: if len(is_right[t]) > 1: assert (is_right[t].cpu().numpy().tolist() == [0,1]) # l_meshes = make_batch_mesh(verts[t][0][None], l_faces[t], colors[t][0][None]) # r_meshes = make_batch_mesh(verts[t][1][None], r_faces[t], colors[t][1][None]) # assert len(l_meshes) == 1 # assert len(r_meshes) == 1 verts[t][0] -= init_trans joints[t][0] -= init_trans tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0) tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj')) # np.save(f'{str(t).zfill(6)}_0.npy', joints[t][0].detach().cpu().numpy()) verts[t][1] -= init_trans joints[t][1] -= init_trans tmesh = vertices_to_trimesh(verts[t][1].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1) tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj')) # np.save(f'{str(t).zfill(6)}_1.npy', joints[t][1].detach().cpu().numpy()) else: assert len(is_right[t]) == 1 if is_right[t] == 0: verts[t][0] -= init_trans joints[t][0] -= init_trans tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), l_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=0) tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_0.obj')) # np.save(f'{str(t).zfill(6)}_0.npy', joints[t][0].detach().cpu().numpy()) elif is_right[t] == 1: verts[t][0] -= init_trans joints[t][0] -= init_trans tmesh = vertices_to_trimesh(verts[t][0].detach().cpu().numpy(), r_faces[t].detach().cpu().numpy(), LIGHT_BLUE, is_right=1) tmesh.export(os.path.join(scene_dir, f'{str(t).zfill(6)}_1.obj')) # np.save(f'{str(t).zfill(6)}_1.npy', joints[t][0].detach().cpu().numpy()) def run_vis( cfg, dataset, out_dir, dev_id, phases=["smooth_fit"], render_views=["src_cam", "above", "side"], make_grid=True, overwrite=False, save_dir=None, render_kps=False, render_layers=False, save_frames=False, **kwargs ): save_dir = out_dir if save_dir is None else save_dir print("OUT_DIR", out_dir) print("SAVE_DIR", save_dir) print("VISUALIZING PHASES", phases) print("RENDERING VIEWS", render_views) print("RENDER_KPS", render_kps) print("OVERWRITE", overwrite) # save input frames inp_vid_path = save_input_frames( dataset, f"{save_dir}/{dataset.seq_name}_input.mp4", fps=cfg.fps, overwrite=True, ) if render_kps: render_keypoints_2d(dataset, save_dir, overwrite=overwrite) if len(render_views) < 1: return out_ext = "/" if render_layers or save_frames else ".mp4" phase_results = {} phase_max_iters = {} for phase in phases: res_dir = os.path.join(out_dir, phase) if phase == "input": res = get_input_dict(dataset) it = f"{0:06d}" elif os.path.isdir(res_dir): res_path_dict = get_results_paths(res_dir) print(f"FOUND {len(res_path_dict)} results in {res_dir}") it = sorted(res_path_dict.keys())[-1] res = load_result(res_path_dict[it])["world"] else: print(f"{res_dir} does not exist, skipping") continue out_name = f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}" mesh_dir = f"{save_dir}/{phase}/{dataset.seq_name}_{it}_meshes" os.makedirs(mesh_dir, exist_ok=True) phase_max_iters[phase] = it out_paths = [f"{out_name}_{view}{out_ext}" for view in render_views] if not overwrite and all(os.path.exists(p) for p in out_paths): print("FOUND OUT PATHS", out_paths) continue phase_results[phase] = out_name, mesh_dir, res if len(phase_results) > 0: out_names, mesh_dir, res_dicts = zip(*phase_results.values()) render_results( cfg, dataset, dev_id, res_dicts, out_names, render_views=render_views, render_layers=render_layers, save_frames=save_frames, **kwargs, ) save_meshes_all(cfg, dataset, res_dicts, dev_id, mesh_dir) if make_grid: for phase, it in phase_max_iters.items(): grid_path = f"{save_dir}/{dataset.seq_name}_{phase}_grid.mp4" vid_paths = [ f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_src_cam.mp4", f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_front.mp4", f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_above.mp4", f"{save_dir}/{dataset.seq_name}_{phase}_final_{it}_side.mp4", ] make_video_grid_2x2( grid_path, vid_paths, fps=cfg.fps, overwrite=True, ) def get_input_dict(dataset): dataset.load_data(interp_input=False) d = dataset.data_dict input_params = { "pose_body": np.stack(d["init_body_pose"], axis=0), "trans": np.stack(d["init_trans"], axis=0), "root_orient": np.stack(d["init_root_orient"], axis=0), } input_params = to_torch(input_params) print({k: v.shape for k, v in input_params.items()}) return input_params def render_keypoints_2d(dataset, save_dir, overwrite=False): """ render 2d keypoints for each track """ dataset.load_data() out_dir = f"{save_dir}/{dataset.seq_name}_joints2d" B, T = dataset.n_tracks, dataset.seq_len if not overwrite and (os.path.isdir(out_dir) and len(os.listdir(out_dir)) >= B * T): print(f"Keypoints already rendered in {out_dir}") return os.makedirs(out_dir, exist_ok=True) for i, tid in enumerate(dataset.track_ids): joints2d = dataset.data_dict["joints2d"][i] # (T, J, 3) for t, sel_img_name in enumerate(dataset.sel_img_names): img = vis_keypoints(joints2d[t : t + 1], dataset.img_size) out_path = f"{out_dir}/{sel_img_name}_{tid}.png" imageio.imwrite(out_path, img) def render_results(cfg, dataset, dev_id, res_dicts, out_names, **kwargs): """ render results for all selected phases """ assert len(res_dicts) == len(out_names) if len(res_dicts) < 1: print("no results to render, skipping") return B = len(dataset) T = dataset.seq_len loader = DataLoader(dataset, batch_size=B, shuffle=False) device = get_device(dev_id) obs_data = move_to(next(iter(loader)), device) cam_data = dataset.get_camera_data() # load models cfg = resolve_cfg_paths(cfg) # Instantiate MANO model mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()} print('initializing MANO model with cfgs:', mano_cfg, 'B*T', B*T) hand_model = MANO(batch_size=B*T, pose2rot=True, **mano_cfg).to(device) vis = init_viewer( dataset.img_size, cam_data["intrins"][0], vis_scale=1.0, bg_paths=dataset.sel_img_paths, fps=cfg.fps, ) # Set 2D keypoints for overlay visualization if render_keypoints is enabled if kwargs.get('render_keypoints', False): # Load 2D keypoints from dataset - assume single track for now # joints2d is list of (T, J, 3) arrays, one per track dataset.load_data() if len(dataset.data_dict["joints2d"]) > 0: joints2d = dataset.data_dict["joints2d"][0] # First track: (T, J, 3) keypoints_seq = [joints2d[t] for t in range(T)] # List of (J, 3) arrays vis.set_keypoints_seq(keypoints_seq) print(f"Loaded {T} frames of 2D keypoints with {joints2d.shape[1]} joints") save_paths_all = [] render_views = kwargs.get('render_views', ['src_cam', 'above', 'side']) # Separate src_cam from other views src_cam_views = [v for v in render_views if v == 'src_cam'] other_views = [v for v in render_views if v != 'src_cam'] for res_dict, out_name in zip(res_dicts, out_names): print(f'preparing results for rendering {out_name}') res_dict = move_to(res_dict, device) # Render src_cam WITH temporal smoothing but WITHOUT trans smoothing if src_cam_views: print(f"Rendering src_cam view WITH temporal smoothing (excluding trans)") scene_dict_no_smooth = prep_result_vis( res_dict, obs_data["vis_mask"], obs_data["track_id"], hand_model, temporal_smooth=cfg.temporal_smooth, # Apply smoothing smooth_trans=False # But don't smooth translation ) # Compute predicted 2D keypoints using THE SAME function as optimization if kwargs.get('render_keypoints', False): from geometry import camera as cam_util from body_model import run_mano # Get the same data as optimization joints3d_op = run_mano( hand_model, res_dict["trans"], res_dict["root_orient"], res_dict["pose_body"], res_dict["is_right"], res_dict.get("betas", None), )["joints"] # (B, T, J, 3) # Get cameras (same as optimization) cam_R = res_dict["cam_R"] # (B, T, 3, 3) cam_t = res_dict["cam_t"] # (B, T, 3) intrins = res_dict["intrins"] # (4,) or (T, 4) or (B, T, 4) # Ensure intrins is (B, T, 4) if intrins.ndim == 1: # (4,) intrins = intrins[None, None].expand(cam_R.shape[0], cam_R.shape[1], -1) # (B, T, 4) elif intrins.ndim == 2: # (T, 4) intrins = intrins[None].expand(cam_R.shape[0], -1, -1) # (B, T, 4) cam_f = intrins[:, :, :2] # (B, T, 2) cam_center = intrins[:, :, 2:] # (B, T, 2) # Reproject using EXACTLY the same function as optimization joints2d_pred = cam_util.reproject(joints3d_op, cam_R, cam_t, cam_f[0], cam_center[0]) # (B, T, J, 2) # Convert to list for first track pred_keypoints_seq = [joints2d_pred[0, t].cpu().numpy() for t in range(joints2d_pred.shape[1])] vis.set_pred_keypoints_seq(pred_keypoints_seq) print(f"Computed {len(pred_keypoints_seq)} frames of predicted 2D keypoints using reproject()") # Render only src_cam view kwargs_src = kwargs.copy() kwargs_src['render_views'] = src_cam_views save_paths_src = animate_scene( vis, scene_dict_no_smooth, out_name, seq_name=dataset.seq_name, **kwargs_src ) save_paths_all.append(save_paths_src) # Render other views WITH temporal smoothing (including trans) if other_views: print(f"Rendering {other_views} views WITH temporal smoothing (including trans)") scene_dict_smooth = prep_result_vis( res_dict, obs_data["vis_mask"], obs_data["track_id"], hand_model, temporal_smooth=cfg.temporal_smooth, # Apply smoothing for other views smooth_trans=True # Smooth translation for other views ) # Render other views kwargs_other = kwargs.copy() kwargs_other['render_views'] = other_views kwargs_other['render_keypoints'] = False # No keypoints on other views save_paths_other = animate_scene( vis, scene_dict_smooth, out_name, seq_name=dataset.seq_name, **kwargs_other ) if not src_cam_views: save_paths_all.append(save_paths_other) vis.close() return save_paths_all def visualize_log(log_dir, dev_id, phases, save_dir=None, **kwargs): print(log_dir) cfg = load_config_from_log(log_dir) print(cfg) print(cfg.data) print(cfg.data.sources) # make sure we get all necessary inputs cfg.data.sources = expand_source_paths(cfg.data.sources) print("SOURCES", cfg.data.sources) # cfg.data.track_ids = "001" dataset = get_dataset_from_cfg(cfg) if len(dataset) < 1: print(f"No tracks in dataset, skipping") return run_vis(cfg, dataset, log_dir, dev_id, phases=phases, save_dir=save_dir, **kwargs) def launch_vis(i, args): log_dir = args.log_dirs[i] dev_id = args.gpus[i % len(args.gpus)] os.environ["EGL_DEVICE_ID"] = str(dev_id) os.environ["PYOPENGL_PLATFORM"] = "egl" if args.save_root is not None: path_name = log_dir.split(args.log_root)[-1].strip("/") exp_name = "-".join(path_name.split("/")[:2]) save_dir = f"{args.save_root}/{exp_name}" os.makedirs(save_dir, exist_ok=True) print('args.save_root: ', args.save_root) visualize_log( log_dir, dev_id, phases=args.phases, save_dir=save_dir, overwrite=args.overwrite, accumulate=args.accumulate, render_kps=args.render_kps, render_layers=args.render_layers, render_views=args.render_views, save_frames=args.save_frames, make_grid=args.grid, ) def main(args): """ visualize all runs in root """ OmegaConf.register_new_resolver("eval", eval) log_dirs = [] for root, subd, files in os.walk(args.log_root): if ".hydra" in subd: log_dirs.append(root) args.log_dirs = log_dirs print(f"FOUND {len(args.log_dirs)} TO RENDER") if len(args.gpus) > 1: from torch.multiprocessing import Pool torch.multiprocessing.set_start_method("spawn") with Pool(processes=len(args.gpus)) as pool: res = pool.starmap( launch_vis, [(i, args) for i in range(len(args.log_dirs))] ) return for i in range(len(args.log_dirs)): launch_vis(i, args) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--log_root", required=True) parser.add_argument("--save_root", default='./') parser.add_argument("--phases", nargs="*", default=["smooth_fit"]) parser.add_argument("--gpus", nargs="*", default=[0]) parser.add_argument( "-rv", "--render_views", nargs="*", default=["src_cam", "front", "above", "side"], ) parser.add_argument("-g", "--grid", action="store_true") parser.add_argument("-rl", "--render_layers", action="store_true") parser.add_argument("-kp", "--render_kps", action="store_true") parser.add_argument("-sf", "--save_frames", action="store_true") parser.add_argument("-ra", "--accumulate", action="store_true") parser.add_argument("-y", "--overwrite", action="store_true") args = parser.parse_args() main(args)