john
1aaef71f52
Document-driven MVP with FastAPI backend, Vue H5, WeChat mini shell, product demo, and Docker dev stack. Co-authored-by: Cursor <cursoragent@cursor.com>
120 lines
3.2 KiB
Python
120 lines
3.2 KiB
Python
"""Pose analysis pipeline — MediaPipe frame extraction with mock fallback."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from app.ai.pose_metrics import FrameLandmarks, build_movement_report, build_screening_report
|
|
from app.mock_data import MOCK_MOVEMENT_REPORT, MOCK_REPORT
|
|
|
|
MAX_FRAMES = 32
|
|
FRAME_STRIDE = 6
|
|
|
|
|
|
def analyze_task(task_type: str, video_path: Path | None = None) -> dict:
|
|
if video_path and video_path.exists():
|
|
mediapipe_result = _analyze_video(video_path, task_type)
|
|
if mediapipe_result:
|
|
return mediapipe_result
|
|
|
|
return _mock_result(task_type)
|
|
|
|
|
|
def _mock_result(task_type: str) -> dict:
|
|
if task_type == "movement_scoring":
|
|
return {
|
|
"engine": "mock",
|
|
"confidence": 0.88,
|
|
"report": {**MOCK_MOVEMENT_REPORT},
|
|
"modelVersion": "mock-movement-v1",
|
|
}
|
|
return {
|
|
"engine": "mock",
|
|
"confidence": 0.88,
|
|
"report": {**MOCK_REPORT},
|
|
"modelVersion": "mock-screening-v1",
|
|
}
|
|
|
|
|
|
def _analyze_video(video_path: Path, task_type: str) -> dict | None:
|
|
frames = _extract_pose_frames(video_path)
|
|
if not frames:
|
|
return None
|
|
|
|
try:
|
|
if task_type == "movement_scoring":
|
|
report = build_movement_report(frames)
|
|
else:
|
|
report = build_screening_report(frames)
|
|
except ValueError:
|
|
return None
|
|
|
|
confidences = []
|
|
if task_type != "movement_scoring":
|
|
for metric in report.get("metrics", []):
|
|
if "confidence" in metric:
|
|
confidences.append(metric["confidence"])
|
|
|
|
return {
|
|
"engine": "mediapipe",
|
|
"confidence": round(mean(confidences) if confidences else 0.9, 2),
|
|
"report": report,
|
|
"modelVersion": "mediapipe-pose-v2",
|
|
"frameCount": len(frames),
|
|
"sourceFile": video_path.name,
|
|
}
|
|
|
|
|
|
def _extract_pose_frames(video_path: Path) -> list[FrameLandmarks]:
|
|
try:
|
|
import cv2
|
|
import mediapipe as mp
|
|
except ImportError:
|
|
return []
|
|
|
|
capture = cv2.VideoCapture(str(video_path))
|
|
if not capture.isOpened():
|
|
return []
|
|
|
|
pose = mp.solutions.pose.Pose(
|
|
static_image_mode=False,
|
|
model_complexity=1,
|
|
min_detection_confidence=0.5,
|
|
min_tracking_confidence=0.5,
|
|
)
|
|
|
|
frames: list[FrameLandmarks] = []
|
|
index = 0
|
|
try:
|
|
while capture.isOpened() and len(frames) < MAX_FRAMES:
|
|
ok, image = capture.read()
|
|
if not ok:
|
|
break
|
|
if index % FRAME_STRIDE != 0:
|
|
index += 1
|
|
continue
|
|
index += 1
|
|
|
|
rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
result = pose.process(rgb)
|
|
if not result.pose_landmarks:
|
|
continue
|
|
|
|
frame: FrameLandmarks = {}
|
|
for idx, landmark in enumerate(result.pose_landmarks.landmark):
|
|
frame[idx] = {
|
|
"x": landmark.x,
|
|
"y": landmark.y,
|
|
"visibility": landmark.visibility,
|
|
}
|
|
frames.append(frame)
|
|
finally:
|
|
pose.close()
|
|
capture.release()
|
|
|
|
return frames
|
|
|
|
|
|
def mean(values: list[float]) -> float:
|
|
return sum(values) / len(values) if values else 0.0
|