| """ |
| Pose engine: deteksi orang (YOLOX, dataset HumanArt) + estimasi pose whole-body |
| 133 titik (RTMW / DWPose, keluarga RTMPose) dengan ONNX Runtime, lalu |
| rendering gaya OpenPose/ControlNet (badan 18 titik + tangan 2x21 + wajah 68). |
| |
| Fokus utama: beberapa orang dalam satu foto yang saling bertumpuk / berhimpitan. |
| Strategi: |
| 1. Deteksi dengan ambang skor sangat rendah dan NMS longgar supaya kotak |
| orang yang tertutup sebagian tidak ikut terhapus. |
| 2. Estimasi pose untuk sejumlah kandidat kotak teratas. |
| 3. Hapus duplikat di tingkat POSE (OKS), bukan hanya di tingkat kotak. |
| 4. Hapus keypoint "bajakan" (dua orang berbagi titik sendi yang persis sama). |
| 5. Pilih orang terbaik; bila jumlah yang lolos ambang ketat kurang dari target, |
| ambil kandidat berikutnya (relaksasi otomatis). |
| """ |
| from __future__ import annotations |
|
|
| import colorsys |
| import math |
| import os |
| import shutil |
| import threading |
| import zipfile |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Dict, List, Optional, Tuple |
|
|
| import cv2 |
| import numpy as np |
|
|
| |
| |
| |
| HF_REPO = "Tau-J/RTMPose" |
| MODEL_DIR = Path(os.environ.get("POSE_MODEL_DIR", Path.home() / ".cache" / "pose-skeleton-models")) |
|
|
| |
| DETECTORS: Dict[str, Tuple[str, Tuple[int, int]]] = { |
| "YOLOX-l (akurasi tinggi, disarankan)": ( |
| "rtmposev1/onnx_sdk/yolox_l_8xb8-300e_humanart-ce1d7a62.zip", (640, 640)), |
| "YOLOX-m (lebih cepat)": ( |
| "rtmposev1/onnx_sdk/yolox_m_8xb8-300e_humanart-c2c7a14a.zip", (640, 640)), |
| "YOLOX-x (terbesar, paling lambat)": ( |
| "rtmposev1/onnx_sdk/yolox_x_8xb8-300e_humanart-a39d44ed.zip", (640, 640)), |
| } |
|
|
| |
| |
| |
| POSE_MODELS: Dict[str, Tuple[str, Tuple[int, int], float]] = { |
| "RTMW-l 384x288 (disarankan)": ( |
| "rtmw/onnx_sdk/rtmw-dw-x-l_simcc-cocktail14_270e-384x288_20231122.zip", (288, 384), 0.1), |
| "RTMW-x 384x288 (terbesar)": ( |
| "rtmw/onnx_sdk/rtmw-x_simcc-cocktail13_pt-ucoco_270e-384x288-0949e3a9_20230925.zip", (288, 384), 1.0), |
| "DWPose-l 384x288 (klasik, sama seperti ComfyUI)": ( |
| "rtmposev1/onnx_sdk/rtmpose-l_simcc-ucoco_dw-ucoco_270e-384x288-2438fd99_20230728.zip", (288, 384), 1.0), |
| "RTMW-m 256x192 (paling cepat)": ( |
| "rtmw/onnx_sdk/rtmw-dw-l-m_simcc-cocktail14_270e-256x192_20231122.zip", (192, 256), 1.0), |
| } |
|
|
| DEFAULT_DETECTOR = "YOLOX-l (akurasi tinggi, disarankan)" |
| DEFAULT_POSE = "RTMW-l 384x288 (disarankan)" |
|
|
| |
| RELAXED_DET_THR = 0.05 |
|
|
| |
| COCO_SIGMAS = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, |
| 1.07, 1.07, .87, .87, .89, .89], dtype=np.float32) / 10.0 |
|
|
| _LOCK = threading.Lock() |
| _SESSIONS: Dict[str, "OnnxModel"] = {} |
|
|
|
|
| |
| |
| |
| def ensure_onnx(zip_rel_path: str) -> str: |
| """Mengunduh zip SDK dari HF (sekali) dan mengekstrak file .onnx di dalamnya.""" |
| MODEL_DIR.mkdir(parents=True, exist_ok=True) |
| stem = Path(zip_rel_path).name[:-4] |
| onnx_path = MODEL_DIR / f"{stem}.onnx" |
| if onnx_path.exists() and onnx_path.stat().st_size > 0: |
| return str(onnx_path) |
|
|
| from huggingface_hub import hf_hub_download |
|
|
| zip_path = hf_hub_download(repo_id=HF_REPO, filename=zip_rel_path) |
| tmp_path = MODEL_DIR / f"{stem}.onnx.part" |
| with zipfile.ZipFile(zip_path) as zf: |
| members = [m for m in zf.namelist() if m.lower().endswith(".onnx")] |
| if not members: |
| raise RuntimeError(f"Tidak ada file .onnx di dalam {zip_rel_path}") |
| |
| members.sort(key=lambda m: (not m.endswith("end2end.onnx"), len(m))) |
| with zf.open(members[0]) as src, open(tmp_path, "wb") as dst: |
| shutil.copyfileobj(src, dst, length=16 * 1024 * 1024) |
| os.replace(tmp_path, onnx_path) |
| return str(onnx_path) |
|
|
|
|
| |
| |
| |
| class OnnxModel: |
| def __init__(self, onnx_path: str): |
| import onnxruntime as ort |
|
|
| so = ort.SessionOptions() |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| providers = ["CPUExecutionProvider"] |
| avail = ort.get_available_providers() |
| if "CUDAExecutionProvider" in avail: |
| providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] |
| self.session = ort.InferenceSession(onnx_path, sess_options=so, providers=providers) |
| self.input_name = self.session.get_inputs()[0].name |
| self.output_names = [o.name for o in self.session.get_outputs()] |
|
|
| def run(self, chw_float: np.ndarray) -> List[np.ndarray]: |
| inp = np.ascontiguousarray(chw_float[None].astype(np.float32)) |
| return self.session.run(self.output_names, {self.input_name: inp}) |
|
|
|
|
| def _get_model(key: str, zip_rel_path: str) -> OnnxModel: |
| with _LOCK: |
| m = _SESSIONS.get(key) |
| if m is None: |
| m = OnnxModel(ensure_onnx(zip_rel_path)) |
| _SESSIONS[key] = m |
| return m |
|
|
|
|
| def get_detector(name: str) -> Tuple[OnnxModel, Tuple[int, int]]: |
| zip_rel, size = DETECTORS[name] |
| return _get_model("det:" + zip_rel, zip_rel), size |
|
|
|
|
| def get_pose_model(name: str) -> Tuple[OnnxModel, Tuple[int, int], float]: |
| zip_rel, size, score_scale = POSE_MODELS[name] |
| return _get_model("pose:" + zip_rel, zip_rel), size, score_scale |
|
|
|
|
| def preload(det_name: str = DEFAULT_DETECTOR, pose_name: str = DEFAULT_POSE) -> None: |
| get_detector(det_name) |
| get_pose_model(pose_name) |
|
|
|
|
| |
| |
| |
| def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> List[int]: |
| x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] |
| areas = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None) |
| order = scores.argsort()[::-1] |
| keep: List[int] = [] |
| while order.size > 0: |
| i = int(order[0]) |
| keep.append(i) |
| if order.size == 1: |
| break |
| rest = order[1:] |
| xx1 = np.maximum(x1[i], x1[rest]) |
| yy1 = np.maximum(y1[i], y1[rest]) |
| xx2 = np.minimum(x2[i], x2[rest]) |
| yy2 = np.minimum(y2[i], y2[rest]) |
| inter = np.clip(xx2 - xx1, 0, None) * np.clip(yy2 - yy1, 0, None) |
| iou = inter / (areas[i] + areas[rest] - inter + 1e-9) |
| order = rest[iou <= iou_thr] |
| return keep |
|
|
|
|
| def box_iou(a: np.ndarray, b: np.ndarray) -> float: |
| xx1, yy1 = max(a[0], b[0]), max(a[1], b[1]) |
| xx2, yy2 = min(a[2], b[2]), min(a[3], b[3]) |
| inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1) |
| ua = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter |
| return float(inter / (ua + 1e-9)) |
|
|
|
|
| def _yolox_preprocess(img_bgr: np.ndarray, input_size: Tuple[int, int]) -> Tuple[np.ndarray, float]: |
| """Letterbox ke (w, h) input model, padding 114 di kanan/bawah. Mengembalikan CHW float32 dan rasio.""" |
| in_w, in_h = input_size |
| padded = np.full((in_h, in_w, 3), 114, dtype=np.uint8) |
| r = min(in_h / img_bgr.shape[0], in_w / img_bgr.shape[1]) |
| new_w, new_h = max(1, int(img_bgr.shape[1] * r)), max(1, int(img_bgr.shape[0] * r)) |
| resized = cv2.resize(img_bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR) |
| padded[:new_h, :new_w] = resized |
| chw = np.ascontiguousarray(padded.transpose(2, 0, 1), dtype=np.float32) |
| return chw, r |
|
|
|
|
| def _yolox_decode(outputs: np.ndarray, input_size: Tuple[int, int]) -> np.ndarray: |
| """Dekode grid/stride YOLOX (keluaran mentah tanpa NMS) -> [cx, cy, w, h, obj, cls...].""" |
| in_w, in_h = input_size |
| grids, strides_all = [], [] |
| for stride in (8, 16, 32): |
| hsize, wsize = in_h // stride, in_w // stride |
| xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) |
| grid = np.stack((xv, yv), 2).reshape(1, -1, 2) |
| grids.append(grid) |
| strides_all.append(np.full((*grid.shape[:2], 1), stride)) |
| grids = np.concatenate(grids, 1).astype(np.float32) |
| strides_all = np.concatenate(strides_all, 1).astype(np.float32) |
| out = outputs.astype(np.float32).copy() |
| out[..., :2] = (out[..., :2] + grids) * strides_all |
| out[..., 2:4] = np.exp(out[..., 2:4]) * strides_all |
| return out |
|
|
|
|
| def detect_people(img_bgr: np.ndarray, det_name: str, score_thr: float, nms_thr: float) -> np.ndarray: |
| """Mengembalikan array (N, 5): x1, y1, x2, y2, skor — koordinat gambar asli, urut skor menurun.""" |
| model, input_size = get_detector(det_name) |
| chw, ratio = _yolox_preprocess(img_bgr, input_size) |
| outputs = model.run(chw) |
| out0 = np.asarray(outputs[0]) |
| if out0.ndim == 2: |
| out0 = out0[None] |
|
|
| if out0.shape[-1] == 5: |
| |
| dets = out0[0] |
| boxes = dets[:, :4] / ratio |
| scores = dets[:, 4] |
| else: |
| pred = _yolox_decode(out0, input_size)[0] |
| boxes_c = pred[:, :4] |
| cls_scores = pred[:, 5:] if pred.shape[1] > 5 else np.ones((pred.shape[0], 1), np.float32) |
| |
| scores = pred[:, 4] * cls_scores[:, 0] |
| boxes = np.empty_like(boxes_c) |
| boxes[:, 0] = boxes_c[:, 0] - boxes_c[:, 2] / 2.0 |
| boxes[:, 1] = boxes_c[:, 1] - boxes_c[:, 3] / 2.0 |
| boxes[:, 2] = boxes_c[:, 0] + boxes_c[:, 2] / 2.0 |
| boxes[:, 3] = boxes_c[:, 1] + boxes_c[:, 3] / 2.0 |
| boxes = boxes / ratio |
|
|
| valid = scores > score_thr |
| boxes, scores = boxes[valid], scores[valid] |
| if boxes.shape[0] == 0: |
| return np.zeros((0, 5), np.float32) |
| keep = _nms(boxes, scores, nms_thr) |
| boxes, scores = boxes[keep], scores[keep] |
| h, w = img_bgr.shape[:2] |
| boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, w - 1) |
| boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, h - 1) |
| dets = np.concatenate([boxes, scores[:, None]], axis=1).astype(np.float32) |
| ok = (dets[:, 2] - dets[:, 0] > 2) & (dets[:, 3] - dets[:, 1] > 2) |
| dets = dets[ok] |
| return dets[np.argsort(-dets[:, 4])] |
|
|
|
|
| |
| |
| |
| MEAN = np.array([123.675, 116.28, 103.53], dtype=np.float32) |
| STD = np.array([58.395, 57.12, 57.375], dtype=np.float32) |
|
|
|
|
| def _bbox_xyxy2cs(bbox: np.ndarray, padding: float = 1.25) -> Tuple[np.ndarray, np.ndarray]: |
| x1, y1, x2, y2 = [float(v) for v in bbox[:4]] |
| center = np.array([(x1 + x2) * 0.5, (y1 + y2) * 0.5], dtype=np.float32) |
| scale = np.array([(x2 - x1) * padding, (y2 - y1) * padding], dtype=np.float32) |
| return center, scale |
|
|
|
|
| def _fix_aspect_ratio(scale: np.ndarray, aspect_ratio: float) -> np.ndarray: |
| w, h = float(scale[0]), float(scale[1]) |
| if w > h * aspect_ratio: |
| return np.array([w, w / aspect_ratio], dtype=np.float32) |
| return np.array([h * aspect_ratio, h], dtype=np.float32) |
|
|
|
|
| def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray: |
| direction = a - b |
| return b + np.array([-direction[1], direction[0]], dtype=np.float32) |
|
|
|
|
| def _get_warp_matrix(center: np.ndarray, scale: np.ndarray, output_size: Tuple[int, int]) -> np.ndarray: |
| src_w = float(scale[0]) |
| dst_w, dst_h = output_size |
| src_dir = np.array([0.0, src_w * -0.5], dtype=np.float32) |
| dst_dir = np.array([0.0, dst_w * -0.5], dtype=np.float32) |
| src = np.zeros((3, 2), dtype=np.float32) |
| src[0] = center |
| src[1] = center + src_dir |
| src[2] = _get_3rd_point(src[0], src[1]) |
| dst = np.zeros((3, 2), dtype=np.float32) |
| dst[0] = [dst_w * 0.5, dst_h * 0.5] |
| dst[1] = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32) + dst_dir |
| dst[2] = _get_3rd_point(dst[0], dst[1]) |
| return cv2.getAffineTransform(src, dst) |
|
|
|
|
| def _simcc_decode(simcc_x: np.ndarray, simcc_y: np.ndarray, split_ratio: float = 2.0): |
| n, k, _ = simcc_x.shape |
| sx = simcc_x.reshape(n * k, -1) |
| sy = simcc_y.reshape(n * k, -1) |
| x_locs = sx.argmax(axis=1) |
| y_locs = sy.argmax(axis=1) |
| locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) |
| vx, vy = sx.max(axis=1), sy.max(axis=1) |
| vals = np.minimum(vx, vy) |
| locs[vals <= 0.0] = -1 |
| return locs.reshape(n, k, 2) / split_ratio, vals.reshape(n, k) |
|
|
|
|
| def _pick_simcc(model: "OnnxModel", outputs: List[np.ndarray], w: int, h: int) -> Tuple[np.ndarray, np.ndarray]: |
| """Memilih tensor simcc_x / simcc_y berdasarkan nama keluaran, atau bentuknya bila nama tidak baku.""" |
| names = [n.lower() for n in getattr(model, "output_names", [])] |
| if "simcc_x" in names and "simcc_y" in names: |
| return np.asarray(outputs[names.index("simcc_x")]), np.asarray(outputs[names.index("simcc_y")]) |
| a, b = np.asarray(outputs[0]), np.asarray(outputs[1]) |
| |
| if a.shape[-1] != b.shape[-1] and abs(a.shape[-1] - 2 * w) > abs(b.shape[-1] - 2 * w): |
| a, b = b, a |
| return a, b |
|
|
|
|
| def estimate_poses(img_rgb: np.ndarray, boxes_xyxy: np.ndarray, pose_name: str) -> Tuple[np.ndarray, np.ndarray]: |
| """Mengembalikan keypoints (N, 133, 2) pada koordinat gambar dan skor (N, 133), skor sudah diskalakan ~0-1.""" |
| model, input_size, score_scale = get_pose_model(pose_name) |
| w, h = input_size |
| kpts_all, scores_all = [], [] |
| for bbox in boxes_xyxy: |
| center, scale = _bbox_xyxy2cs(np.asarray(bbox, dtype=np.float32), padding=1.25) |
| scale = _fix_aspect_ratio(scale, w / h) |
| warp = _get_warp_matrix(center, scale, (w, h)) |
| crop = cv2.warpAffine(img_rgb, warp, (int(w), int(h)), flags=cv2.INTER_LINEAR) |
| crop = (crop.astype(np.float32) - MEAN) / STD |
| outputs = model.run(crop.transpose(2, 0, 1)) |
| simcc_x, simcc_y = _pick_simcc(model, outputs, w, h) |
| locs, vals = _simcc_decode(simcc_x, simcc_y) |
| kp = locs[0] / np.array([w, h], dtype=np.float32) * scale + center - scale / 2.0 |
| kpts_all.append(kp.astype(np.float32)) |
| scores_all.append((vals[0] * score_scale).astype(np.float32)) |
| if not kpts_all: |
| return np.zeros((0, 133, 2), np.float32), np.zeros((0, 133), np.float32) |
| return np.stack(kpts_all), np.stack(scores_all) |
|
|
|
|
| |
| |
| |
| @dataclass(eq=False) |
| class Person: |
| box: np.ndarray |
| det_score: float |
| kpts: np.ndarray |
| scores: np.ndarray |
| relaxed: bool = False |
| removed_joints: int = 0 |
|
|
| @property |
| def area(self) -> float: |
| return float(max(1.0, (self.box[2] - self.box[0]) * (self.box[3] - self.box[1]))) |
|
|
| @property |
| def center(self) -> np.ndarray: |
| return np.array([(self.box[0] + self.box[2]) / 2.0, (self.box[1] + self.box[3]) / 2.0], np.float32) |
|
|
| def n_visible_body(self, thr: float) -> int: |
| return int((self.scores[:17] > thr).sum()) |
|
|
| def mean_body_score(self, thr: float) -> float: |
| s = self.scores[:17] |
| m = s > thr |
| return float(s[m].mean()) if m.any() else 0.0 |
|
|
|
|
| def oks_similarity(a: Person, b: Person, thr: float) -> Tuple[float, int]: |
| """OKS atas sendi badan yang terlihat pada keduanya. Area = perkiraan area segmentasi (0.53 x area kotak).""" |
| vis = (a.scores[:17] > thr) & (b.scores[:17] > thr) |
| n = int(vis.sum()) |
| if n < 3: |
| return 0.0, n |
| d2 = ((a.kpts[:17] - b.kpts[:17]) ** 2).sum(-1) |
| area = 0.53 * 0.5 * (a.area + b.area) |
| var = (2.0 * COCO_SIGMAS) ** 2 |
| e = d2 / (2.0 * var * (area + 1e-6)) |
| return float(np.exp(-e[vis]).mean()), n |
|
|
|
|
| def suppress_duplicate_persons(persons: List[Person], kpt_thr: float, oks_thr: float) -> List[Person]: |
| """NMS tingkat pose: instance dengan OKS tinggi terhadap instance berperingkat lebih tinggi = duplikat.""" |
| kept: List[Person] = [] |
| for cand in persons: |
| dup = False |
| for k in kept: |
| o, _ = oks_similarity(cand, k, kpt_thr) |
| if o > oks_thr: |
| dup = True |
| break |
| if not dup: |
| kept.append(cand) |
| return kept |
|
|
|
|
| def remove_shared_joints(persons: List[Person], kpt_thr: float, rel_tol: float = 0.05) -> None: |
| """Dua orang yang 'berbagi' titik sendi persis sama (indeks sama, jarak sangat dekat) => |
| titik itu dipertahankan untuk orang yang kotaknya paling 'memiliki' titik tersebut |
| (jarak ternormalisasi ke pusat kotak lebih kecil); bila seimbang, skor keypoint yang menentukan.""" |
| for i in range(len(persons)): |
| for j in range(i + 1, len(persons)): |
| a, b = persons[i], persons[j] |
| tol = max(3.0, rel_tol * math.sqrt(min(a.area, b.area))) |
| d = np.linalg.norm(a.kpts - b.kpts, axis=-1) |
| both = (a.scores > kpt_thr) & (b.scores > kpt_thr) & (d < tol) |
| if not both.any(): |
| continue |
| sa, sb = math.sqrt(a.area), math.sqrt(b.area) |
| for jidx in np.flatnonzero(both): |
| na = float(np.linalg.norm(a.kpts[jidx] - a.center)) / sa |
| nb = float(np.linalg.norm(b.kpts[jidx] - b.center)) / sb |
| if abs(na - nb) < 0.15: |
| loser = a if a.scores[jidx] < b.scores[jidx] else b |
| else: |
| loser = a if na > nb else b |
| loser.scores[jidx] = 0.0 |
| loser.removed_joints += 1 |
|
|
|
|
| |
| _LR_GROUPS = [((17, 20), (20, 23)), ((91, 112), (112, 133))] |
|
|
|
|
| def remove_lr_duplicates(p: Person, kpt_thr: float, rel_tol: float = 0.025) -> None: |
| """Artefak umum top-down: tangan (atau kaki) kiri dan kanan diprediksi di tempat yang persis sama |
| karena salah satunya tersembunyi. Jika kedua kelompok hampir berimpit, kelompok dengan skor rata-rata |
| lebih rendah dianggap halusinasi dan dimatikan.""" |
| tol = max(4.0, rel_tol * math.sqrt(p.area)) |
| for (l0, l1), (r0, r1) in _LR_GROUPS: |
| ls, rs = p.scores[l0:l1], p.scores[r0:r1] |
| if not (ls > kpt_thr).any() or not (rs > kpt_thr).any(): |
| continue |
| d = np.linalg.norm(p.kpts[l0:l1] - p.kpts[r0:r1], axis=-1) |
| if float(d.mean()) < tol: |
| lo0, lo1 = (l0, l1) if float(ls.mean()) < float(rs.mean()) else (r0, r1) |
| n_on = int((p.scores[lo0:lo1] > kpt_thr).sum()) |
| p.scores[lo0:lo1] = 0.0 |
| p.removed_joints += n_on |
|
|
|
|
| @dataclass |
| class PipelineResult: |
| persons: List[Person] |
| n_candidates: int |
| n_after_dedup: int |
| notes: List[str] = field(default_factory=list) |
|
|
|
|
| def run_pipeline(img_rgb: np.ndarray, |
| det_name: str = DEFAULT_DETECTOR, |
| pose_name: str = DEFAULT_POSE, |
| n_target: int = 2, |
| det_thr: float = 0.3, |
| nms_thr: float = 0.75, |
| kpt_thr: float = 0.3, |
| oks_thr: float = 0.75, |
| dedup_joints: bool = True, |
| min_body_visible: int = 3, |
| max_candidates: Optional[int] = None) -> PipelineResult: |
| notes: List[str] = [] |
| h, w = img_rgb.shape[:2] |
| img_bgr = np.ascontiguousarray(img_rgb[..., ::-1]) |
|
|
| |
| dets = detect_people(img_bgr, det_name, score_thr=min(RELAXED_DET_THR, det_thr), nms_thr=nms_thr) |
| if dets.shape[0] == 0: |
| notes.append("Detektor tidak menemukan orang; seluruh foto dipakai sebagai satu kotak.") |
| dets = np.array([[0, 0, w - 1, h - 1, 1.0]], dtype=np.float32) |
| n_candidates = int(dets.shape[0]) |
|
|
| |
| pre: List[np.ndarray] = [] |
| for d in dets: |
| if all(box_iou(d, k) < 0.9 for k in pre): |
| pre.append(d) |
| dets = np.stack(pre) |
|
|
| if max_candidates is None: |
| max_candidates = max(6, 2 * n_target + 4) if n_target > 0 else 12 |
| dets = dets[:max_candidates] |
|
|
| |
| kpts, scores = estimate_poses(img_rgb, dets[:, :4], pose_name) |
| persons = [Person(box=dets[i, :4].copy(), det_score=float(dets[i, 4]), |
| kpts=kpts[i].copy(), scores=scores[i].copy()) for i in range(len(dets))] |
|
|
| |
| persons = suppress_duplicate_persons(persons, kpt_thr, oks_thr) |
| n_after_dedup = len(persons) |
|
|
| |
| |
| if dedup_joints: |
| for p in persons: |
| remove_lr_duplicates(p, kpt_thr) |
| if len(persons) > 1: |
| remove_shared_joints(persons, kpt_thr) |
|
|
| |
| persons = [p for p in persons if p.n_visible_body(kpt_thr) >= min_body_visible] |
|
|
| |
| strong = [p for p in persons if p.det_score >= det_thr] |
| if n_target > 0: |
| final = strong[:n_target] |
| if len(final) < n_target: |
| extra = [p for p in persons if p not in final][: n_target - len(final)] |
| for p in extra: |
| p.relaxed = True |
| if extra: |
| notes.append( |
| f"{len(extra)} orang ditambahkan lewat relaksasi ambang (skor deteksi < {det_thr:.2f}) " |
| f"agar mencapai target {n_target} orang.") |
| final += extra |
| if len(final) < n_target: |
| notes.append(f"Hanya {len(final)} orang yang dapat ditemukan (target {n_target}).") |
| else: |
| final = strong |
|
|
| |
| final.sort(key=lambda p: float(p.center[0])) |
| return PipelineResult(persons=final, n_candidates=n_candidates, n_after_dedup=n_after_dedup, notes=notes) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| _OP_FROM_COCO = {0: 0, 2: 6, 3: 8, 4: 10, 5: 5, 6: 7, 7: 9, 8: 12, 9: 14, 10: 16, |
| 11: 11, 12: 13, 13: 15, 14: 2, 15: 1, 16: 4, 17: 3} |
|
|
|
|
| @dataclass(eq=False) |
| class OpenPosePerson: |
| body: np.ndarray |
| body_s: np.ndarray |
| face: np.ndarray |
| face_s: np.ndarray |
| lhand: np.ndarray |
| lhand_s: np.ndarray |
| rhand: np.ndarray |
| rhand_s: np.ndarray |
| feet: np.ndarray |
| feet_s: np.ndarray |
| box: np.ndarray |
| det_score: float |
| relaxed: bool |
| removed_joints: int |
|
|
|
|
| def to_openpose(p: Person, kpt_thr: float, scale: float = 1.0) -> OpenPosePerson: |
| k = p.kpts * scale |
| s = p.scores |
| body = np.zeros((18, 2), np.float32) |
| body_s = np.zeros(18, np.float32) |
| for op_i, coco_i in _OP_FROM_COCO.items(): |
| body[op_i] = k[coco_i] |
| body_s[op_i] = s[coco_i] |
| |
| if s[5] > kpt_thr and s[6] > kpt_thr: |
| body[1] = (k[5] + k[6]) / 2.0 |
| body_s[1] = min(s[5], s[6]) |
| return OpenPosePerson( |
| body=body, body_s=body_s, |
| face=k[23:91].copy(), face_s=s[23:91].copy(), |
| lhand=k[91:112].copy(), lhand_s=s[91:112].copy(), |
| rhand=k[112:133].copy(), rhand_s=s[112:133].copy(), |
| feet=k[17:23].copy(), feet_s=s[17:23].copy(), |
| box=p.box * scale, det_score=p.det_score, relaxed=p.relaxed, |
| removed_joints=p.removed_joints, |
| ) |
|
|
|
|
| def _flat(kp: np.ndarray, sc: np.ndarray, thr: float) -> List[float]: |
| out: List[float] = [] |
| for (x, y), c in zip(kp, sc): |
| if c > thr: |
| out += [round(float(x), 2), round(float(y), 2), round(float(c), 3)] |
| else: |
| out += [0.0, 0.0, 0.0] |
| return out |
|
|
|
|
| def openpose_json(people: List[OpenPosePerson], width: int, height: int, kpt_thr: float) -> List[dict]: |
| """Format JSON OpenPose (kompatibel dengan POSE_KEYPOINT ComfyUI / OpenPose Editor): list of frames.""" |
| plist = [] |
| for idx, p in enumerate(people): |
| |
| face70 = np.concatenate([p.face, p.body[[14, 15]]], axis=0) |
| face70_s = np.concatenate([p.face_s, p.body_s[[14, 15]]], axis=0) |
| plist.append({ |
| "person_id": [idx], |
| "pose_keypoints_2d": _flat(p.body, p.body_s, kpt_thr), |
| "face_keypoints_2d": _flat(face70, face70_s, kpt_thr), |
| "hand_left_keypoints_2d": _flat(p.lhand, p.lhand_s, kpt_thr), |
| "hand_right_keypoints_2d": _flat(p.rhand, p.rhand_s, kpt_thr), |
| "foot_keypoints_2d": _flat(p.feet, p.feet_s, kpt_thr), |
| "bbox_xyxy": [round(float(v), 1) for v in p.box], |
| "detection_score": round(float(p.det_score), 3), |
| }) |
| return [{"version": 1.3, "people": plist, "canvas_width": int(width), "canvas_height": int(height)}] |
|
|
|
|
| |
| |
| |
| LIMB_SEQ = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [10, 11], |
| [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], [1, 16], [16, 18]] |
| BODY_COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], |
| [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], |
| [0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] |
| HAND_EDGES = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], |
| [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] |
|
|
|
|
| def _pt(v: np.ndarray) -> Tuple[int, int]: |
| return int(round(float(v[0]))), int(round(float(v[1]))) |
|
|
|
|
| def draw_bodies(canvas: np.ndarray, people: List[OpenPosePerson], kpt_thr: float, s: float) -> np.ndarray: |
| stickwidth = max(1, int(round(4 * s))) |
| radius = max(1, int(round(4 * s))) |
| for p in people: |
| for i, (a, b) in enumerate(LIMB_SEQ): |
| ia, ib = a - 1, b - 1 |
| if p.body_s[ia] <= kpt_thr or p.body_s[ib] <= kpt_thr: |
| continue |
| (x1, y1), (x2, y2) = p.body[ia], p.body[ib] |
| mx, my = (x1 + x2) / 2.0, (y1 + y2) / 2.0 |
| length = math.hypot(x1 - x2, y1 - y2) |
| angle = math.degrees(math.atan2(y1 - y2, x1 - x2)) |
| poly = cv2.ellipse2Poly((int(mx), int(my)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) |
| cv2.fillConvexPoly(canvas, poly, BODY_COLORS[i]) |
| canvas = (canvas.astype(np.float32) * 0.6).astype(np.uint8) |
| for p in people: |
| for i in range(18): |
| if p.body_s[i] > kpt_thr: |
| cv2.circle(canvas, _pt(p.body[i]), radius, BODY_COLORS[i], thickness=-1) |
| return canvas |
|
|
|
|
| def draw_hands(canvas: np.ndarray, people: List[OpenPosePerson], kpt_thr: float, s: float) -> np.ndarray: |
| thickness = max(1, int(round(2 * s))) |
| radius = max(1, int(round(3 * s))) |
| n_edges = len(HAND_EDGES) |
| for p in people: |
| for hand, hand_s in ((p.lhand, p.lhand_s), (p.rhand, p.rhand_s)): |
| for ie, (a, b) in enumerate(HAND_EDGES): |
| if hand_s[a] <= kpt_thr or hand_s[b] <= kpt_thr: |
| continue |
| r, g, bl = colorsys.hsv_to_rgb(ie / float(n_edges), 1.0, 1.0) |
| color = (int(r * 255), int(g * 255), int(bl * 255)) |
| cv2.line(canvas, _pt(hand[a]), _pt(hand[b]), color, thickness=thickness) |
| for i in range(21): |
| if hand_s[i] > kpt_thr: |
| cv2.circle(canvas, _pt(hand[i]), radius, (0, 0, 255), thickness=-1) |
| return canvas |
|
|
|
|
| def draw_faces(canvas: np.ndarray, people: List[OpenPosePerson], kpt_thr: float, s: float) -> np.ndarray: |
| radius = max(1, int(round(2 * s))) |
| for p in people: |
| for i in range(p.face.shape[0]): |
| if p.face_s[i] > kpt_thr: |
| cv2.circle(canvas, _pt(p.face[i]), radius, (255, 255, 255), thickness=-1) |
| return canvas |
|
|
|
|
| def render_pose_map(people: List[OpenPosePerson], width: int, height: int, kpt_thr: float, |
| line_scale: float = 1.0, draw_body: bool = True, draw_hand: bool = True, |
| draw_face: bool = True, background: Optional[np.ndarray] = None) -> np.ndarray: |
| canvas = np.zeros((height, width, 3), np.uint8) if background is None else background.copy() |
| s = line_scale * max(1.0, max(width, height) / 512.0) |
| if draw_body: |
| canvas = draw_bodies(canvas, people, kpt_thr, s) |
| if draw_hand: |
| canvas = draw_hands(canvas, people, kpt_thr, s) |
| if draw_face: |
| canvas = draw_faces(canvas, people, kpt_thr, s) |
| return canvas |
|
|
|
|
| def render_overlay(img_rgb: np.ndarray, pose_map: np.ndarray, people: List[OpenPosePerson], |
| alpha_bg: float = 0.45) -> np.ndarray: |
| """Foto digelapkan + skeleton ditimpa + kotak & nomor orang untuk pengecekan.""" |
| base = (img_rgb.astype(np.float32) * alpha_bg).astype(np.uint8) |
| mask = pose_map.any(axis=-1) |
| base[mask] = pose_map[mask] |
| s = max(1.0, max(base.shape[:2]) / 512.0) |
| for idx, p in enumerate(people): |
| x1, y1, x2, y2 = [int(round(float(v))) for v in p.box] |
| color = (255, 255, 255) if not p.relaxed else (255, 200, 0) |
| cv2.rectangle(base, (x1, y1), (x2, y2), color, max(1, int(round(s)))) |
| label = f"#{idx + 1} {p.det_score:.2f}" + (" (relaksasi)" if p.relaxed else "") |
| cv2.putText(base, label, (x1 + 3, max(12, y1 - 4)), cv2.FONT_HERSHEY_SIMPLEX, |
| 0.45 * s, color, max(1, int(round(s))), cv2.LINE_AA) |
| return base |
|
|
