wsy182
c8fd057129
- All settings moved to config.yaml - configs/load.py reads from config.yaml with env var overrides - Environment variables still work for backward compatibility - Added pyyaml to requirements
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
|
|
import cv2
|
|
from aiortc.mediastreams import MediaStreamError
|
|
from loguru import logger
|
|
|
|
from app.audio.rep_announcer import RepAnnouncer
|
|
from app.exercises.dead_bug.detector import DeadBugDetector
|
|
from app.rendering.window_display import close_window, is_esc_pressed, show_frame
|
|
from configs.load import (
|
|
EXTENSION_CONFIRM_FRAMES,
|
|
MODEL_PATH,
|
|
PREFER_GPU,
|
|
PROCESS_EVERY_N_FRAMES,
|
|
REP_ANNOUNCER_ENABLED,
|
|
REP_ANNOUNCER_RATE,
|
|
REP_ANNOUNCER_VOLUME,
|
|
RESET_CONFIRM_FRAMES,
|
|
VISIBILITY_THRESHOLD,
|
|
)
|
|
|
|
|
|
def _format_pose_debug(pose_result) -> str:
|
|
metrics = pose_result.metrics
|
|
if metrics is None:
|
|
return "metrics=None"
|
|
return (
|
|
f"side={pose_result.side}, standard={pose_result.is_standard}, "
|
|
f"angles(le={metrics.left_elbow_angle:.1f}, re={metrics.right_elbow_angle:.1f}, "
|
|
f"lk={metrics.left_knee_angle:.1f}, rk={metrics.right_knee_angle:.1f}), "
|
|
f"extended(la={metrics.left_arm_extended}, ra={metrics.right_arm_extended}, "
|
|
f"ll={metrics.left_leg_extended}, rl={metrics.right_leg_extended})"
|
|
)
|
|
|
|
|
|
class VideoReceiver:
|
|
def __init__(self, track) -> None:
|
|
self._track = track
|
|
|
|
async def run(self) -> None:
|
|
logger.info("Start receiving video frames, process_every_n={}", PROCESS_EVERY_N_FRAMES)
|
|
|
|
frame_count = 0
|
|
processed_count = 0
|
|
detector = DeadBugDetector(
|
|
model_path=MODEL_PATH,
|
|
visibility_threshold=VISIBILITY_THRESHOLD,
|
|
extension_confirm_frames=EXTENSION_CONFIRM_FRAMES,
|
|
reset_confirm_frames=RESET_CONFIRM_FRAMES,
|
|
prefer_gpu=PREFER_GPU,
|
|
)
|
|
announcer = RepAnnouncer(
|
|
enabled=REP_ANNOUNCER_ENABLED,
|
|
rate=REP_ANNOUNCER_RATE,
|
|
volume=REP_ANNOUNCER_VOLUME,
|
|
)
|
|
last_announced_rep = 0
|
|
last_pose_result = None
|
|
last_annotated = None
|
|
|
|
try:
|
|
while True:
|
|
frame = await self._track.recv()
|
|
frame_count += 1
|
|
raw_img = frame.to_ndarray(format="bgr24")
|
|
timestamp_ms = int(frame.time * 1000) if frame.time is not None else frame_count * 33
|
|
|
|
if frame_count % PROCESS_EVERY_N_FRAMES == 0 or last_pose_result is None:
|
|
processed_count += 1
|
|
last_annotated, last_pose_result = detector.process_frame(raw_img, timestamp_ms)
|
|
if last_pose_result.rep_count > last_announced_rep:
|
|
last_announced_rep = last_pose_result.rep_count
|
|
announcer.announce_count(last_announced_rep)
|
|
|
|
display_img = last_annotated if last_annotated is not None else raw_img
|
|
show_frame(display_img)
|
|
|
|
if frame_count % 100 == 0:
|
|
logger.info(
|
|
"Received {} frames, processed={}, raw_shape={}, reps={}, phase={}, feedback={}, {}",
|
|
frame_count,
|
|
processed_count,
|
|
raw_img.shape,
|
|
last_pose_result.rep_count if last_pose_result is not None else 0,
|
|
last_pose_result.phase.value if last_pose_result is not None else "none",
|
|
" | ".join(last_pose_result.feedback) if last_pose_result is not None else "",
|
|
_format_pose_debug(last_pose_result) if last_pose_result is not None else "metrics=None",
|
|
)
|
|
|
|
if is_esc_pressed():
|
|
logger.info("ESC pressed, closing display")
|
|
break
|
|
|
|
except asyncio.CancelledError:
|
|
logger.info("Video receive task cancelled")
|
|
except MediaStreamError:
|
|
logger.info("Video track ended")
|
|
except Exception as e:
|
|
logger.exception(f"Video receive error: {e!r}")
|
|
finally:
|
|
announcer.close()
|
|
detector.close()
|
|
close_window()
|