f8037f7b85
The torchvision detector drop-in, headless pyrender/OpenGL shims, and the detector+smoke tests existed only on ultra's disk inside gitignored .engine/. Captured byte-identical copies into engine_patches/ (mirrors the HSMR subtree) and moved setup_hsmr.sh to the repo root (tracked, path-relative) with a --patch step that overlays them onto the clone. Acceptance PASSED: fresh shallow clone of HSMR in a temp dir + './setup_hsmr.sh --patch' + validated venv -> test_detector.py green (7/7 demo imgs, 15 patches), no hand edits. A re-clone can no longer silently drop the ARM fixes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
76 lines
3.4 KiB
Python
76 lines
3.4 KiB
Python
"""
|
|
Person detector — torchvision drop-in replacement for HSMR's detectron2/ViTDet.
|
|
|
|
Why: detectron2 + chumpy won't build from source on Apple Silicon (the Lane B blocker),
|
|
and this module's `import detectron2` was a load-time landmine for everything that touches
|
|
`lib.kits.hsmr_demo` (smoke_test.py, the future pose_engine.py). torchvision ships a
|
|
COCO-pretrained Faster R-CNN with ungated weights and no build step.
|
|
|
|
Same public contract as the original, so run_demo.py / smoke_test.py / pose_engine.py
|
|
import it unchanged:
|
|
|
|
build_detector(...) -> callable(raw_imgs) -> (dets, downsample_ratios)
|
|
|
|
dets[i] : dict of CPU tensors (empty tensors if no person) —
|
|
pred_classes : all 0 (== hsmr_demo CLASS_HUMAN_ID)
|
|
scores : float
|
|
pred_boxes : (N, 4) left-upper-right-bottom pixels
|
|
downsample_ratios : one float per image; boxes are in native-resolution pixels -> 1.0
|
|
|
|
Contract is checked by test_detector.py (runs today, no SKEL needed).
|
|
|
|
# ponytail: ~35 lines replaces the whole detectron2 stack. The ViT-H HMR forward pass is
|
|
# the runtime cost, not detection, so we don't chase detector throughput here.
|
|
"""
|
|
import numpy as np
|
|
import torch
|
|
from tqdm import tqdm
|
|
from torchvision.models.detection import (
|
|
fasterrcnn_resnet50_fpn_v2,
|
|
FasterRCNN_ResNet50_FPN_V2_Weights,
|
|
)
|
|
|
|
_COCO_PERSON = 1 # torchvision COCO label id for "person" (0 == background)
|
|
|
|
|
|
class _PersonDetector:
|
|
"""Callable over a list of RGB HxWx3 images (what HSMR's load_inputs yields).
|
|
uint8 or float, [0,255] or [0,1] are all handled."""
|
|
|
|
def __init__(self, device='cpu', score_thresh=0.25):
|
|
# score_thresh 0.25 mirrors the old detectron2 test_score_thresh; hsmr_demo's
|
|
# _img_det2patches re-filters at 0.5, so this only widens recall a little.
|
|
weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
|
|
self.model = fasterrcnn_resnet50_fpn_v2(
|
|
weights=weights, box_score_thresh=score_thresh,
|
|
).eval().to(device)
|
|
self.device = device
|
|
|
|
@torch.no_grad()
|
|
def __call__(self, raw_imgs):
|
|
dets, ratios = [], []
|
|
for img in tqdm(raw_imgs, desc='Detecting'):
|
|
t = torch.as_tensor(np.ascontiguousarray(img), device=self.device).float()
|
|
if float(t.max()) > 1.5: # [0,255] -> [0,1]
|
|
t = t / 255.0
|
|
t = t.permute(2, 0, 1) # HWC -> CHW
|
|
out = self.model([t])[0]
|
|
keep = out['labels'] == _COCO_PERSON
|
|
dets.append({
|
|
'pred_classes': torch.zeros(int(keep.sum()), dtype=torch.long), # human -> 0
|
|
'scores' : out['scores'][keep].cpu(),
|
|
'pred_boxes' : out['boxes'][keep].cpu(), # xyxy == left-upper-right-bottom
|
|
})
|
|
ratios.append(1.0) # detected at native resolution
|
|
return dets, ratios
|
|
|
|
|
|
def build_detector(batch_size=1, max_img_size=512, device='cpu'):
|
|
"""Signature-compatible with the old detectron2 builder. batch_size / max_img_size are
|
|
accepted for parity but unused — detection runs per-frame at native resolution.
|
|
|
|
# ponytail: detector pinned to CPU regardless of `device`. It's cheap next to the
|
|
# ViT-H HMR pass, and CPU sidesteps the MPS coverage gaps in RoIAlign/NMS. Flip to
|
|
# `device` only if per-frame detection ever dominates on long videos."""
|
|
return _PersonDetector(device='cpu')
|