vr-hmr / debug_unity_data.py

Upload folder using huggingface_hub

7e120dd 8 months ago

6.64 kB

#!/usr/bin/env python3
"""
Diagnostic script to check Unity dataset coordinate consistency.
Checks:
- Rotation consistency between `smpl_params_c`, `smpl_params_w`, and `T_w2c`
- `cam_angvel` matches the convention used in preprocessing
- SMPL forward-kinematics consistency between camera/world parameters
Run with the gvhmr env python:
  /root/miniconda3/envs/gvhmr/bin/python debug_unity_data.py
"""
import torch
import numpy as np
from pathlib import Path
from scipy.spatial.transform import Rotation as R

def axis_angle_to_matrix(aa):
    """Convert axis-angle to rotation matrix (numpy)."""
    return R.from_rotvec(aa).as_matrix()

def check_single_sequence(pt_path):
    """Check a single .pt file for coordinate consistency."""
    print(f"\n{'='*80}")
    print(f"Checking: {pt_path.name}")
    print(f"{'='*80}")

    data = torch.load(pt_path, map_location="cpu")

    # Extract key data
    smpl_c = data["smpl_params_c"]
    smpl_w = data["smpl_params_w"]
    T_w2c = data["T_w2c"].numpy()

    # Check first frame
    idx = 0
    print(f"\n[Frame {idx}]")

    # Ground truth
    go_c_gt = smpl_c["global_orient"][idx].numpy()  # (3,) axis-angle
    go_w_gt = smpl_w["global_orient"][idx].numpy()  # (3,) axis-angle

    # Convert to matrices
    R_c_gt = axis_angle_to_matrix(go_c_gt)  # Pelvis in camera frame
    R_w_gt = axis_angle_to_matrix(go_w_gt)  # Pelvis in world frame
    R_w2c = T_w2c[idx, :3, :3]  # World to camera
    R_c2w = R_w2c.T  # Camera to world

    # Verify: R_w = R_c2w @ R_c
    R_w_reconstructed = R_c2w @ R_c_gt

    # Compare
    R_diff = R_w_reconstructed @ R_w_gt.T
    angle_err_deg = np.linalg.norm(R.from_matrix(R_diff).as_rotvec()) * 180.0 / np.pi

    print(f"Ground truth global_orient_c (axis-angle): {go_c_gt}")
    print(f"Ground truth global_orient_w (axis-angle): {go_w_gt}")
    print(f"\nReconstruction test: R_w = R_c2w @ R_c")
    print(f"  Rotation error: {angle_err_deg:.4f}°")

    if angle_err_deg > 1.0:
        print(f"  ❌ ERROR: Rotation mismatch > 1°!")
        print(f"  R_w (ground truth):\n{R_w_gt}")
        print(f"  R_w (reconstructed):\n{R_w_reconstructed}")
    else:
        print(f"  ✅ OK: Rotations are consistent")

    # Check cam_angvel computation (should match preprocess convention)
    print(f"\n[Camera Angular Velocity Check]")
    cam_ok = True
    if "cam_angvel" in data:
        cam_angvel = data["cam_angvel"]  # (L, 6) - 6D rotation
        print(f"  cam_angvel shape: {cam_angvel.shape}")
        print(f"  cam_angvel[0]: {cam_angvel[0].numpy()}")

        # Manually compute cam_angvel and compare.
        # Convention (see `tools/demo/process_dataset.py:compute_velocity`):
        #   cam_angvel[0] = [1,0,0, 0,1,0]  (identity, rotation6d)
        #   cam_angvel[i] = rot6d(R_i @ R_{i-1}^T)
        from genmo.utils.rotation_conversions import matrix_to_rotation_6d
        R_w2c_t = torch.from_numpy(T_w2c[:, :3, :3]).float()
        L = int(R_w2c_t.shape[0])
        cam_angvel_manual = torch.zeros((L, 6), dtype=torch.float32)
        cam_angvel_manual[0] = cam_angvel_manual.new_tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0])
        if L > 1:
            R_diff_manual = R_w2c_t[1:] @ R_w2c_t[:-1].transpose(-1, -2)
            cam_angvel_manual[1:] = matrix_to_rotation_6d(R_diff_manual)

        diff = (cam_angvel - cam_angvel_manual).abs().max()
        print(f"  Manual vs stored cam_angvel max diff: {diff:.6f}")
        if diff > 1e-4:
            print(f"  ❌ WARNING: cam_angvel mismatch!")
            cam_ok = False
        else:
            print(f"  ✅ OK: cam_angvel matches manual computation")

    # Check SMPL forward kinematics consistency
    print(f"\n[SMPL FK Check]")
    fk_ok = True
    try:
        from third_party.GVHMR.hmr4d.utils.smplx_utils import make_smplx
        smplx_model = make_smplx("supermotion").eval()

        with torch.no_grad():
            # Incam SMPL
            out_c = smplx_model(
                global_orient=smpl_c["global_orient"][idx:idx+1],
                body_pose=smpl_c["body_pose"][idx:idx+1],
                betas=smpl_c["betas"][idx:idx+1],
                transl=smpl_c["transl"][idx:idx+1]
            )
            joints_c = out_c.joints[0, :22].numpy()  # (22, 3)

            # Global SMPL
            out_w = smplx_model(
                global_orient=smpl_w["global_orient"][idx:idx+1],
                body_pose=smpl_w["body_pose"][idx:idx+1],
                betas=smpl_w["betas"][idx:idx+1],
                transl=smpl_w["transl"][idx:idx+1]
            )
            joints_w = out_w.joints[0, :22].numpy()  # (22, 3)

        # Transform camera->world using T_w2c (world->camera):
        #   x_c = R_w2c x_w + t_w2c
        # => x_w = R_w2c^T (x_c - t_w2c)
        t_w2c = T_w2c[idx, :3, 3]
        joints_c2w = (R_c2w @ (joints_c - t_w2c).T).T

        # Compare
        joint_err = np.linalg.norm(joints_c2w - joints_w, axis=-1).mean()
        print(f"  Mean joint error (incam→world vs world GT): {joint_err:.4f}m")

        if joint_err > 0.05:
            print(f"  ❌ ERROR: Joint mismatch > 5cm!")
            fk_ok = False
        else:
            print(f"  ✅ OK: SMPL joints are consistent")

    except Exception as e:
        print(f"  ⚠️  Could not run SMPL FK check: {e}")
        fk_ok = False

    # Consider the clip consistent only if all checks are reasonable.
    ok_rot = angle_err_deg < 1.0
    return ok_rot and cam_ok and fk_ok

