Spaces:
Runtime error
Runtime error
| import cv2 | |
| import mediapipe as mp | |
| import numpy as np | |
| from tqdm import tqdm | |
| class MediaPipeProcessor: | |
| def __init__(self, config): | |
| self.config = config.get('mediapipe_config', {}) | |
| self.setup_mediapipe() | |
| def setup_mediapipe(self): | |
| """Configura os modelos do MediaPipe""" | |
| self.mp_holistic = mp.solutions.holistic | |
| self.mp_drawing = mp.solutions.drawing_utils | |
| self.mp_drawing_styles = mp.solutions.drawing_styles | |
| self.holistic = self.mp_holistic.Holistic( | |
| static_image_mode=self.config.get('static_image_mode', False), | |
| model_complexity=self.config.get('model_complexity', 1), | |
| smooth_landmarks=self.config.get('smooth_landmarks', True), | |
| min_detection_confidence=self.config.get('min_detection_confidence', 0.5), | |
| min_tracking_confidence=self.config.get('min_tracking_confidence', 0.5) | |
| ) | |
| def process_video(self, video_path): | |
| """Processa o vídeo e extrai keypoints""" | |
| cap = cv2.VideoCapture(video_path) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| keypoints_data = [] | |
| print(f"Extraindo keypoints de {total_frames} frames...") | |
| for frame_idx in tqdm(range(total_frames), desc="Processando frames"): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Converter BGR para RGB | |
| frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| results = self.holistic.process(frame_rgb) | |
| frame_keypoints = self.extract_frame_keypoints(results) | |
| frame_keypoints['frame_index'] = frame_idx | |
| frame_keypoints['timestamp'] = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0 | |
| keypoints_data.append(frame_keypoints) | |
| cap.release() | |
| return keypoints_data | |
| def extract_frame_keypoints(self, results): | |
| """Extrai keypoints de um frame""" | |
| keypoints = {} | |
| # Face landmarks | |
| if results.face_landmarks: | |
| keypoints['face_landmarks'] = self.landmarks_to_array( | |
| results.face_landmarks.landmark, 468 | |
| ) | |
| # Pose landmarks | |
| if results.pose_landmarks: | |
| keypoints['pose_landmarks'] = self.landmarks_to_array( | |
| results.pose_landmarks.landmark, 33 | |
| ) | |
| # Left hand landmarks | |
| if results.left_hand_landmarks: | |
| keypoints['left_hand_landmarks'] = self.landmarks_to_array( | |
| results.left_hand_landmarks.landmark, 21 | |
| ) | |
| # Right hand landmarks | |
| if results.right_hand_landmarks: | |
| keypoints['right_hand_landmarks'] = self.landmarks_to_array( | |
| results.right_hand_landmarks.landmark, 21 | |
| ) | |
| return keypoints | |
| def landmarks_to_array(self, landmarks, expected_count): | |
| """Converte landmarks do MediaPipe para array numpy""" | |
| if not landmarks or len(landmarks) != expected_count: | |
| return np.zeros((expected_count, 3)) | |
| return np.array([[lm.x, lm.y, lm.z] for lm in landmarks]) |