Spaces:
Running
Running
| """Architecture-aware stroke model forward (CNN-LSTM vs pose+transformer stacks).""" | |
| from __future__ import annotations | |
| from typing import List | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from api import state | |
| from api.model_loader import ARCH_CNN_LSTM | |
| def imagenet_normalize_btc_hw(x01: torch.Tensor) -> torch.Tensor: | |
| """Normalize clip tensor from [0,1] to ImageNet stats. Shape (B, T, 3, H, W).""" | |
| mean = torch.tensor([0.485, 0.456, 0.406], device=x01.device, dtype=x01.dtype).view(1, 1, 3, 1, 1) | |
| std = torch.tensor([0.229, 0.224, 0.225], device=x01.device, dtype=x01.dtype).view(1, 1, 3, 1, 1) | |
| return (x01 - mean) / std | |
| def joint_seq_btj3_from_rgb_frames( | |
| segment_frames_rgb: List[np.ndarray], | |
| pose_estimator, | |
| device: torch.device | str, | |
| dtype: torch.dtype = torch.float32, | |
| ) -> torch.Tensor: | |
| """Build (1, T, 33, 3) pose tensor from RGB frames (uint8 or float).""" | |
| T = len(segment_frames_rgb) | |
| out = torch.zeros(1, T, 33, 3, device=device, dtype=dtype) | |
| for t, fr in enumerate(segment_frames_rgb): | |
| if fr.dtype != np.uint8: | |
| if float(fr.max()) <= 1.0: | |
| fr = (np.clip(fr, 0, 1) * 255.0).astype(np.uint8) | |
| else: | |
| fr = fr.astype(np.uint8) | |
| bgr = cv2.cvtColor(fr, cv2.COLOR_RGB2BGR) | |
| res = pose_estimator.process_frame(bgr) | |
| lm_list = pose_estimator.get_landmarks_as_list(res) | |
| if not lm_list or len(lm_list[0]) < 33: | |
| continue | |
| person = lm_list[0] | |
| for j in range(33): | |
| d = person[j] | |
| out[0, t, j, 0] = float(d["x"]) | |
| out[0, t, j, 1] = float(d["y"]) | |
| out[0, t, j, 2] = float(d["z"]) | |
| return out | |
| def run_stroke_model( | |
| segment_tensor_01: torch.Tensor, | |
| segment_frames_rgb: List[np.ndarray], | |
| device: torch.device | str | None = None, | |
| ) -> dict[str, torch.Tensor]: | |
| """ | |
| Run loaded stroke model on a 16-frame window. | |
| ``segment_tensor_01``: (1, 16, 3, 224, 224) float in [0, 1]. | |
| ``segment_frames_rgb``: same frames as numpy RGB (for per-frame MediaPipe). | |
| """ | |
| dev = device or state.device | |
| arch = getattr(state, "model_architecture", ARCH_CNN_LSTM) | |
| model = state.model | |
| pe = state.pose_estimator | |
| if arch == ARCH_CNN_LSTM: | |
| return model(segment_tensor_01.to(dev)) | |
| joint = joint_seq_btj3_from_rgb_frames(segment_frames_rgb, pe, dev, dtype=segment_tensor_01.dtype) | |
| x = imagenet_normalize_btc_hw(segment_tensor_01.to(dev)) | |
| return model(x, joint) | |