wsy182
e86c2301ec
- Stripped _ENV_MAP and _apply_env_overrides from configs/load.py - Cleaned up unused imports in video_receiver.py - Restored MediaPipe env suppressors in app/main.py - Removed .env.example (replaced by config.yaml)
96 lines
3.9 KiB
Python
96 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
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 config
|
|
|
|
|
|
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={}", config.video.process_every_n_frames)
|
|
|
|
frame_count = 0
|
|
processed_count = 0
|
|
detector = DeadBugDetector(
|
|
model_path=config.model.resolved_path,
|
|
visibility_threshold=config.dead_bug.visibility_threshold,
|
|
extension_confirm_frames=config.dead_bug.extension_confirm_frames,
|
|
reset_confirm_frames=config.dead_bug.reset_confirm_frames,
|
|
prefer_gpu=config.model.prefer_gpu,
|
|
)
|
|
announcer = RepAnnouncer(
|
|
enabled=config.audio.rep_announcer_enabled,
|
|
rate=config.audio.rep_announcer_rate,
|
|
volume=config.audio.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 % config.video.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()
|