def main():
    dataset_root = Path("./processed_dataset")
    feat_dir = dataset_root / "genmo_features"

    if not feat_dir.exists():
        print(f"Error: {feat_dir} not found!")
        return

    pt_files = sorted(list(feat_dir.glob("*.pt")))
    print(f"Found {len(pt_files)} sequences")

    if len(pt_files) == 0:
        print("No .pt files found!")
        return

    # Check first 3 sequences
    num_check = min(3, len(pt_files))
    all_ok = True

    for i in range(num_check):
        ok = check_single_sequence(pt_files[i])
        all_ok = all_ok and ok

    print(f"\n{'='*80}")
    if all_ok:
        print("✅ All checks passed! Data appears consistent.")
        print("\nIf training still has high loss, the issue is likely:")
        print("  1. Model architecture/hyperparameters")
        print("  2. Normalization statistics mismatch")
        print("  3. Sequence length handling during training")
    else:
        print("❌ Data consistency issues found!")
        print("\nThis explains the high training loss.")
        print("You need to fix the coordinate system in process_dataset.py")
    print(f"{'='*80}\n")

if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	Diagnostic script to check Unity dataset coordinate consistency.

	Checks:
	- Rotation consistency between `smpl_params_c`, `smpl_params_w`, and `T_w2c`
	- `cam_angvel` matches the convention used in preprocessing
	- SMPL forward-kinematics consistency between camera/world parameters

	Run with the gvhmr env python:
	/root/miniconda3/envs/gvhmr/bin/python debug_unity_data.py
	"""
	import torch
	import numpy as np
	from pathlib import Path
	from scipy.spatial.transform import Rotation as R

	def axis_angle_to_matrix(aa):
	"""Convert axis-angle to rotation matrix (numpy)."""
	return R.from_rotvec(aa).as_matrix()

	def check_single_sequence(pt_path):
	"""Check a single .pt file for coordinate consistency."""
	print(f"\n{'='*80}")
	print(f"Checking: {pt_path.name}")
	print(f"{'='*80}")

	data = torch.load(pt_path, map_location="cpu")

	# Extract key data
	smpl_c = data["smpl_params_c"]
	smpl_w = data["smpl_params_w"]
	T_w2c = data["T_w2c"].numpy()

	# Check first frame
	idx = 0
	print(f"\n[Frame {idx}]")

	# Ground truth
	go_c_gt = smpl_c["global_orient"][idx].numpy() # (3,) axis-angle
	go_w_gt = smpl_w["global_orient"][idx].numpy() # (3,) axis-angle

	# Convert to matrices
	R_c_gt = axis_angle_to_matrix(go_c_gt) # Pelvis in camera frame
	R_w_gt = axis_angle_to_matrix(go_w_gt) # Pelvis in world frame
	R_w2c = T_w2c[idx, :3, :3] # World to camera
	R_c2w = R_w2c.T # Camera to world

	# Verify: R_w = R_c2w @ R_c
	R_w_reconstructed = R_c2w @ R_c_gt

	# Compare
	R_diff = R_w_reconstructed @ R_w_gt.T
	angle_err_deg = np.linalg.norm(R.from_matrix(R_diff).as_rotvec()) * 180.0 / np.pi

	print(f"Ground truth global_orient_c (axis-angle): {go_c_gt}")
	print(f"Ground truth global_orient_w (axis-angle): {go_w_gt}")
	print(f"\nReconstruction test: R_w = R_c2w @ R_c")
	print(f" Rotation error: {angle_err_deg:.4f}°")

	if angle_err_deg > 1.0:
	print(f" ❌ ERROR: Rotation mismatch > 1°!")
	print(f" R_w (ground truth):\n{R_w_gt}")
	print(f" R_w (reconstructed):\n{R_w_reconstructed}")
	else:
	print(f" ✅ OK: Rotations are consistent")

	# Check cam_angvel computation (should match preprocess convention)
	print(f"\n[Camera Angular Velocity Check]")
	cam_ok = True
	if "cam_angvel" in data:
	cam_angvel = data["cam_angvel"] # (L, 6) - 6D rotation
	print(f" cam_angvel shape: {cam_angvel.shape}")
	print(f" cam_angvel[0]: {cam_angvel[0].numpy()}")

	# Manually compute cam_angvel and compare.
	# Convention (see `tools/demo/process_dataset.py:compute_velocity`):
	# cam_angvel[0] = [1,0,0, 0,1,0] (identity, rotation6d)
	# cam_angvel[i] = rot6d(R_i @ R_{i-1}^T)
	from genmo.utils.rotation_conversions import matrix_to_rotation_6d
	R_w2c_t = torch.from_numpy(T_w2c[:, :3, :3]).float()
	L = int(R_w2c_t.shape[0])
	cam_angvel_manual = torch.zeros((L, 6), dtype=torch.float32)
	cam_angvel_manual[0] = cam_angvel_manual.new_tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0])
	if L > 1:
	R_diff_manual = R_w2c_t[1:] @ R_w2c_t[:-1].transpose(-1, -2)
	cam_angvel_manual[1:] = matrix_to_rotation_6d(R_diff_manual)

	diff = (cam_angvel - cam_angvel_manual).abs().max()
	print(f" Manual vs stored cam_angvel max diff: {diff:.6f}")
	if diff > 1e-4:
	print(f" ❌ WARNING: cam_angvel mismatch!")
	cam_ok = False
	else:
	print(f" ✅ OK: cam_angvel matches manual computation")

	# Check SMPL forward kinematics consistency
	print(f"\n[SMPL FK Check]")
	fk_ok = True
	try:
	from third_party.GVHMR.hmr4d.utils.smplx_utils import make_smplx
	smplx_model = make_smplx("supermotion").eval()

	with torch.no_grad():
	# Incam SMPL
	out_c = smplx_model(
	global_orient=smpl_c["global_orient"][idx:idx+1],
	body_pose=smpl_c["body_pose"][idx:idx+1],
	betas=smpl_c["betas"][idx:idx+1],
	transl=smpl_c["transl"][idx:idx+1]
	)
	joints_c = out_c.joints[0, :22].numpy() # (22, 3)

	# Global SMPL
	out_w = smplx_model(
	global_orient=smpl_w["global_orient"][idx:idx+1],
	body_pose=smpl_w["body_pose"][idx:idx+1],
	betas=smpl_w["betas"][idx:idx+1],
	transl=smpl_w["transl"][idx:idx+1]
	)
	joints_w = out_w.joints[0, :22].numpy() # (22, 3)

	# Transform camera->world using T_w2c (world->camera):
	# x_c = R_w2c x_w + t_w2c
	# => x_w = R_w2c^T (x_c - t_w2c)
	t_w2c = T_w2c[idx, :3, 3]
	joints_c2w = (R_c2w @ (joints_c - t_w2c).T).T

	# Compare
	joint_err = np.linalg.norm(joints_c2w - joints_w, axis=-1).mean()
	print(f" Mean joint error (incam→world vs world GT): {joint_err:.4f}m")

	if joint_err > 0.05:
	print(f" ❌ ERROR: Joint mismatch > 5cm!")
	fk_ok = False
	else:
	print(f" ✅ OK: SMPL joints are consistent")

	except Exception as e:
	print(f" ⚠️ Could not run SMPL FK check: {e}")
	fk_ok = False

	# Consider the clip consistent only if all checks are reasonable.
	ok_rot = angle_err_deg < 1.0
	return ok_rot and cam_ok and fk_ok

	def main():
	dataset_root = Path("./processed_dataset")
	feat_dir = dataset_root / "genmo_features"

	if not feat_dir.exists():
	print(f"Error: {feat_dir} not found!")
	return

	pt_files = sorted(list(feat_dir.glob("*.pt")))
	print(f"Found {len(pt_files)} sequences")

	if len(pt_files) == 0:
	print("No .pt files found!")
	return

	# Check first 3 sequences
	num_check = min(3, len(pt_files))
	all_ok = True

	for i in range(num_check):
	ok = check_single_sequence(pt_files[i])
	all_ok = all_ok and ok

	print(f"\n{'='*80}")
	if all_ok:
	print("✅ All checks passed! Data appears consistent.")
	print("\nIf training still has high loss, the issue is likely:")
	print(" 1. Model architecture/hyperparameters")
	print(" 2. Normalization statistics mismatch")
	print(" 3. Sequence length handling during training")
	else:
	print("❌ Data consistency issues found!")
	print("\nThis explains the high training loss.")
	print("You need to fix the coordinate system in process_dataset.py")
	print(f"{'='*80}\n")

	if __name__ == "__main__":
	main()