import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader, random_split import pandas as pd import numpy as np from transformers import BertTokenizer import os from pose_format import Pose import matplotlib.pyplot as plt from matplotlib import animation from fastdtw import fastdtw # Keep this import from scipy.spatial.distance import cosine from config import MAX_TEXT_LEN, TARGET_NUM_FRAMES, BATCH_SIZE, TEACHER_FORCING_RATIO, SMOOTHING_ENABLED # ===== KEYPOINT SELECTION ===== selected_keypoint_indices = list(np.r_[0:25, 501:522, 522:543]) NUM_KEYPOINTS = len(selected_keypoint_indices) POSE_DIM = NUM_KEYPOINTS * 3 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ===== SMOOTHING ===== def selective_smoothing(preds): smoothed = preds.clone() body_indices = slice(0, 25 * 3) for t in range(1, preds.shape[1] - 1): smoothed[:, t, body_indices] = ( 0.25 * preds[:, t - 1, body_indices] + 0.5 * preds[:, t, body_indices] + 0.25 * preds[:, t + 1, body_indices] ) return smoothed # ===== HAND INDEX SETUP ===== hand_indices = list(range(15 * 3, POSE_DIM)) # hand joints = after body joints in flattened 3D vector joint_weights = torch.ones(POSE_DIM).to(device) joint_weights[hand_indices] *= 3.0 # ===== GLOBAL NORMALIZATION ===== def compute_global_mean_std(pose_folder, csv_file): data = pd.read_csv(csv_file) all_poses = [] # Store valid masks separately to ensure correct normalization all_masks = [] for filename in data["filename"]: pose_path = os.path.join(pose_folder, filename) with open(pose_path, "rb") as f: pose = Pose.read(f.read()) keypoints = np.array(selected_keypoint_indices) # (T, 1, K, 3) -> (T, K, 3) pose_data = np.squeeze(pose.body.data, axis=1)[:, keypoints, :] # (T, 1, K) -> (T, K) confidence = np.squeeze(pose.body.confidence, axis=1)[:, keypoints] # Reshape to (T, K*3) pose_data_flat = pose_data.reshape(pose_data.shape[0], -1) # Reshape confidence to (T, K*3) - repeat confidence for each coordinate confidence_flat = np.repeat(confidence, 3, axis=1) # Create a mask based on confidence for the flattened data mask_flat = (confidence_flat > 0.5).astype(np.float32) # Append the full pose data and mask for interpolation later all_poses.append(pose_data_flat) all_masks.append(mask_flat) # Pad or interpolate all poses and masks to a fixed length (TARGET_NUM_FRAMES) padded_poses = [] padded_masks = [] for pose_data_flat, mask_flat in zip(all_poses, all_masks): current_frames = pose_data_flat.shape[0] if current_frames < TARGET_NUM_FRAMES: pad_len = TARGET_NUM_FRAMES - current_frames pose_pad = np.zeros((pad_len, POSE_DIM)) mask_pad = np.zeros((pad_len, POSE_DIM)) # Pad mask with zeros padded_pose = np.concatenate([pose_data_flat, pose_pad], axis=0) padded_mask = np.concatenate([mask_flat, mask_pad], axis=0) else: indices = np.linspace(0, current_frames - 1, TARGET_NUM_FRAMES).astype(int) padded_pose = pose_data_flat[indices] padded_mask = mask_flat[indices] padded_poses.append(padded_pose) padded_masks.append(padded_mask) # Stack all padded poses and masks: (Total_Samples * TARGET_NUM_FRAMES, POSE_DIM) stacked_poses = np.vstack(padded_poses) stacked_masks = np.vstack(padded_masks) # Compute mean and std using the masks to only include valid points # Weighted average using mask as weights mean = np.sum(stacked_poses * stacked_masks, axis=0) / (np.sum(stacked_masks, axis=0) + 1e-8) # Add epsilon for stability # Compute variance, then sqrt for std variance = np.sum(stacked_masks * (stacked_poses - mean)**2, axis=0) / (np.sum(stacked_masks, axis=0) + 1e-8) std = np.sqrt(variance) std[std == 0] = 1e-8 # Avoid division by zero return mean, std #POSE_FOLDER = "/content/drive/MyDrive/pose/words/ase" CSV_FILE = "annotated.csv" mean_path = "global_mean.npy" std_path = "global_std.npy" if os.path.exists(mean_path) and os.path.exists(std_path): print("Loading global mean and std from file.") GLOBAL_MEAN = np.load(mean_path) GLOBAL_STD = np.load(std_path) else: print("Computing global mean and std from dataset.") GLOBAL_MEAN, GLOBAL_STD = compute_global_mean_std(POSE_FOLDER, CSV_FILE) # Save the computed mean and std # Ensure they are not MaskedArrays if the computation somehow produced them # If compute_global_mean_std is modified to return standard arrays, this is redundant but safe if isinstance(GLOBAL_MEAN, np.ma.MaskedArray): GLOBAL_MEAN = GLOBAL_MEAN.data if isinstance(GLOBAL_STD, np.ma.MaskedArray): GLOBAL_STD = GLOBAL_STD.data np.save(mean_path, GLOBAL_MEAN) np.save(std_path, GLOBAL_STD) GLOBAL_MEAN_T = torch.tensor(GLOBAL_MEAN).float().to(device) GLOBAL_STD_T = torch.tensor(GLOBAL_STD).float().to(device) class TextToPoseDataset(Dataset): def __init__(self, csv_file, pose_folder, tokenizer, is_train=True): self.data = pd.read_csv(csv_file) self.pose_folder = pose_folder self.tokenizer = tokenizer self.is_train = is_train # enable augment only during training def __len__(self): return len(self.data) def load_pose_data_and_mask(self, filename): pose_path = os.path.join(self.pose_folder, filename) with open(pose_path, "rb") as f: pose = Pose.read(f.read()) keypoints = np.array(selected_keypoint_indices) pose_data = np.squeeze(pose.body.data, axis=1)[:, keypoints, :] confidence = np.squeeze(pose.body.confidence, axis=1)[:, keypoints] return pose_data, confidence def apply_augmentations(self, pose_data, confidence): T = pose_data.shape[0] # Temporal warp (resample frame indices with small noise) if T > TARGET_NUM_FRAMES and np.random.rand() < 0.5: indices = np.linspace(0, T - 1, TARGET_NUM_FRAMES) jitter = np.random.uniform(-0.5, 0.5, size=indices.shape) indices = np.clip(indices + jitter, 0, T - 1).astype(int) pose_data = pose_data[indices] confidence = confidence[indices] # Mirror (flip X-axis) if np.random.rand() < 0.3: pose_data[..., 0] *= -1 # Jitter (small Gaussian noise) if np.random.rand() < 0.3: pose_data += np.random.normal(0, 0.02, pose_data.shape) return pose_data, confidence def __getitem__(self, idx): row = self.data.iloc[idx] filename = row["filename"] text = row["text"] input_ids = self.tokenizer( text, padding="max_length", truncation=True, max_length=MAX_TEXT_LEN, return_tensors="pt" ) pose_data, confidence = self.load_pose_data_and_mask(filename) if self.is_train: pose_data, confidence = self.apply_augmentations(pose_data, confidence) # OLD Flatten pose_data_flat = pose_data.reshape(pose_data.shape[0], -1) confidence_flat = np.repeat(confidence, 3, axis=1) mask_flat = (confidence_flat > 0.5).astype(np.float32) # Pad or warp to fixed length current_frames = pose_data_flat.shape[0] if current_frames < TARGET_NUM_FRAMES: pad_len = TARGET_NUM_FRAMES - current_frames pose_pad = np.zeros((pad_len, POSE_DIM)) mask_pad = np.zeros((pad_len, POSE_DIM)) padded_pose = np.concatenate([pose_data_flat, pose_pad], axis=0) padded_mask = np.concatenate([mask_flat, mask_pad], axis=0) else: indices = np.linspace(0, current_frames - 1, TARGET_NUM_FRAMES).astype(int) padded_pose = pose_data_flat[indices] padded_mask = mask_flat[indices] # Normalize normalized_pose = (padded_pose - GLOBAL_MEAN) / GLOBAL_STD return ( input_ids.input_ids.squeeze(0), input_ids.attention_mask.squeeze(0), torch.tensor(normalized_pose).float(), torch.tensor(padded_mask).float(), text ) def collate_fn(batch): input_ids, attn_masks, poses, masks, words = zip(*batch) return ( torch.stack(input_ids), torch.stack(attn_masks), torch.stack(poses), torch.stack(masks), list(words) )