| from __future__ import annotations | |
| from typing import Sequence | |
| import numpy as np | |
| def get_largest_bbox_indices(bboxes: Sequence[Sequence[float]], num_bboxes: int = 2) -> list[int]: | |
| """Return indices of the largest bboxes by area (normalized coords are fine).""" | |
| if num_bboxes <= 0: | |
| return [] | |
| bboxes_with_area: list[tuple[int, float]] = [] | |
| for i, bbox in enumerate(bboxes): | |
| try: | |
| x1, y1, x2, y2 = bbox | |
| except Exception: | |
| continue | |
| area = max(0.0, float(x2) - float(x1)) * max(0.0, float(y2) - float(y1)) | |
| bboxes_with_area.append((i, area)) | |
| bboxes_with_area.sort(key=lambda x: x[1], reverse=True) | |
| return [idx for idx, _area in bboxes_with_area[: min(num_bboxes, len(bboxes_with_area))]] | |
| def change_poses_to_limit_num(poses, bboxes, num_bboxes: int = 2): | |
| """Trim DWpose outputs to the top-N bboxes per frame (SCAIL-Pose multi-human helper).""" | |
| bboxes = list(bboxes) | |
| for frame_idx, (pose, bbox_list) in enumerate(zip(poses, bboxes)): | |
| if not bbox_list: | |
| continue | |
| largest_indices = get_largest_bbox_indices(bbox_list, num_bboxes) | |
| if not largest_indices: | |
| continue | |
| bodies = pose.get("bodies", {}) | |
| if "candidate" in bodies: | |
| bodies["candidate"] = bodies["candidate"][largest_indices] | |
| if "subset" in bodies: | |
| bodies["subset"] = bodies["subset"][largest_indices] | |
| pose["bodies"] = bodies | |
| faces = pose.get("faces", None) | |
| if isinstance(faces, np.ndarray): | |
| pose["faces"] = faces[largest_indices] | |
| elif isinstance(faces, list): | |
| pose["faces"] = [faces[i] for i in largest_indices if i < len(faces)] | |
| hands = pose.get("hands", None) | |
| hand_indices = [j for i in largest_indices for j in (2 * i, 2 * i + 1)] | |
| if isinstance(hands, np.ndarray): | |
| pose["hands"] = hands[hand_indices] | |
| elif isinstance(hands, list): | |
| pose["hands"] = [hands[i] for i in hand_indices if i < len(hands)] | |
| bboxes[frame_idx] = [bbox_list[i] for i in largest_indices if i < len(bbox_list)] | |
| return poses, bboxes | |