189 lines
6.0 KiB
Python
189 lines
6.0 KiB
Python
import os
|
|
import imageio
|
|
import glob
|
|
import json
|
|
import subprocess
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from util.tensor import move_to, detach_all, to_torch
|
|
from util.logger import Logger
|
|
|
|
|
|
def get_results_paths(res_dir):
|
|
"""
|
|
get the iterations of all saved results in res_dir
|
|
:param res_dir (str) result dir
|
|
returns a dict of iter to result path
|
|
"""
|
|
res_files = sorted(glob.glob(f"{res_dir}/*_results.npz"))
|
|
print(f"found {len(res_files)} results in {res_dir}")
|
|
|
|
path_dict = {}
|
|
for res_file in res_files:
|
|
it, name, _ = os.path.basename(res_file).split("_")[-3:]
|
|
assert name in ["world", "prior"]
|
|
if it not in path_dict:
|
|
path_dict[it] = {}
|
|
path_dict[it][name] = res_file
|
|
return path_dict
|
|
|
|
|
|
def load_result(res_path_dict):
|
|
"""
|
|
load all saved results for a given iteration
|
|
:param res_path_dict (dict) paths to relevant results
|
|
returns dict of results
|
|
"""
|
|
res_dict = {}
|
|
for name, path in res_path_dict.items():
|
|
res = np.load(path)
|
|
res_dict[name] = to_torch({k: res[k] for k in res.files})
|
|
return res_dict
|
|
|
|
|
|
def save_initial_predictions(model, out_dir, seq_name):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
with torch.no_grad():
|
|
print('go into get_optim_result from output.py save_initial_predictions')
|
|
pred_dict = model.get_optim_result()
|
|
pred_dict = move_to(detach_all(pred_dict), "cpu")
|
|
|
|
for name, results in pred_dict.items():
|
|
out_path = f"{out_dir}/{seq_name}_{0:06d}_init_{name}_results.npz"
|
|
Logger.log(f"saving initial results to {out_path}")
|
|
np.savez(out_path, **results)
|
|
# print('cam_R in save_initial_predictions: ', results['cam_R'].shape, results['cam_R'])
|
|
# print((results['cam_R'][0] == results['cam_R'][1]).all())
|
|
# print((results['cam_R'][1] == results['cam_R'][2]).all())
|
|
# print((results['cam_R'][2] == results['cam_R'][3]).all())
|
|
# print((results['cam_R'][3] == results['cam_R'][4]).all())
|
|
# print((results['cam_R'][4] == results['cam_R'][5]).all())
|
|
# print((results['cam_R'][5] == results['cam_R'][6]).all())
|
|
|
|
|
|
def save_input_poses(dataset, out_dir, seq_name, name="phalp"):
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
dataset.load_data(interp_input=False)
|
|
d = dataset.data_dict
|
|
res = {
|
|
"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),
|
|
}
|
|
print({k: v.shape for k, v in res.items()})
|
|
out_path = f"{out_dir}/{seq_name}_{0:06d}_{name}_world_results.npz"
|
|
Logger.log(f"saving inputs to {out_path}")
|
|
np.savez(out_path, **res)
|
|
|
|
|
|
def save_input_frames_ffmpeg(dataset, out_dir, name="input", fps=30, overwrite=False):
|
|
vid_path = f"{out_dir}/{name}.mp4"
|
|
if not overwrite and os.path.isfile(vid_path):
|
|
return
|
|
|
|
list_path = f"{out_dir}/src_paths.txt"
|
|
with open(list_path, "w") as f:
|
|
f.write("\n".join([f"file '{p}'" for p in dataset.sel_img_paths]))
|
|
|
|
filter_str = f"-framerate {fps} -vf 'format=yuv420p'"
|
|
cmd = (
|
|
f"ffmpeg -loglevel error -f concat -safe 0 "
|
|
f"-i {list_path} -c:v libx264 {filter_str} {vid_path} -y"
|
|
)
|
|
print(cmd)
|
|
subprocess.call(cmd, shell=True, stdin=subprocess.PIPE)
|
|
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
|
return vid_path
|
|
|
|
|
|
def save_input_frames(dataset, vid_path, fps=30, overwrite=False):
|
|
if not overwrite and os.path.isfile(vid_path):
|
|
return
|
|
|
|
writer = imageio.get_writer(vid_path, fps=fps)
|
|
for path in dataset.sel_img_paths:
|
|
writer.append_data(imageio.imread(path))
|
|
writer.close()
|
|
print(f"SAVED {len(dataset.sel_img_paths)} INPUT FRAMES TO {vid_path}")
|
|
return vid_path
|
|
|
|
|
|
def load_track_info(path):
|
|
assert os.path.isfile(path)
|
|
|
|
with open(path, "r") as f:
|
|
track_info = json.load(f)
|
|
|
|
meta = track_info["meta"]
|
|
tracks = track_info["tracks"]
|
|
track_ids, idcs, vis_masks = map(
|
|
torch.as_tensor,
|
|
zip(
|
|
*[
|
|
(int(tid), info["index"], info["vis_mask"])
|
|
for tid, info in tracks.items()
|
|
]
|
|
),
|
|
)
|
|
track_ids = track_ids[idcs]
|
|
vis_masks = vis_masks[idcs]
|
|
data_interval = meta["data_interval"]
|
|
seq_interval = meta["seq_interval"]
|
|
return track_ids, vis_masks, data_interval, seq_interval
|
|
|
|
|
|
def save_track_info(dataset, out_dir):
|
|
# track indices are relative to the entire sequence
|
|
track_ids = dataset.track_ids
|
|
track_dict = {}
|
|
for i, (tid, mask) in enumerate(zip(dataset.track_ids, dataset.track_vis_masks)):
|
|
track_dict[int(tid)] = {"index": i, "vis_mask": mask.tolist()}
|
|
|
|
out_dict = {
|
|
"tracks": track_dict,
|
|
"meta": {
|
|
"seq_interval": (int(dataset.start_idx), int(dataset.end_idx)),
|
|
"data_interval": (int(dataset.data_start), int(dataset.data_end)),
|
|
},
|
|
}
|
|
|
|
with open(f"{out_dir}/track_info.json", "w") as f:
|
|
json.dump(out_dict, f)
|
|
print("SAVED TRACK INFO")
|
|
|
|
|
|
def load_camera_json(path):
|
|
assert os.path.isfile(path)
|
|
|
|
with open(path, "r") as f:
|
|
cam_data = json.load(f)
|
|
R = torch.as_tensor(cam_data["rotation"]).reshape(-1, 3, 3)
|
|
t = torch.as_tensor(cam_data["translation"])
|
|
intrins = torch.as_tensor(cam_data["intrinsics"])
|
|
return R.float(), t.float(), intrins.float()
|
|
|
|
|
|
def save_camera_json(path, cam_R, cam_t, intrins):
|
|
"""
|
|
:param path
|
|
:param cam_R (N, 3, 3)
|
|
:param cam_t (N, 3)
|
|
:param intrins (N, 4)
|
|
"""
|
|
N = len(cam_R)
|
|
T = torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype=torch.float32)
|
|
cam_R = torch.einsum("ij,bjk->bik", T, cam_R)
|
|
cam_t = torch.einsum("ij,bj->bi", T, cam_t)
|
|
with open(path, "w") as f:
|
|
json.dump(
|
|
{
|
|
"rotation": cam_R.reshape(N, 9).tolist(),
|
|
"translation": cam_t.tolist(),
|
|
"intrinsics": intrins.tolist(),
|
|
},
|
|
f,
|
|
indent=1,
|
|
)
|