Buckets:
| import contextlib | |
| import os | |
| import sys | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from .. import paths | |
| from ..hands import Hand | |
| from ..skeleton import MANO | |
| HAWOR_ROOT = paths.ROOT / "external" / "HaWoR" | |
| DETECTION_THRESHOLD = 0.2 | |
| class HandTrack: | |
| side: str | |
| keypoints: np.ndarray | |
| valid: np.ndarray | |
| root_orient: np.ndarray | |
| hand_pose: np.ndarray | |
| betas: np.ndarray | |
| trans: np.ndarray | |
| def to_hand(self, camera): | |
| rotation, position = camera[:, :3, :3], camera[:, None, :3, 3] | |
| world = np.einsum("tij,tkj->tki", rotation, self.keypoints) + position | |
| return Hand.from_keypoints(self.side, world, MANO, self.valid) | |
| FIELDS = ("keypoints", "valid", "root_orient", "hand_pose", "betas", "trans") | |
| def save_tracks(path, tracks): | |
| np.savez_compressed( | |
| path, | |
| **{ | |
| f"{side}/{name}": getattr(t, name) | |
| for side, t in tracks.items() | |
| for name in FIELDS | |
| }, | |
| ) | |
| def load_tracks(path): | |
| data = np.load(path) | |
| return { | |
| side: HandTrack(side, **{name: data[f"{side}/{name}"] for name in FIELDS}) | |
| for side in ("left", "right") | |
| } | |
| def _inside(root): | |
| previous = os.getcwd() | |
| os.chdir(root) | |
| sys.path.insert(0, str(root)) | |
| try: | |
| yield | |
| finally: | |
| sys.path.remove(str(root)) | |
| os.chdir(previous) | |
| def extract_frames(video_path, folder): | |
| folder.mkdir(parents=True, exist_ok=True) | |
| capture = cv2.VideoCapture(str(video_path)) | |
| files = [] | |
| while True: | |
| ok, frame = capture.read() | |
| if not ok: | |
| break | |
| files.append(str(folder / f"{len(files):04d}.jpg")) | |
| cv2.imwrite(files[-1], frame) | |
| capture.release() | |
| return files | |
| class HaWoR: | |
| def __init__(self, root=HAWOR_ROOT, device="cuda"): | |
| self.root = Path(root) | |
| self.device = device | |
| with _inside(self.root): | |
| from hawor.configs import get_config | |
| from lib.models.hawor import HAWOR | |
| config = get_config( | |
| str(self.root / "weights/hawor/model_config.yaml"), update_cachedir=True | |
| ) | |
| config.defrost() | |
| config.MODEL.BBOX_SHAPE = [192, 256] | |
| config.freeze() | |
| checkpoint = self.root / "weights/hawor/checkpoints/hawor.ckpt" | |
| self.model = ( | |
| HAWOR.load_from_checkpoint(str(checkpoint), strict=False, cfg=config) | |
| .to(device) | |
| .eval() | |
| ) | |
| def track(self, video_path, focal, work_dir): | |
| files = extract_frames(video_path, Path(work_dir) / "frames") | |
| height, width = cv2.imread(files[0]).shape[:2] | |
| with _inside(self.root): | |
| from lib.pipeline.tools import detect_track | |
| _, tracks = detect_track(files, thresh=DETECTION_THRESHOLD) | |
| by_side = self._split(tracks.item()) | |
| return { | |
| side: self._estimate( | |
| side, by_side[side], files, focal, (width / 2, height / 2) | |
| ) | |
| for side in ("left", "right") | |
| } | |
| def _split(tracks): | |
| by_side = {"left": [], "right": []} | |
| for track in tracks.values(): | |
| valid = np.array([t["det"] for t in track]) | |
| is_right = np.concatenate([t["det_handedness"] for t in track])[valid] | |
| by_side["left" if is_right.sum() / len(is_right) < 0.5 else "right"].extend( | |
| track | |
| ) | |
| return { | |
| side: sorted(track, key=lambda t: t["frame"]) | |
| for side, track in by_side.items() | |
| } | |
| def _estimate(self, side, track, files, focal, center): | |
| from hawor.utils.process import run_mano, run_mano_left | |
| from hawor.utils.rotation import rotation_matrix_to_angle_axis | |
| from lib.eval_utils.custom_utils import interpolate_bboxes | |
| from lib.pipeline.tools import parse_chunks | |
| count = len(files) | |
| out = HandTrack( | |
| side=side, | |
| keypoints=np.zeros((count, 21, 3)), | |
| valid=np.zeros(count, dtype=bool), | |
| root_orient=np.zeros((count, 3)), | |
| hand_pose=np.zeros((count, 45)), | |
| betas=np.zeros((count, 10)), | |
| trans=np.zeros((count, 3)), | |
| ) | |
| valid = np.array([t["det"] for t in track]) | |
| if valid.sum() < 2: | |
| return out | |
| boxes = np.concatenate([t["det_box"] for t in track]) | |
| filled = np.where(np.any(boxes != 0, axis=1))[0] | |
| first, last = filled[0], filled[-1] | |
| boxes[first : last + 1] = interpolate_bboxes(boxes[first : last + 1]) | |
| valid[first : last + 1] = True | |
| boxes = boxes[first : last + 1] | |
| frames = np.array([t["frame"] for t in track])[valid] | |
| flip = side == "left" | |
| run = run_mano_left if flip else run_mano | |
| for frame_chunk, box_chunk in zip(*parse_chunks(frames, boxes, min_len=1)): | |
| results = self.model.inference( | |
| np.array(files)[frame_chunk], | |
| box_chunk, | |
| img_focal=focal, | |
| img_center=list(center), | |
| do_flip=flip, | |
| ) | |
| root = rotation_matrix_to_angle_axis(results["pred_rotmat"][None, :, 0]) | |
| pose = rotation_matrix_to_angle_axis(results["pred_rotmat"][None, :, 1:]) | |
| if flip: | |
| root[..., 1:] *= -1 | |
| pose[..., 1:] *= -1 | |
| trans = results["pred_trans"][None, :, 0] | |
| betas = results["pred_shape"][None, :] | |
| joints = run(trans, root, pose, betas=betas)["joints"][0].cpu().numpy() | |
| out.keypoints[frame_chunk] = joints | |
| out.valid[frame_chunk] = True | |
| out.root_orient[frame_chunk] = root[0].numpy() | |
| out.hand_pose[frame_chunk] = pose[0].reshape(len(frame_chunk), 45).numpy() | |
| out.betas[frame_chunk] = betas[0].numpy() | |
| out.trans[frame_chunk] = trans[0].numpy() | |
| return out | |
Xet Storage Details
- Size:
- 6.06 kB
- Xet hash:
- d76d4bf59a8426a873d069085c443bebc75b85ff0ed898ae4bbddcc034213351
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.