File size: 6,138 Bytes
7e120dd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
#!/usr/bin/env python3
"""
Check internal consistency of Unity processed samples.

This script compares:
  - saved `smpl_params_w` vs `smpl_params_c` + `T_w2c`-derived world params
  - saved `cam_angvel` (if present) vs recomputed from `T_w2c`

If these disagree by non-trivial amounts, fine-tuning can "break global" because the
network is asked to fit mutually inconsistent targets/condition signals.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Dict, Any, Iterable

import torch

from genmo.utils.geo_transform import compute_cam_angvel, normalize_T_w2c
from genmo.utils.rotation_conversions import axis_angle_to_matrix, matrix_to_axis_angle


def _load_pt(path: Path) -> Dict[str, Any]:
    return torch.load(path, map_location="cpu", weights_only=False)


def _axis_angle_angle_deg(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Geodesic angle between rotations in axis-angle form. Shapes (..., 3). Returns (...,) degrees."""
    Ra = axis_angle_to_matrix(a.float())
    Rb = axis_angle_to_matrix(b.float())
    Rrel = Ra.transpose(-1, -2) @ Rb
    aa = matrix_to_axis_angle(Rrel)
    return aa.norm(dim=-1) * (180.0 / torch.pi)


def _derive_world_from_c_and_T(
    smpl_params_c: Dict[str, torch.Tensor], T_w2c: torch.Tensor
) -> Dict[str, torch.Tensor]:
    """
    Derive world translation/orientation from camera params and T_w2c.

    Assumes:
      p_w = R_c2w * p_c + t_c2w
      R_w = R_c2w * R_c
    """
    R_w2c = T_w2c[:, :3, :3].float()  # (L,3,3)
    t_w2c = T_w2c[:, :3, 3].float()  # (L,3)
    R_c2w = R_w2c.transpose(-1, -2)
    t_c2w = -torch.einsum("fij,fj->fi", R_c2w, t_w2c)

    transl_c = smpl_params_c["transl"].float()
    transl_w = torch.einsum("fij,fj->fi", R_c2w, transl_c) + t_c2w

    R_c = axis_angle_to_matrix(smpl_params_c["global_orient"].float())
    R_w = torch.einsum("fij,fjk->fik", R_c2w, R_c)
    go_w = matrix_to_axis_angle(R_w)

    return {"transl": transl_w, "global_orient": go_w}


def _maybe_slice(x: Any, start: int, end: int) -> Any:
    if isinstance(x, dict):
        return {k: _maybe_slice(v, start, end) for k, v in x.items()}
    if isinstance(x, torch.Tensor):
        return x[start:end]
    return x


def _iter_paths(root: Path, glob_pat: str) -> Iterable[Path]:
    if root.is_file():
        yield root
        return
    yield from sorted(root.glob(glob_pat))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--root",
        type=Path,
        default=Path("processed_dataset/genmo_features"),
        help="A single .pt file or a directory containing .pt files.",
    )
    ap.add_argument("--glob", type=str, default="*.pt")
    ap.add_argument("--max", type=int, default=10, help="Max files to check")
    ap.add_argument("--frame0_only", action="store_true", help="Check only frame 0")
    args = ap.parse_args()

    paths = list(_iter_paths(args.root, args.glob))
    if not paths:
        raise SystemExit(f"No files found under {args.root} with glob {args.glob!r}")
    paths = paths[: max(1, int(args.max))]

    all_stats = []
    for p in paths:
        d = _load_pt(p)
        smpl_c = d["smpl_params_c"]
        smpl_w = d["smpl_params_w"]
        T_w2c = d["T_w2c"]

        if args.frame0_only:
            smpl_c = _maybe_slice(smpl_c, 0, 1)
            smpl_w = _maybe_slice(smpl_w, 0, 1)
            T_w2c = _maybe_slice(T_w2c, 0, 1)

        derived = _derive_world_from_c_and_T(smpl_c, T_w2c)

        diff_t = smpl_w["transl"].float() - derived["transl"]  # (L, 3)
        dt = diff_t.norm(dim=-1)  # (L,)
        diff_t_abs_mean = diff_t.abs().mean(dim=0)
        diff_t_abs_max = diff_t.abs().max(dim=0)[0]
        dgo = _axis_angle_angle_deg(smpl_w["global_orient"].float(), derived["global_orient"])  # (L,)

        # cam_angvel consistency
        cam_av_saved = d.get("cam_angvel", None)
        cam_av_deg = None
        if cam_av_saved is not None:
            cam_av_saved = cam_av_saved[: T_w2c.shape[0]]
            Tw = T_w2c.float()
            if Tw.ndim == 3 and Tw.shape[-2:] == (4, 4) and Tw.shape[0] >= 2:
                Tw = normalize_T_w2c(Tw)
                cam_av_re = compute_cam_angvel(Tw[:, :3, :3])  # (L,6)
                # matrix_to_rotation_6d convention: compare directly in 6D space
                cam_av_deg = (cam_av_saved.float() - cam_av_re.float()).abs().mean().item()

        stat = {
            "file": p.name,
            "max_transl_err_m": float(dt.max().item()),
            "mean_transl_err_m": float(dt.mean().item()),
            "mean_abs_dx": float(diff_t_abs_mean[0].item()),
            "mean_abs_dy": float(diff_t_abs_mean[1].item()),
            "mean_abs_dz": float(diff_t_abs_mean[2].item()),
            "max_abs_dx": float(diff_t_abs_max[0].item()),
            "max_abs_dy": float(diff_t_abs_max[1].item()),
            "max_abs_dz": float(diff_t_abs_max[2].item()),
            "max_go_err_deg": float(dgo.max().item()),
            "mean_go_err_deg": float(dgo.mean().item()),
            "cam_angvel_absmean_diff": None if cam_av_deg is None else float(cam_av_deg),
        }
        all_stats.append(stat)

        print(
            f"{p.name}: "
            f"transl_err mean={stat['mean_transl_err_m']:.4f}m max={stat['max_transl_err_m']:.4f}m | "
            f"|Δt| mean_abs(x,y,z)=({stat['mean_abs_dx']:.4f},{stat['mean_abs_dy']:.4f},{stat['mean_abs_dz']:.4f}) "
            f"max_abs(x,y,z)=({stat['max_abs_dx']:.4f},{stat['max_abs_dy']:.4f},{stat['max_abs_dz']:.4f}) | "
            f"go_err mean={stat['mean_go_err_deg']:.2f}° max={stat['max_go_err_deg']:.2f}°"
            + (
                ""
                if stat["cam_angvel_absmean_diff"] is None
                else f" | cam_angvel_absmean_diff={stat['cam_angvel_absmean_diff']:.6f}"
            )
        )

    # Aggregate
    max_transl = max(s["max_transl_err_m"] for s in all_stats)
    max_go = max(s["max_go_err_deg"] for s in all_stats)
    print(f"\nAggregate over {len(all_stats)} files: max_transl_err_m={max_transl:.4f}, max_go_err_deg={max_go:.2f}")


if __name__ == "__main__":
    main()