diff --git a/data_only_viz/arkit_skeleton.py b/data_only_viz/arkit_skeleton.py index ae9102a..656a45a 100644 --- a/data_only_viz/arkit_skeleton.py +++ b/data_only_viz/arkit_skeleton.py @@ -6,6 +6,8 @@ the GPU line buffer. """ from __future__ import annotations +import math + def bones_from_parents(parents: list[int]) -> list[tuple[int, int]]: """(child, parent) pairs for every joint with a valid parent.""" @@ -13,14 +15,20 @@ def bones_from_parents(parents: list[int]) -> list[tuple[int, int]]: return [(i, p) for i, p in enumerate(parents) if 0 <= p < n] -def arkit_segments(arr2d, valid, parents): +def arkit_segments(arr2d, valid, parents, max_len: float = 0.5, bounds: float = 0.1): """Return (ax, ay, bx, by) for each bone with both endpoints valid. - arr2d: indexable of (x, y) normalized [0,1], length == len(parents). - valid: indexable of bool (length == len(parents)) or None to keep all. + arr2d: indexable of (x, y) normalized [0,1], length == len(parents). + valid: indexable of bool (length == len(parents)) or None to keep all. parents: parent index per joint (-1 = root). + max_len: maximum bone length in normalized units; bones longer than this + are skipped (garbage coords during ARKit tracking loss). + bounds: margin outside [0,1] that is still accepted; endpoints outside + [-bounds, 1+bounds] in x or y are rejected. """ n = len(parents) + lo = -bounds + hi = 1.0 + bounds segs: list[tuple[float, float, float, float]] = [] for child, parent in bones_from_parents(parents): if child >= n or parent >= n: @@ -29,5 +37,13 @@ def arkit_segments(arr2d, valid, parents): continue ax, ay = float(arr2d[child][0]), float(arr2d[child][1]) bx, by = float(arr2d[parent][0]), float(arr2d[parent][1]) + # Reject out-of-frame endpoints + if not (lo <= ax <= hi and lo <= ay <= hi): + continue + if not (lo <= bx <= hi and lo <= by <= hi): + continue + # Reject implausibly long bones (tracking-loss garbage) + if math.hypot(bx - ax, by - ay) > max_len: + continue segs.append((ax, ay, bx, by)) return segs diff --git a/data_only_viz/tests/test_arkit_skeleton.py b/data_only_viz/tests/test_arkit_skeleton.py index 381744d..f31c6f7 100644 --- a/data_only_viz/tests/test_arkit_skeleton.py +++ b/data_only_viz/tests/test_arkit_skeleton.py @@ -13,15 +13,16 @@ def test_bones_skip_out_of_range(): def test_segments_basic(): - arr2d = [(0.0, 0.0), (0.5, 0.5), (1.0, 1.0)] + # Realistic bone lengths < 0.5 (horizontal spine-like chain) + arr2d = [(0.0, 0.0), (0.2, 0.0), (0.4, 0.0)] valid = [True, True, True] parents = [-1, 0, 1] segs = arkit_segments(arr2d, valid, parents) - assert segs == [(0.5, 0.5, 0.0, 0.0), (1.0, 1.0, 0.5, 0.5)] + assert segs == [(0.2, 0.0, 0.0, 0.0), (0.4, 0.0, 0.2, 0.0)] def test_segments_drop_invalid_endpoint(): - arr2d = [(0.0, 0.0), (0.5, 0.5), (1.0, 1.0)] + arr2d = [(0.0, 0.0), (0.2, 0.0), (0.4, 0.0)] valid = [True, False, True] # joint 1 invalid parents = [-1, 0, 1] # bone (1,0) drops (1 invalid); bone (2,1) drops (1 invalid) @@ -29,6 +30,53 @@ def test_segments_drop_invalid_endpoint(): def test_segments_valid_none_keeps_all(): - arr2d = [(0.0, 0.0), (0.5, 0.5)] + arr2d = [(0.0, 0.0), (0.3, 0.0)] # bone length 0.3 < max_len=0.5 parents = [-1, 0] - assert arkit_segments(arr2d, None, parents) == [(0.5, 0.5, 0.0, 0.0)] + assert arkit_segments(arr2d, None, parents) == [(0.3, 0.0, 0.0, 0.0)] + + +# --------------------------------------------------------------------------- +# New filter tests (bounds + max_len) +# --------------------------------------------------------------------------- + +def test_segments_out_of_bounds_endpoint_skipped(): + """Endpoint outside [-bounds, 1+bounds] → bone skipped.""" + # Joint 1 at x=-0.2 is outside [-0.1, 1.1] + arr2d = [(0.5, 0.5), (-0.2, 0.5), (0.7, 0.5)] + valid = [True, True, True] + parents = [-1, 0, 1] + # bone (1,0): joint 1 out of bounds → skip + # bone (2,1): joint 1 out of bounds → skip + segs = arkit_segments(arr2d, valid, parents, bounds=0.1) + assert segs == [] + + +def test_segments_overlong_bone_skipped(): + """Bone length > max_len → bone skipped.""" + # Bone from (0.0, 0.5) to (1.0, 0.5): length=1.0 > max_len=0.5 + arr2d = [(0.0, 0.5), (1.0, 0.5)] + valid = [True, True] + parents = [-1, 0] + segs = arkit_segments(arr2d, valid, parents, max_len=0.5) + assert segs == [] + + +def test_segments_normal_skeleton_passes_with_defaults(): + """Realistic bones (length < 0.5, coords in [0,1]) pass all filters.""" + arr2d = [(0.5, 0.3), (0.5, 0.6), (0.5, 0.8)] # vertical spine-like + valid = [True, True, True] + parents = [-1, 0, 1] + segs = arkit_segments(arr2d, valid, parents) + # Both bones length = 0.3 < max_len=0.5; all coords in [0,1] ⊂ [-0.1,1.1] + assert len(segs) == 2 + assert segs[0] == (0.5, 0.6, 0.5, 0.3) + assert segs[1] == (0.5, 0.8, 0.5, 0.6) + + +def test_segments_in_bounds_endpoint_passes(): + """Endpoint exactly at boundary (0.0, 1.0) is kept.""" + arr2d = [(0.5, 0.5), (0.5, 0.2)] # bone length 0.3 < 0.5 + valid = [True, True] + parents = [-1, 0] + segs = arkit_segments(arr2d, valid, parents, bounds=0.1) + assert len(segs) == 1