File size: 3,910 Bytes
8e5456b | 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 | """Keypoint layout, so a trimmed keypoint set can be trained without hardcoding 128.
The full DWPose layout is 128 keypoints: body 0:18, face 18:86, left hand 86:107,
right hand 107:128. Dropping keypoints (e.g. knees/ankles, which these clips crop
out) changes the model's input dimension, so the group boundaries can no longer be
constants.
Keypoints are always re-ordered into body / face / lhand / rhand blocks, so group
ranges stay contiguous slices -- only the boundaries move. A layout is persisted as
`layout.json` next to the packed memmaps.
"""
import json
import os
# --- the full DWPose layout, used when no layout.json is present -------------
FULL_GROUPS = {"body": (0, 18), "face": (18, 86), "lhand": (86, 107), "rhand": (107, 128)}
FULL_NKP = 128
# OpenPose-18 body indices, for reference when choosing what to drop:
# 0 nose 1 neck 2 Rshoulder 3 Relbow 4 Rwrist 5 Lshoulder 6 Lelbow 7 Lwrist
# 8 Rhip 9 Rknee 10 Rankle 11 Lhip 12 Lknee 13 Lankle
# 14 Reye 15 Leye 16 Rear 17 Lear
# Knees/ankles are detected in only ~19% / ~0.1% of Full_TriVis frames.
LOWER_BODY = (9, 10, 12, 13)
PRESETS = {
"full": (),
"upper": LOWER_BODY, # drop knees + ankles, keep hips (100% valid)
"upper_nohip": LOWER_BODY + (8, 11),
}
class Layout:
"""Which original DWPose keypoints are kept, and where the groups sit."""
def __init__(self, keep, groups, name="full"):
self.keep = list(keep) # original indices, in new order
self.groups = {k: tuple(v) for k, v in groups.items()}
self.name = name
self.n_kpts = len(self.keep)
self.dim = self.n_kpts * 2
# original index -> new index (None if dropped)
self.old2new = {o: n for n, o in enumerate(self.keep)}
# ------------------------------------------------------------------ build
@classmethod
def full(cls):
return cls(list(range(FULL_NKP)), FULL_GROUPS, "full")
@classmethod
def from_drop(cls, drop, name="custom"):
"""Build a layout that drops `drop` (original indices), keeping group order."""
drop = set(int(d) for d in drop)
keep, groups, cur = [], {}, 0
for g, (a, b) in FULL_GROUPS.items():
idx = [i for i in range(a, b) if i not in drop]
keep.extend(idx)
groups[g] = (cur, cur + len(idx))
cur += len(idx)
return cls(keep, groups, name)
@classmethod
def preset(cls, name):
if name not in PRESETS:
raise ValueError(f"unknown layout preset {name!r}; choices {list(PRESETS)}")
return cls.full() if name == "full" else cls.from_drop(PRESETS[name], name)
# ------------------------------------------------------------------- i/o
def save(self, data_dir):
with open(os.path.join(data_dir, "layout.json"), "w") as f:
json.dump({"name": self.name, "keep": self.keep,
"groups": {k: list(v) for k, v in self.groups.items()},
"n_kpts": self.n_kpts, "dim": self.dim}, f, indent=2)
@classmethod
def load(cls, data_dir):
p = os.path.join(data_dir, "layout.json")
if not os.path.exists(p):
return cls.full() # datasets packed before layouts existed
with open(p) as f:
d = json.load(f)
return cls(d["keep"], d["groups"], d.get("name", "custom"))
# ---------------------------------------------------------------- helpers
def metric_groups(self):
"""Group ranges for reporting, with 'hands' merging both hands."""
g = dict(self.groups)
out = {"all": (0, self.n_kpts), "body": g["body"], "face": g["face"],
"hands": (g["lhand"][0], g["rhand"][1])}
return out
def __repr__(self):
return (f"Layout({self.name}, {self.n_kpts} kpts, dim {self.dim}, "
f"groups {self.groups})")
|