Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+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