File size: 2,769 Bytes
fbd9366 | 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 | """Hand seeds for the canonical T-Rex 250-point CoTracker layout.
No appearance auto-repair. Tune prompts in ``sam2_prompt_hands.py``.
"""
from __future__ import annotations
import numpy as np
from trex_track.layout import NUM_HEAD_PER_HAND
NUM_HAND_POINTS = NUM_HEAD_PER_HAND
__all__ = [
"NUM_HAND_POINTS",
"wrist_hand_50",
"head_hands_50",
# Compatibility aliases for callers of the old 25-point implementation.
"wrist_hand_25",
"head_hands_25",
]
def _stable_spatial_order(points_xy: np.ndarray) -> np.ndarray:
"""Give sampled slots deterministic top-to-bottom, then left-to-right IDs."""
points = np.asarray(points_xy, dtype=np.float32).reshape(-1, 2)
order = np.lexsort((points[:, 0], points[:, 1]))
return points[order]
def wrist_hand_50(
predictor,
rgb_hwc: np.ndarray,
view: str,
*,
n_points: int = NUM_HAND_POINTS,
seed: int | None = 0,
) -> tuple[np.ndarray, np.ndarray, str]:
"""One SAM2 pass → 50 stably ordered queries on a wrist hand mask.
Background 5×5 grid is NOT produced here — callers prepend ``make_image_grid``.
"""
from trex_track.sam2_prompt_hands import sample_points_on_mask, sam2_wrist_mask
mask, tag, _meta = sam2_wrist_mask(predictor, rgb_hwc, view)
rng = np.random.default_rng(seed)
pts = sample_points_on_mask(mask, n_points, rng=rng, distal_bias=True)
h, w = rgb_hwc.shape[:2]
pts[:, 0] = np.clip(pts[:, 0], 0, w - 1)
pts[:, 1] = np.clip(pts[:, 1], 0, h - 1)
return _stable_spatial_order(pts), mask.astype(bool), tag
def head_hands_50(
predictor,
rgb_hwc: np.ndarray,
*,
n_points: int = NUM_HAND_POINTS,
seed: int | None = 0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, str]:
"""One SAM2 pass per hand → 100 queries (left 50, then right 50)."""
from trex_track.sam2_prompt_hands import sample_points_on_mask, sam2_head_hand_mask
rng = np.random.default_rng(seed)
left, tag_l, _ = sam2_head_hand_mask(predictor, rgb_hwc, "left")
right, tag_r, _ = sam2_head_hand_mask(predictor, rgb_hwc, "right")
pts_l = _stable_spatial_order(
sample_points_on_mask(left, n_points, rng=rng, distal_bias=True)
)
pts_r = _stable_spatial_order(
sample_points_on_mask(right, n_points, rng=rng, distal_bias=True)
)
h, w = rgb_hwc.shape[:2]
queries = np.concatenate([pts_l, pts_r], axis=0).astype(np.float32)
queries[:, 0] = np.clip(queries[:, 0], 0, w - 1)
queries[:, 1] = np.clip(queries[:, 1], 0, h - 1)
tag = f"{tag_l}|{tag_r}"
return queries, left.astype(bool), right.astype(bool), tag
# Keep old imports working while making their new 50-point behavior explicit.
wrist_hand_25 = wrist_hand_50
head_hands_25 = head_hands_50
|