Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
from .dataset import *
|
||||
from . import tools
|
||||
+482
@@ -0,0 +1,482 @@
|
||||
import os
|
||||
import glob
|
||||
import typing
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import json
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from body_model import MANO_JOINTS
|
||||
assert len(MANO_JOINTS) == 16
|
||||
import time
|
||||
|
||||
from util.logger import Logger
|
||||
from geometry.camera import invert_camera
|
||||
|
||||
from .tools import read_keypoints, read_mask_path, load_mano_preds, load_keypoints_with_interp
|
||||
from .vidproc import preprocess_cameras, preprocess_frames, preprocess_tracks
|
||||
|
||||
|
||||
"""
|
||||
Define data-related constants
|
||||
"""
|
||||
# DEFAULT_GROUND = np.array([0.0, -1.0, 0.0, -0.5])
|
||||
|
||||
# XXX: TEMPORARY CONSTANTS
|
||||
SHOT_PAD = 0
|
||||
MIN_SEQ_LEN = 10
|
||||
MAX_NUM_TRACKS = 12
|
||||
MIN_TRACK_LEN = 60
|
||||
MIN_KEYP_CONF = 0.4
|
||||
|
||||
|
||||
def get_dataset_from_cfg(cfg):
|
||||
args = cfg.data
|
||||
if not args.use_cams:
|
||||
args.sources.cameras = ""
|
||||
|
||||
args.sources = expand_source_paths(args.sources)
|
||||
print("DATA SOURCES", args.sources)
|
||||
a = time.time()
|
||||
check_data_sources(args, cfg)
|
||||
print('end of check_data_sources, time: ', time.time() - a)
|
||||
|
||||
return MultiPeopleDataset(
|
||||
args.sources,
|
||||
args.seq,
|
||||
tid_spec=args.track_ids,
|
||||
shot_idx=args.shot_idx,
|
||||
start_idx=int(args.start_idx),
|
||||
end_idx=int(args.end_idx),
|
||||
is_static=cfg.is_static,
|
||||
split_cameras=args.get("split_cameras", True),
|
||||
)
|
||||
|
||||
|
||||
def expand_source_paths(data_sources):
|
||||
return {k: get_data_source(v) for k, v in data_sources.items()}
|
||||
|
||||
|
||||
def get_data_source(source):
|
||||
matches = glob.glob(source)
|
||||
if len(matches) < 1:
|
||||
print(f"{source} does not exist")
|
||||
return source # return anyway for default values
|
||||
if len(matches) > 1:
|
||||
raise ValueError(f"{source} is not unique")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def check_data_sources(args, cfg):
|
||||
print()
|
||||
a = time.time()
|
||||
if args.type == "video":
|
||||
preprocess_frames(args.sources.images, args.src_path, **args.frame_opts)
|
||||
b = time.time()
|
||||
preprocess_tracks(cfg.datatype, args.sources.images, args.sources.tracks, args.sources.shots, gpu=cfg.gpu)
|
||||
c = time.time()
|
||||
preprocess_cameras(args, overwrite=args.get("overwrite_cams", False))
|
||||
d = time.time()
|
||||
print('frame, hand track, slam camera: ', b-a, c-b, d-c)
|
||||
|
||||
|
||||
class MultiPeopleDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data_sources: typing.Dict,
|
||||
seq_name,
|
||||
tid_spec="all",
|
||||
shot_idx=0,
|
||||
start_idx=0,
|
||||
end_idx=-1,
|
||||
is_static=False,
|
||||
pad_shot=False,
|
||||
split_cameras=True,
|
||||
):
|
||||
self.is_static = is_static
|
||||
self.seq_name = seq_name
|
||||
self.data_sources = data_sources
|
||||
self.split_cameras = split_cameras
|
||||
|
||||
# select only images in the desired shot
|
||||
img_files, _ = get_shot_img_files(
|
||||
self.data_sources["shots"], shot_idx, pad_shot
|
||||
)
|
||||
end_idx = end_idx if end_idx > 0 else len(img_files)
|
||||
self.data_start, self.data_end = start_idx, end_idx
|
||||
img_files = img_files[start_idx:end_idx]
|
||||
self.img_names = [get_name(f) for f in img_files]
|
||||
self.num_imgs = len(self.img_names)
|
||||
|
||||
img_dir = self.data_sources["images"]
|
||||
assert os.path.isdir(img_dir)
|
||||
self.img_paths = [os.path.join(img_dir, f) for f in img_files]
|
||||
img_h, img_w = imageio.imread(self.img_paths[0]).shape[:2]
|
||||
self.img_size = img_w, img_h
|
||||
print(f"USING TOTAL {self.num_imgs} {img_w}x{img_h} IMGS")
|
||||
|
||||
# find the tracks in the video
|
||||
track_root = self.data_sources["tracks"]
|
||||
if tid_spec == "all" or tid_spec.startswith("longest"):
|
||||
n_tracks = MAX_NUM_TRACKS
|
||||
if tid_spec.startswith("longest"):
|
||||
n_tracks = int(tid_spec.split("-")[1])
|
||||
# get the longest tracks in the selected shot
|
||||
track_ids = sorted(os.listdir(track_root))
|
||||
track_paths = [
|
||||
[f"{track_root}/{tid}/{name}_keypoints.json" for name in self.img_names]
|
||||
for tid in track_ids
|
||||
]
|
||||
track_lens = [
|
||||
len(list(filter(os.path.isfile, paths))) for paths in track_paths
|
||||
]
|
||||
print("raw TRACK IDS and LENGTHS", track_ids, track_lens)
|
||||
track_ids = [
|
||||
track_ids[i]
|
||||
for i in range(len(track_lens))# np.argsort(track_lens)[::-1]
|
||||
if track_lens[i] > MIN_TRACK_LEN
|
||||
]
|
||||
print("TRACK LENGTHS", track_ids, track_lens)
|
||||
track_ids = track_ids[:n_tracks]
|
||||
else:
|
||||
track_ids = [f"{int(tid):03d}" for tid in tid_spec.split("-")]
|
||||
|
||||
print("TRACK IDS", track_ids)
|
||||
|
||||
self.track_ids = track_ids
|
||||
self.n_tracks = len(track_ids)
|
||||
self.track_dirs = [os.path.join(track_root, tid) for tid in track_ids]
|
||||
|
||||
# keep a list of frame index masks of whether a track is available in a frame
|
||||
sidx = np.inf
|
||||
eidx = -1
|
||||
self.track_vis_masks = []
|
||||
for pred_dir in self.track_dirs:
|
||||
kp_paths = [f"{pred_dir}/{x}_keypoints.json" for x in self.img_names]
|
||||
has_kp = [os.path.isfile(x) for x in kp_paths]
|
||||
|
||||
# keep track of which frames this track is visible in
|
||||
vis_mask = np.array(has_kp)
|
||||
idcs = np.where(vis_mask)[0]
|
||||
if len(idcs) > 0:
|
||||
si, ei = min(idcs), max(idcs)
|
||||
sidx = min(sidx, si)
|
||||
eidx = max(eidx, ei)
|
||||
self.track_vis_masks.append(vis_mask)
|
||||
|
||||
eidx = max(eidx + 1, 0)
|
||||
sidx = min(sidx, eidx)
|
||||
print("START", sidx, "END", eidx)
|
||||
self.start_idx = sidx
|
||||
self.end_idx = eidx
|
||||
self.seq_len = eidx - sidx
|
||||
self.seq_intervals = [(sidx, eidx) for _ in track_ids]
|
||||
|
||||
self.sel_img_paths = self.img_paths[sidx:eidx]
|
||||
self.sel_img_names = self.img_names[sidx:eidx]
|
||||
|
||||
# used to cache data
|
||||
self.data_dict = {}
|
||||
self.cam_data = None
|
||||
|
||||
def __len__(self):
|
||||
return self.n_tracks
|
||||
|
||||
def load_data(self, interp_input=True):
|
||||
if len(self.data_dict) > 0:
|
||||
return
|
||||
|
||||
# load camera data
|
||||
self.load_camera_data()
|
||||
# get data for each track
|
||||
data_out = {
|
||||
"mask_paths": [],
|
||||
"floor_plane": [],
|
||||
"joints2d": [],
|
||||
"vis_mask": [],
|
||||
"track_interval": [],
|
||||
"init_body_pose": [],
|
||||
"init_body_shape": [],
|
||||
"init_root_orient": [],
|
||||
"init_trans": [],
|
||||
"is_right": []
|
||||
}
|
||||
|
||||
# create batches of sequences
|
||||
# each batch is a track for a person
|
||||
T = self.seq_len
|
||||
sidx, eidx = self.start_idx, self.end_idx
|
||||
for i, tid in enumerate(self.track_ids):
|
||||
if len(self.track_ids) == 2:
|
||||
assert int(i) == int(tid), f'{i}/{tid}/{self.track_ids}'
|
||||
# load mask of visible frames for this track
|
||||
vis_mask = self.track_vis_masks[i][sidx:eidx] # (T)
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
track_s, track_e = min(vis_idcs), max(vis_idcs) + 1
|
||||
data_out["track_interval"].append([track_s, track_e])
|
||||
|
||||
vis_mask = get_ternary_mask(vis_mask)
|
||||
data_out["vis_mask"].append(vis_mask)
|
||||
|
||||
# load 2d keypoints for visible frames with interpolation
|
||||
kp_paths = [
|
||||
f"{self.track_dirs[i]}/{x}_keypoints.json" for x in self.sel_img_names
|
||||
]
|
||||
# (T, J, 3) (x, y, conf) - with interpolation for missing frames
|
||||
joints2d_data = load_keypoints_with_interp(kp_paths, interp=interp_input)
|
||||
# print('joints2d_data: ', joints2d_data.shape)
|
||||
# print('joints2d_data[:, :, [2]]: ', np.unique(joints2d_data[:, :, [2]]))
|
||||
# raise ValueError
|
||||
# Discard bad ViTPose detections
|
||||
joints2d_data[:, :, 2] = 1.0
|
||||
joints2d_data[
|
||||
np.repeat(joints2d_data[:, :, [2]] < MIN_KEYP_CONF, 3, axis=2)
|
||||
] = 0
|
||||
# Set all confidence values to 1.0 for optimization
|
||||
data_out["joints2d"].append(joints2d_data)
|
||||
|
||||
# load single image mano predictions
|
||||
pred_paths = [
|
||||
f"{self.track_dirs[i]}/{x}_mano.json" for x in self.sel_img_names
|
||||
]
|
||||
pose_init, orient_init, trans_init, betas_init, is_right = load_mano_preds(
|
||||
pred_paths, tid=tid, interp=interp_input
|
||||
)
|
||||
|
||||
n_joints = len(MANO_JOINTS) - 1 # 15
|
||||
data_out["init_body_pose"].append(pose_init)
|
||||
data_out["init_body_shape"].append(betas_init)
|
||||
data_out["init_root_orient"].append(orient_init)
|
||||
data_out["init_trans"].append(trans_init)
|
||||
data_out['is_right'].append(is_right)
|
||||
|
||||
# DEBUG: Print hand type for each track
|
||||
print(f"DEBUG: Track {i} (tid={tid}): is_right={is_right[0] if len(is_right) > 0 else 'empty'}, {len(pred_paths)} frames")
|
||||
|
||||
# data_out["floor_plane"].append(DEFAULT_GROUND[:3] * DEFAULT_GROUND[3:])
|
||||
|
||||
self.data_dict = data_out
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if len(self.data_dict) < 1:
|
||||
self.load_data()
|
||||
|
||||
obs_data = dict()
|
||||
|
||||
# 2D keypoints
|
||||
joint2d_data = self.data_dict["joints2d"][idx]
|
||||
obs_data["joints2d"] = torch.Tensor(joint2d_data)
|
||||
|
||||
# single frame predictions
|
||||
obs_data["init_body_pose"] = torch.Tensor(self.data_dict["init_body_pose"][idx])
|
||||
obs_data["init_body_shape"] = torch.Tensor(self.data_dict["init_body_shape"][idx])
|
||||
obs_data["init_root_orient"] = torch.Tensor(
|
||||
self.data_dict["init_root_orient"][idx]
|
||||
)
|
||||
obs_data["init_trans"] = torch.Tensor(self.data_dict["init_trans"][idx])
|
||||
obs_data["is_right"] = torch.Tensor(self.data_dict["is_right"][idx])
|
||||
|
||||
# floor plane
|
||||
# obs_data["floor_plane"] = torch.Tensor(self.data_dict["floor_plane"][idx])
|
||||
|
||||
# the frames the track is visible in
|
||||
obs_data["vis_mask"] = torch.Tensor(self.data_dict["vis_mask"][idx])
|
||||
|
||||
# the frames used in this subsequence
|
||||
obs_data["seq_interval"] = torch.Tensor(list(self.seq_intervals[idx])).to(
|
||||
torch.int
|
||||
)
|
||||
# the start and end interval of available keypoints
|
||||
obs_data["track_interval"] = torch.Tensor(
|
||||
self.data_dict["track_interval"][idx]
|
||||
).int()
|
||||
|
||||
print(len(self.track_ids), self.track_ids[idx], obs_data["is_right"].shape, obs_data["seq_interval"].shape)
|
||||
|
||||
obs_data["track_id"] = int(self.track_ids[idx])
|
||||
# print(self.track_ids[idx], obs_data["is_right"])
|
||||
# print(self.track_ids[idx] == obs_data["is_right"])
|
||||
# is_right should be constant across time for each track
|
||||
assert torch.all(obs_data["is_right"] == obs_data["is_right"][0]), "is_right should be constant"
|
||||
assert int(self.track_ids[idx]) == int(obs_data["is_right"][0]), f"Track ID {self.track_ids[idx]} != is_right {obs_data['is_right'][0]}"
|
||||
obs_data["seq_name"] = self.seq_name
|
||||
return obs_data
|
||||
|
||||
def load_camera_data(self):
|
||||
cam_dir = self.data_sources["cameras"]
|
||||
data_interval = 0, -1
|
||||
if self.split_cameras:
|
||||
data_interval = self.data_start, self.data_end
|
||||
track_interval = self.start_idx, self.end_idx
|
||||
self.cam_data = CameraData(
|
||||
cam_dir, self.seq_len, self.img_size, self.is_static, data_interval, track_interval
|
||||
)
|
||||
|
||||
def get_camera_data(self):
|
||||
if self.cam_data is None:
|
||||
raise ValueError
|
||||
return self.cam_data.as_dict()
|
||||
|
||||
|
||||
class CameraData(object):
|
||||
def __init__(
|
||||
self, cam_dir, seq_len, img_size, is_static, data_interval=[0, -1], track_interval=[0, -1]
|
||||
):
|
||||
self.img_size = img_size
|
||||
self.cam_dir = cam_dir
|
||||
self.is_static = is_static
|
||||
|
||||
# inclusive exclusive
|
||||
data_start, data_end = data_interval
|
||||
if data_end < 0:
|
||||
data_end += seq_len + 1
|
||||
data_len = data_end - data_start
|
||||
|
||||
# start and end indices are with respect to the data interval
|
||||
sidx, eidx = track_interval
|
||||
if eidx < 0:
|
||||
eidx += data_len + 1
|
||||
|
||||
self.sidx, self.eidx = sidx + data_start, eidx + data_start
|
||||
print(self.sidx, self.eidx, data_start, data_end, data_interval, track_interval)
|
||||
#
|
||||
self.seq_len = self.eidx - self.sidx
|
||||
|
||||
self.load_data()
|
||||
|
||||
def load_data(self):
|
||||
# camera info
|
||||
sidx, eidx = self.sidx, self.eidx
|
||||
img_w, img_h = self.img_size
|
||||
fpath = os.path.join(self.cam_dir, "cameras.npz")
|
||||
if os.path.isfile(fpath) and not self.is_static:
|
||||
Logger.log(f"Loading cameras from {fpath}...")
|
||||
cam_R, cam_t, intrins, width, height = load_cameras_npz(fpath, self.seq_len)
|
||||
scale = img_w / width
|
||||
self.intrins = scale * intrins[sidx:eidx]
|
||||
# move first camera to origin
|
||||
# R0, t0 = invert_camera(cam_R[sidx], cam_t[sidx])
|
||||
# self.cam_R = torch.einsum("ij,...jk->...ik", R0, cam_R[sidx:eidx])
|
||||
# self.cam_t = t0 + torch.einsum("ij,...j->...i", R0, cam_t[sidx:eidx])
|
||||
# t0 = -cam_t[sidx:eidx].mean(dim=0) + torch.randn(3) * 0.1
|
||||
t0 = -cam_t[sidx:sidx+1] + torch.randn(3) * 0.1
|
||||
self.cam_R = cam_R[sidx:eidx]
|
||||
self.cam_t = cam_t[sidx:eidx] - t0
|
||||
else:
|
||||
# raise ValueError
|
||||
Logger.log(f"WARNING: {fpath} does not exist, using static cameras...")
|
||||
default_focal = 0.5 * (img_h + img_w)
|
||||
self.intrins = torch.tensor(
|
||||
[default_focal, default_focal, img_w / 2, img_h / 2]
|
||||
)[None].repeat(self.seq_len, 1)
|
||||
|
||||
self.cam_R = torch.eye(3)[None].repeat(self.seq_len, 1, 1)
|
||||
self.cam_t = torch.zeros(self.seq_len, 3)
|
||||
|
||||
Logger.log(f"Images have {img_w}x{img_h}, intrins {self.intrins[0]}")
|
||||
print("CAMERA DATA", self.cam_R.shape, self.cam_t.shape, self.intrins[0])
|
||||
|
||||
def world2cam(self):
|
||||
return self.cam_R, self.cam_t
|
||||
|
||||
def cam2world(self):
|
||||
print('dataset.py, cam2world: ', self.cam_R)
|
||||
R = self.cam_R.transpose(-1, -2)
|
||||
t = -torch.einsum("bij,bj->bi", R, self.cam_t)
|
||||
print('dataset.py, cam2world: ', R)
|
||||
return R, t
|
||||
|
||||
def as_dict(self):
|
||||
return {
|
||||
"cam_R": self.cam_R, # (T, 3, 3)
|
||||
"cam_t": self.cam_t, # (T, 3)
|
||||
"intrins": self.intrins, # (T, 4)
|
||||
"static": self.is_static, # bool
|
||||
}
|
||||
|
||||
|
||||
def get_ternary_mask(vis_mask):
|
||||
# get the track start and end idcs relative to the filtered interval
|
||||
vis_mask = torch.as_tensor(vis_mask)
|
||||
vis_idcs = torch.where(vis_mask)[0]
|
||||
track_s, track_e = min(vis_idcs), max(vis_idcs) + 1
|
||||
# -1 = track out of scene, 0 = occlusion, 1 = visible
|
||||
vis_mask = vis_mask.float()
|
||||
vis_mask[:track_s] = -1
|
||||
vis_mask[track_e:] = -1
|
||||
return vis_mask
|
||||
|
||||
|
||||
def get_shot_img_files(shots_path, shot_idx, shot_pad=SHOT_PAD):
|
||||
assert os.path.isfile(shots_path)
|
||||
with open(shots_path, "r") as f:
|
||||
shots_dict = json.load(f)
|
||||
img_names = sorted(shots_dict.keys())
|
||||
N = len(img_names)
|
||||
shot_mask = np.array([shots_dict[x] == shot_idx for x in img_names])
|
||||
|
||||
idcs = np.where(shot_mask)[0]
|
||||
if shot_pad > 0: # drop the frames before/after shot change
|
||||
raise ValueError
|
||||
if min(idcs) > 0:
|
||||
idcs = idcs[shot_pad:]
|
||||
if len(idcs) > 0 and max(idcs) < N - 1:
|
||||
idcs = idcs[:-shot_pad]
|
||||
if len(idcs) < MIN_SEQ_LEN:
|
||||
raise ValueError("shot is too short for optimization")
|
||||
|
||||
shot_mask = np.zeros(N, dtype=bool)
|
||||
shot_mask[idcs] = 1
|
||||
|
||||
sel_paths = [img_names[i] for i in idcs]
|
||||
print(f"FOUND {len(idcs)}/{len(shots_dict)} FRAMES FOR SHOT {shot_idx}")
|
||||
return sel_paths, idcs
|
||||
|
||||
|
||||
def load_cameras_npz(camera_path, v):
|
||||
assert os.path.splitext(camera_path)[-1] == ".npz"
|
||||
|
||||
cam_data = np.load(camera_path)
|
||||
height, width, focal = (
|
||||
int(cam_data["height"]),
|
||||
int(cam_data["width"]),
|
||||
float(cam_data["focal"]),
|
||||
)
|
||||
|
||||
w2c = torch.from_numpy(cam_data["w2c"]) # (N, 4, 4)
|
||||
cam_R = w2c[:, :3, :3] # (N, 3, 3)
|
||||
cam_t = w2c[:, :3, 3] # (N, 3)
|
||||
N = len(w2c)
|
||||
|
||||
if "intrins" in cam_data:
|
||||
intrins = torch.from_numpy(cam_data["intrins"].astype(np.float32))
|
||||
else:
|
||||
intrins = torch.tensor([focal, focal, width / 2, height / 2])[None].repeat(N, 1)
|
||||
|
||||
print(f"Loaded {N} cameras")
|
||||
return cam_R, cam_t, intrins, width, height
|
||||
|
||||
|
||||
def is_image(x):
|
||||
return (x.endswith(".png") or x.endswith(".jpg")) and not x.startswith(".")
|
||||
|
||||
|
||||
def get_name(x):
|
||||
return os.path.splitext(os.path.basename(x))[0]
|
||||
|
||||
|
||||
def split_name(x, suffix):
|
||||
return os.path.basename(x).split(suffix)[0]
|
||||
|
||||
|
||||
def get_names_in_dir(d, suffix):
|
||||
files = [split_name(x, suffix) for x in glob.glob(f"{d}/*{suffix}")]
|
||||
return sorted(files)
|
||||
|
||||
|
||||
def batch_join(parent, names, suffix=""):
|
||||
return [os.path.join(parent, f"{n}{suffix}") for n in names]
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
import os
|
||||
import json
|
||||
import functools
|
||||
|
||||
import numpy as np
|
||||
|
||||
from body_model import OP_NUM_JOINTS
|
||||
from scipy.interpolate import interp1d
|
||||
from scipy.spatial.transform import Rotation, Slerp
|
||||
|
||||
|
||||
def read_keypoints(keypoint_fn):
|
||||
"""
|
||||
Only reads body keypoint data of first person.
|
||||
"""
|
||||
empty_kps = np.zeros((OP_NUM_JOINTS, 3), dtype=np.float32)
|
||||
if not os.path.isfile(keypoint_fn):
|
||||
return empty_kps
|
||||
|
||||
with open(keypoint_fn) as keypoint_file:
|
||||
data = json.load(keypoint_file)
|
||||
|
||||
if len(data["people"]) == 0:
|
||||
print("WARNING: Found no keypoints in %s! Returning zeros!" % (keypoint_fn))
|
||||
return empty_kps
|
||||
|
||||
person_data = data["people"][0]
|
||||
body_keypoints = np.array(person_data["pose_keypoints_2d"], dtype=np.float32)
|
||||
body_keypoints = body_keypoints.reshape([-1, 3])
|
||||
return body_keypoints
|
||||
|
||||
|
||||
def load_keypoints_with_interp(kp_paths, interp=True):
|
||||
"""
|
||||
Load 2D keypoints from paths and interpolate missing frames between tmin and tmax.
|
||||
|
||||
Args:
|
||||
kp_paths: List of paths to keypoint JSON files
|
||||
interp: Whether to interpolate missing frames
|
||||
|
||||
Returns:
|
||||
joints2d_data: (T, J, 3) array of keypoints with interpolation
|
||||
"""
|
||||
# Load all keypoints
|
||||
joints2d_data = np.stack([read_keypoints(p) for p in kp_paths], axis=0).astype(np.float32)
|
||||
|
||||
if not interp:
|
||||
return joints2d_data
|
||||
|
||||
# Find which frames have valid keypoints (file exists and has data)
|
||||
vis_mask = np.array([os.path.isfile(p) for p in kp_paths])
|
||||
# Also check if keypoints are not all zeros
|
||||
has_data = ~np.all(joints2d_data == 0, axis=(1, 2))
|
||||
vis_mask = vis_mask & has_data
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
|
||||
if len(vis_idcs) == 0 or len(vis_idcs) == 1:
|
||||
# No interpolation needed
|
||||
return joints2d_data
|
||||
|
||||
# Interpolate only between tmin and tmax (same as MANO interpolation)
|
||||
T, J, _ = joints2d_data.shape
|
||||
tmin, tmax = min(vis_idcs), max(vis_idcs) + 1
|
||||
times = np.arange(tmin, tmax)
|
||||
|
||||
# Interpolate each joint's x, y coordinates separately
|
||||
for j in range(J):
|
||||
# Interpolate x coordinate
|
||||
x_interp = interp1d(vis_idcs, joints2d_data[vis_idcs, j, 0],
|
||||
kind='linear', bounds_error=False)
|
||||
# Interpolate y coordinate
|
||||
y_interp = interp1d(vis_idcs, joints2d_data[vis_idcs, j, 1],
|
||||
kind='linear', bounds_error=False)
|
||||
|
||||
# Interpolate for frames between tmin and tmax
|
||||
for t in times:
|
||||
if t not in vis_idcs:
|
||||
# Interpolate
|
||||
joints2d_data[t, j, 0] = x_interp(t)
|
||||
joints2d_data[t, j, 1] = y_interp(t)
|
||||
|
||||
# Set confidence as average of neighboring visible frames
|
||||
prev_vis = vis_idcs[vis_idcs < t]
|
||||
next_vis = vis_idcs[vis_idcs > t]
|
||||
if len(prev_vis) > 0 and len(next_vis) > 0:
|
||||
conf_prev = joints2d_data[prev_vis[-1], j, 2]
|
||||
conf_next = joints2d_data[next_vis[0], j, 2]
|
||||
joints2d_data[t, j, 2] = (conf_prev + conf_next) / 2.0 * 0.8 # Slightly lower confidence for interpolated
|
||||
elif len(prev_vis) > 0:
|
||||
joints2d_data[t, j, 2] = joints2d_data[prev_vis[-1], j, 2] * 0.8
|
||||
elif len(next_vis) > 0:
|
||||
joints2d_data[t, j, 2] = joints2d_data[next_vis[0], j, 2] * 0.8
|
||||
|
||||
return joints2d_data
|
||||
|
||||
|
||||
def read_mask_path(path):
|
||||
mask_path = None
|
||||
if not os.path.isfile(path):
|
||||
return mask_path
|
||||
|
||||
with open(path, "r") as f:
|
||||
data = json.load(path)
|
||||
|
||||
person_data = data["people"][0]
|
||||
if "mask_path" in person_data:
|
||||
mask_path = person_data["mask_path"]
|
||||
|
||||
return mask_path
|
||||
|
||||
|
||||
def read_mano_preds(pred_path, tid, num_betas=10):
|
||||
"""
|
||||
reads the betas, body_pose, global orientation and translation of a mano prediction
|
||||
exported from phalp outputs
|
||||
returns betas (10,), body_pose (23, 3), global_orientation (3,), translation (3,)
|
||||
"""
|
||||
pose = np.zeros((15, 3))
|
||||
rot = np.zeros(3)
|
||||
trans = np.zeros(3)
|
||||
betas = np.zeros(num_betas)
|
||||
if not os.path.isfile(pred_path):
|
||||
return pose, rot, trans, betas, int(tid)
|
||||
|
||||
with open(pred_path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
pose = np.array(data["body_pose"], dtype=np.float32)
|
||||
rot = np.array(data["global_orient"], dtype=np.float32)
|
||||
trans = np.array(data["cam_trans"], dtype=np.float32)
|
||||
betas = np.array(data["betas"], dtype=np.float32)
|
||||
is_right = np.array(data["is_right"], dtype=np.float32)
|
||||
|
||||
return pose, rot, trans, betas, is_right
|
||||
|
||||
|
||||
def load_mano_preds(pred_paths, tid, interp=True, num_betas=10):
|
||||
vis_mask = np.array([os.path.isfile(x) for x in pred_paths])
|
||||
vis_idcs = np.where(vis_mask)[0]
|
||||
|
||||
# load single image mano predictions
|
||||
stack_fnc = functools.partial(np.stack, axis=0)
|
||||
# (N, 16, 3), (N, 3), (N, 3), (N, 10)
|
||||
pose, orient, trans, betas, is_right = map(
|
||||
stack_fnc, zip(*[read_mano_preds(p, tid, num_betas=num_betas) for p in pred_paths])
|
||||
)
|
||||
|
||||
assert len(np.where(is_right!=int(tid))[0]) == 0
|
||||
if not interp:
|
||||
return pose, orient, trans, betas, is_right
|
||||
|
||||
# interpolate the occluded tracks
|
||||
orient_slerp = Slerp(vis_idcs, Rotation.from_rotvec(orient[vis_idcs]))
|
||||
trans_interp = interp1d(vis_idcs, trans[vis_idcs], axis=0)
|
||||
betas_interp = interp1d(vis_idcs, betas[vis_idcs], axis=0)
|
||||
|
||||
tmin, tmax = min(vis_idcs), max(vis_idcs) + 1
|
||||
times = np.arange(tmin, tmax)
|
||||
orient[times] = orient_slerp(times).as_rotvec()
|
||||
trans[times] = trans_interp(times)
|
||||
betas[times] = betas_interp(times)
|
||||
|
||||
# interpolate for each joint angle
|
||||
for i in range(pose.shape[1]):
|
||||
pose_slerp = Slerp(vis_idcs, Rotation.from_rotvec(pose[vis_idcs, i]))
|
||||
pose[times, i] = pose_slerp(times).as_rotvec()
|
||||
|
||||
return pose, orient, trans, betas, is_right
|
||||
+284
@@ -0,0 +1,284 @@
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import cv2
|
||||
|
||||
import preproc.launch_hamer as hamer
|
||||
from preproc.launch_slam import split_frames_shots, get_command, check_intrins
|
||||
from preproc.extract_frames import video_to_frames, split_frame
|
||||
|
||||
|
||||
def is_nonempty(d):
|
||||
return os.path.isdir(d) and len(os.listdir(d)) > 0
|
||||
|
||||
|
||||
def preprocess_frames(img_dir, src_path, overwrite=False, **kwargs):
|
||||
if not overwrite and is_nonempty(img_dir):
|
||||
print(f"FOUND {len(os.listdir(img_dir))} FRAMES in {img_dir}")
|
||||
return
|
||||
print(f"EXTRACTING FRAMES FROM {src_path} TO {img_dir}")
|
||||
print(kwargs)
|
||||
|
||||
# out = video_to_frames(src_path, img_dir, overwrite=overwrite, **kwargs)
|
||||
out = split_frame(src_path, img_dir, overwrite=overwrite, **kwargs)
|
||||
assert out == 0, "FAILED FRAME EXTRACTION"
|
||||
|
||||
|
||||
def preprocess_tracks(datatype, img_dir, track_dir, shot_dir, gpu, overwrite=False):
|
||||
"""
|
||||
:param img_dir
|
||||
:param track_dir, expected format: res_root/track_name/sequence
|
||||
:param shot_dir, expected format: res_root/shot_name/sequence
|
||||
"""
|
||||
if not overwrite and is_nonempty(track_dir):
|
||||
print(f"FOUND TRACKS IN {track_dir}")
|
||||
return
|
||||
|
||||
print(f"RUNNING HAMER ON {img_dir}")
|
||||
track_root, seq = os.path.split(track_dir.rstrip("/"))
|
||||
res_root, track_name = os.path.split(track_root)
|
||||
shot_name = shot_dir.rstrip("/").split("/")[-2]
|
||||
|
||||
hamer.process_seq(
|
||||
[gpu],
|
||||
res_root,
|
||||
seq,
|
||||
img_dir,
|
||||
track_name=track_name,
|
||||
shot_name=shot_name,
|
||||
datatype=datatype,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
|
||||
|
||||
def load_vipe_cameras(vipe_dir, seq_name, img_dir, start=0, end=-1):
|
||||
"""
|
||||
Load VIPE camera outputs and convert to DROID-SLAM format
|
||||
|
||||
Args:
|
||||
vipe_dir: Path to VIPE results directory
|
||||
seq_name: Sequence name
|
||||
img_dir: Image directory to get image size
|
||||
start: Start frame index
|
||||
end: End frame index
|
||||
|
||||
Returns:
|
||||
w2c: (N, 4, 4) world-to-camera matrices
|
||||
intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
|
||||
"""
|
||||
# Load VIPE pose and intrinsics
|
||||
vipe_pose_path = os.path.join(vipe_dir, "pose", f"{seq_name}.npz")
|
||||
vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{seq_name}.npz")
|
||||
|
||||
if not os.path.exists(vipe_pose_path):
|
||||
raise FileNotFoundError(f"VIPE pose file not found: {vipe_pose_path}")
|
||||
if not os.path.exists(vipe_intrins_path):
|
||||
raise FileNotFoundError(f"VIPE intrinsics file not found: {vipe_intrins_path}")
|
||||
|
||||
print(f"Loading VIPE cameras from {vipe_dir}")
|
||||
pose_data = np.load(vipe_pose_path)
|
||||
c2w = pose_data['data'] # (N, 4, 4) camera-to-world
|
||||
pose_inds = pose_data['inds'] # (N,) frame indices
|
||||
|
||||
intrins_data = np.load(vipe_intrins_path)
|
||||
intrins = intrins_data['data'] # (N, 4) [fx, fy, cx, cy]
|
||||
intrins_inds = intrins_data['inds'] # (N,) frame indices
|
||||
|
||||
# Verify indices match
|
||||
assert np.array_equal(pose_inds, intrins_inds), "Pose and intrinsics indices don't match!"
|
||||
|
||||
# Get image size from first image
|
||||
image_files = sorted([f for f in os.listdir(img_dir)
|
||||
if f.endswith(('.png', '.jpg', '.jpeg'))])
|
||||
if not image_files:
|
||||
# Infer from intrinsics (cx, cy should be roughly at center)
|
||||
img_width = int(intrins[0, 2] * 2)
|
||||
img_height = int(intrins[0, 3] * 2)
|
||||
print(f"Inferred image size from intrinsics: {img_width}x{img_height}")
|
||||
else:
|
||||
img_path = os.path.join(img_dir, image_files[0])
|
||||
img = cv2.imread(img_path)
|
||||
img_height, img_width = img.shape[:2]
|
||||
print(f"Got image size from images: {img_width}x{img_height}")
|
||||
|
||||
# Select frames based on start/end
|
||||
if end < 0:
|
||||
end = len(c2w)
|
||||
c2w = c2w[start:end]
|
||||
intrins = intrins[start:end]
|
||||
|
||||
# Convert camera-to-world to world-to-camera
|
||||
w2c = np.linalg.inv(c2w)
|
||||
|
||||
# Add width and height to intrinsics
|
||||
N = len(w2c)
|
||||
intrins_full = np.zeros((N, 6), dtype=np.float32)
|
||||
intrins_full[:, :4] = intrins
|
||||
intrins_full[:, 4] = img_width
|
||||
intrins_full[:, 5] = img_height
|
||||
|
||||
print(f"Loaded {N} VIPE camera poses")
|
||||
return w2c, intrins_full
|
||||
|
||||
|
||||
def save_vipe_cameras_as_droid(output_dir, w2c, intrins_full):
|
||||
"""
|
||||
Save VIPE cameras in DROID-SLAM format
|
||||
|
||||
Args:
|
||||
output_dir: Output directory
|
||||
w2c: (N, 4, 4) world-to-camera matrices
|
||||
intrins_full: (N, 6) intrinsics [fx, fy, cx, cy, W, H]
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Extract parameters
|
||||
W, H = intrins_full[0, 4], intrins_full[0, 5]
|
||||
focal = intrins_full[:, :2].mean()
|
||||
|
||||
print(f"Saving VIPE cameras to {output_dir}")
|
||||
print(f"Image size: {int(W)}x{int(H)}, focal: {focal:.2f}")
|
||||
|
||||
# Save cameras.npz (main format used by Dyn-HaMR)
|
||||
np.savez(
|
||||
f"{output_dir}/cameras.npz",
|
||||
height=H,
|
||||
width=W,
|
||||
focal=focal,
|
||||
intrins=intrins_full[:, :4],
|
||||
w2c=w2c,
|
||||
)
|
||||
print(f"Saved cameras.npz with {len(w2c)} frames")
|
||||
|
||||
|
||||
def run_vipe(video_path, vipe_dir, vipe_root=None, vipe_python=None, vipe_pipeline=None):
|
||||
"""Run VIPE without shell activation, using the project venv by default.
|
||||
|
||||
``VIPE_PYTHON`` or ``vipe_python`` can select a separate interpreter.
|
||||
Input/output paths are resolved before changing the child working directory.
|
||||
"""
|
||||
project_root = Path(__file__).resolve().parents[4]
|
||||
root = Path(vipe_root or project_root / "third_party" / "vipe").expanduser().resolve()
|
||||
video = Path(video_path).expanduser().absolute()
|
||||
output = Path(vipe_dir).expanduser().absolute()
|
||||
configured_python = vipe_python or os.environ.get("VIPE_PYTHON")
|
||||
project_python = project_root / ".venv" / "bin" / "python"
|
||||
python = Path(configured_python or (project_python if project_python.is_file() else sys.executable)).expanduser().absolute()
|
||||
if not video.is_file():
|
||||
raise FileNotFoundError(f"VIPE input video not found: {video}")
|
||||
if not (root / "vipe" / "cli" / "main.py").is_file():
|
||||
raise FileNotFoundError(f"VIPE source not found in {root}; set data.vipe_root")
|
||||
if not python.is_file():
|
||||
raise FileNotFoundError(f"VIPE interpreter not found: {python}; set VIPE_PYTHON")
|
||||
cmd = [str(python), "-m", "vipe.cli.main", "infer", str(video), "--output", str(output)]
|
||||
if vipe_pipeline:
|
||||
cmd.extend(["--pipeline", str(vipe_pipeline)])
|
||||
print(f"Running VIPE in {root}: {shlex.join(cmd)}", flush=True)
|
||||
out = subprocess.call(cmd, cwd=str(root))
|
||||
if out != 0:
|
||||
print(f"WARNING: VIPE failed with exit code {out}")
|
||||
else:
|
||||
print("VIPE completed successfully")
|
||||
return out
|
||||
|
||||
|
||||
def preprocess_cameras(cfg, overwrite=False):
|
||||
if not overwrite and is_nonempty(cfg.sources.cameras):
|
||||
print(f"FOUND CAMERAS IN {cfg.sources.cameras}")
|
||||
return
|
||||
|
||||
# Check if we should use VIPE instead of DROID-SLAM
|
||||
use_vipe = cfg.get("use_vipe", False)
|
||||
vipe_dir = cfg.get("vipe_dir", None)
|
||||
|
||||
if use_vipe and vipe_dir is not None:
|
||||
# Check if VIPE results exist for this sequence
|
||||
vipe_pose_path = os.path.join(vipe_dir, "pose", f"{cfg.seq}.npz")
|
||||
vipe_intrins_path = os.path.join(vipe_dir, "intrinsics", f"{cfg.seq}.npz")
|
||||
|
||||
# If VIPE results don't exist, run VIPE
|
||||
if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
|
||||
print(f"VIPE results not found for sequence '{cfg.seq}', running VIPE...")
|
||||
|
||||
# The source location is independent of the output directory.
|
||||
vipe_root = cfg.get("vipe_root", None)
|
||||
video_path = cfg.get("src_path", None)
|
||||
|
||||
if video_path is None or not os.path.exists(video_path):
|
||||
raise FileNotFoundError(
|
||||
f"Video path not found: {video_path}\n"
|
||||
f"Cannot run VIPE. Please provide a valid 'src_path' in your config."
|
||||
)
|
||||
|
||||
# Run VIPE
|
||||
out = run_vipe(video_path, vipe_dir, vipe_root, cfg.get("vipe_python", None), cfg.get("vipe_pipeline", None))
|
||||
if out != 0:
|
||||
raise RuntimeError(
|
||||
f"VIPE failed with exit code {out}\n"
|
||||
f"Check the command and error above. Select an installed Python "
|
||||
f"with VIPE_PYTHON or data.vipe_python; conda is not required."
|
||||
)
|
||||
|
||||
# Load VIPE results (after potentially running VIPE)
|
||||
if not (os.path.exists(vipe_pose_path) and os.path.exists(vipe_intrins_path)):
|
||||
raise FileNotFoundError(
|
||||
f"VIPE results not found after execution:\n"
|
||||
f" Pose: {vipe_pose_path}\n"
|
||||
f" Intrinsics: {vipe_intrins_path}\n"
|
||||
f"VIPE may have failed silently. Please check VIPE logs."
|
||||
)
|
||||
|
||||
print(f"USING VIPE CAMERAS FROM {vipe_dir}")
|
||||
img_dir = cfg.sources.images
|
||||
map_dir = cfg.sources.cameras
|
||||
|
||||
# Get frame range
|
||||
subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
|
||||
shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
|
||||
start, end = subseqs[shot_idx]
|
||||
|
||||
if not cfg.split_cameras:
|
||||
# only run on specified segment within shot
|
||||
end = start + cfg.end_idx
|
||||
start = start + cfg.start_idx
|
||||
|
||||
# Load and convert VIPE cameras
|
||||
w2c, intrins_full = load_vipe_cameras(vipe_dir, cfg.seq, img_dir, start, end)
|
||||
save_vipe_cameras_as_droid(map_dir, w2c, intrins_full)
|
||||
return
|
||||
|
||||
# Default: use DROID-SLAM
|
||||
print(f"RUNNING SLAM ON {cfg.seq}")
|
||||
img_dir = cfg.sources.images
|
||||
map_dir = cfg.sources.cameras
|
||||
subseqs, shot_idcs = split_frames_shots(cfg.sources.images, cfg.sources.shots)
|
||||
print(shot_idcs, cfg.shot_idx, np.where(shot_idcs == cfg.shot_idx), cfg.sources.images, cfg.sources.shots)
|
||||
print(subseqs)
|
||||
shot_idx = np.where(shot_idcs == cfg.shot_idx)[0][0]
|
||||
# run on selected shot
|
||||
start, end = subseqs[shot_idx]
|
||||
if not cfg.split_cameras:
|
||||
# only run on specified segment within shot
|
||||
end = start + cfg.end_idx
|
||||
start = start + cfg.start_idx
|
||||
intrins_path = cfg.sources.get("intrins", None)
|
||||
if intrins_path is not None:
|
||||
intrins_path = check_intrins(cfg.type, cfg.root, intrins_path, cfg.seq, cfg.split)
|
||||
|
||||
print('img_dir, map_dir, start, end, intrins_path', img_dir, map_dir, start, end, intrins_path)
|
||||
# raise ValueERRROR
|
||||
cmd = get_command(
|
||||
img_dir,
|
||||
map_dir,
|
||||
start=start,
|
||||
end=end,
|
||||
intrins_path=intrins_path,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
print(cmd)
|
||||
gpu = os.environ.get("CUDA_VISIBLE_DEVICES", 0)
|
||||
out = subprocess.call(f"CUDA_VISIBLE_DEVICES={gpu} {cmd}", shell=True)
|
||||
assert out == 0, "SLAM FAILED"
|
||||
Reference in New Issue
Block a user