File size: 3,416 Bytes
3b18ebd | 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 | # -*- coding: utf-8 -*-
"""即時抽點 —— 與組員的 extract_features.py 完全一致。
★ 這個檔案存在的唯一理由：訓練資料怎麼抽的，即時推論就要一模一樣抽。
節點順序、左右手判定、缺值填 NaN 的方式，任何一點不同都會讓模型失準。
所以這裡直接沿用 mp.solutions.holistic，而不是新版 Tasks API
（Tasks API 的 HandLandmarker 用 handedness 分類器判左右，
holistic 則是從 pose 推 ROI —— 兩者對左右手的認定會不一樣）。
因此本檔需要**舊版 mediapipe**（holistic 在 0.10.15 之後被移除）：
pip install "mediapipe==0.10.14"
建議跟你的 TSL 生成專案分開用不同的 venv，避免和 0.10.35 打架。
請跟組員確認他跑 extract_features.py 時用的是哪個版本，直接對齊那個。
"""
from __future__ import annotations
import numpy as np
import config as C
# 與 extract_features.py 相同：pose 0:12, 左手 12:33, 右手 33:54, 虛擬頸點 54
POSE_INDICES = C.POSE_MP_INDICES
NUM_NODES = C.NUM_POINTS
def _missing(count):
return [[np.nan, np.nan, np.nan] for _ in range(count)]
def extract_landmarks(results):
"""把 holistic 的結果轉成一格 (55, 3)。邏輯與 extract_features.py 相同。"""
frame_data = []
if results.pose_landmarks:
for index in POSE_INDICES:
lm = results.pose_landmarks.landmark[index]
frame_data.append([lm.x, lm.y, lm.z])
else:
frame_data.extend(_missing(len(POSE_INDICES)))
for landmarks in (results.left_hand_landmarks, results.right_hand_landmarks):
if landmarks:
frame_data.extend([lm.x, lm.y, lm.z] for lm in landmarks.landmark)
else:
frame_data.extend(_missing(21))
if results.pose_landmarks:
left = results.pose_landmarks.landmark[11]
right = results.pose_landmarks.landmark[12]
frame_data.append([
(left.x + right.x) / 2.0,
(left.y + right.y) / 2.0,
(left.z + right.z) / 2.0,
])
else:
frame_data.extend(_missing(1))
return frame_data
class HolisticExtractor:
"""包一層 context manager，供即時推論逐格呼叫。"""
def __init__(self, model_complexity: int = 1):
import mediapipe as mp
if not hasattr(mp.solutions, "holistic"):
raise RuntimeError(
"這個 mediapipe 版本沒有 solutions.holistic。\n"
'請安裝舊版：pip install "mediapipe==0.10.14"'
)
self._mp = mp
self.holistic = mp.solutions.holistic.Holistic(
static_image_mode=False,
model_complexity=model_complexity,
enable_segmentation=False,
refine_face_landmarks=False,
min_detection_confidence=0.5,
min_tracking_confidence=0.5,
)
def __call__(self, frame_bgr) -> np.ndarray:
"""輸入**未鏡像**的 BGR 影格，回傳 (55, 3)，未偵測到為 NaN。"""
import cv2
image = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
image.flags.writeable = False
results = self.holistic.process(image)
return np.asarray(extract_landmarks(results), dtype=np.float32)
def close(self):
self.holistic.close()
|