import os, torch, random import numpy as np from pathlib import Path from torch.utils.data import Dataset, Sampler, DataLoader from functools import partial from typing import List, Dict, Tuple, Optional, Union from torch_geometric.data import Data, Batch from mydataset import npz_2_data from skel_pose_graph import SkelPoseGraph, find_feet class SequenceDataset(Dataset): """ Dataset that loads complete motion sequences with text descriptions. Each item is a complete motion sequence rather than a single frame. """ copy_orig_contact = False def __init__(self, text_dir: Optional[str] = None, seq_length: Optional[int] = None): """ Initialize the sequence dataset. Args: text_dir: Directory containing text description files (.txt) seq_length: Fixed sequence length. If None, load full sequences. If specified, sequences will be truncated or padded to this length. """ # Skeleton and pose data self.skel_list = [] self.pose_list = [] # File information self.filepaths = [] self.frame_cnts = [] self.start_frames = [] self.end_frames = [] # Motion set related info self.mi_ri_2_fi = [] # mi: semantic motion index (same mi means semantically identical motion) # ri: 0<=ri list of text descriptions for this motion # Sequence length configuration self.seq_length = seq_length def add_data(self, lo, go, qb, edges, q, p, qv, pv, pprev, c, r, filepath, mi, tf, text=None): """Add motion data to the dataset.""" # Convert to SkelData and PoseData sd, pdl = npz_2_data(lo, go, qb, edges, q, p, qv, pv, pprev, c, r, tf) self.skel_list.append(sd) self.pose_list.extend(pdl) # Update file info fi = len(self.filepaths) nFrame = q.shape[0] start = sum(self.frame_cnts) end = start + nFrame self.filepaths.append(filepath) self.frame_cnts.append(nFrame) self.start_frames.append(start) self.end_frames.append(end) # Update motion set info assert len(self.mi_ri_2_fi) >= mi if len(self.mi_ri_2_fi) == mi: # New semantic motion set self.mi_ri_2_fi.append([]) # Add text for new motion (only once per motion index) self.mi_2_text.append(text if text is not None else []) else: # Ensure consistent frame count among retargeted dataset orig_fi = self.mi_ri_2_fi[mi][0] orig_nFrame = self.frame_cnts[orig_fi] assert orig_nFrame == nFrame, f"Frame count mismatch: {orig_nFrame} vs {nFrame}" if self.copy_orig_contact: # Copy contact from the original motion (optional) lf, rf = find_feet(sd) orig_sd = self.skel_list[orig_fi] orig_start = self.start_frames[orig_fi] orig_pdl = self.pose_list[orig_start : orig_start + orig_nFrame] orig_lf, orig_rf = find_feet(orig_sd) for pdi, orig_pdi in zip(pdl, orig_pdl): pdi.c[lf] = orig_pdi.c[orig_lf] pdi.c[rf] = orig_pdi.c[orig_rf] self.mi_ri_2_fi[mi].append(fi) def add_data_from_npz(self, mi, npz_fp, bvh_fp=None, text=None): """Load data from an NPZ file.""" data = np.load(npz_fp) if bvh_fp is None: bvh_fp = npz_fp # placeholder self.add_data(**data, filepath=bvh_fp, mi=mi, text=text) def load_data_dir_pairs(self, data_dir): """ Load paired motion data from a directory. The directory should contain: - pair.txt: List of (source, target) motion pairs - *.npz files: Motion data files """ pair_path = os.path.join(data_dir, "pair.txt") assert os.path.exists(pair_path), f"{pair_path} does not exist" bvh_prefix = os.path.join(os.path.dirname(data_dir), "bvh") src_id_map = {} with open(pair_path, "r") as pair_file: for line in pair_file: if line.strip() == "": continue src_rel_path, dst_rel_path = line.strip().split() if src_rel_path in src_id_map: src_id = src_id_map[src_rel_path] else: # New source src_id = len(self.mi_ri_2_fi) src_id_map[src_rel_path] = src_id npz_fp = os.path.join(data_dir, src_rel_path) bvh_fp = os.path.join(bvh_prefix, Path(src_rel_path).with_suffix("")) # Load text if text_dir is provided text = self._load_text_for_file(src_rel_path) self.add_data_from_npz(src_id, npz_fp, bvh_fp, text=text) if dst_rel_path in src_id_map: continue else: npz_fp = os.path.join(data_dir, dst_rel_path) bvh_fp = os.path.join(bvh_prefix, Path(dst_rel_path).with_suffix("")) # No text for retargeted motions (they share text with source) self.add_data_from_npz(src_id, npz_fp, bvh_fp, text=None) def _load_text_for_file(self, file_rel_path): """ Load text descriptions for a motion file. Text files are stored in self.text_dir with the same name as motion file but .txt extension. Each line in the text file contains one text description in the format: caption#tokens#start_time#end_time Args: file_rel_path: Relative path to the motion file (e.g., 'test/006888.bvh.npz') Returns: List of text description dictionaries, or empty list if no text file found """ if self.text_dir is None: return [] # Extract filename without extensions basename = os.path.basename(file_rel_path) # '006888.bvh.npz' filename = Path(basename).stem # '006888.bvh' if filename.endswith('.bvh'): filename = filename[:-4] # '006888' text_file_path = os.path.join(self.text_dir, filename + '.txt') if not os.path.exists(text_file_path): return [] text_list = [] try: with open(text_file_path, 'r', encoding='utf-8') as f: for line_num, line in enumerate(f.readlines(), 1): line = line.strip() if not line: continue # Parse format: caption#tokens#start_time#end_time parts = line.split('#') if len(parts) < 4: print(f"Warning: Malformed line {line_num} in {text_file_path}: expected 4 parts, got {len(parts)}") continue try: text_dict = { 'caption': parts[0].strip(), 'tokens': parts[1].strip().split(' ') if parts[1].strip() else [], 'start_time': float(parts[2].strip()) if parts[2].strip() else 0.0, 'end_time': float(parts[3].strip()) if parts[3].strip() else 0.0, } text_list.append(text_dict) except ValueError as e: print(f"Warning: Failed to parse line {line_num} in {text_file_path}: {e}") continue except Exception as e: print(f"Warning: Failed to load text from {text_file_path}: {e}") return [] return text_list def get_mi_ri_fi_graph(self, mi, ri, frame): """Get skeleton pose graph for a specific motion, rig, and frame.""" fi = self.mi_ri_2_fi[mi][ri] si = self.skel_list[fi] pi = self.pose_list[self.start_frames[fi] + frame] return SkelPoseGraph(si, pi) def get_sequence_graphs(self, mi, ri, start_frame=0, length=None): """ Get a sequence of graphs for a specific motion and rig. Args: mi: Motion index ri: Rig index start_frame: Starting frame index length: Number of frames to load. If None, load all remaining frames. Returns: List of SkelPoseGraph objects """ fi = self.mi_ri_2_fi[mi][ri] frame_cnt = self.frame_cnts[fi] if length is None: length = frame_cnt - start_frame end_frame = min(start_frame + length, frame_cnt) graphs = [] for frame in range(start_frame, end_frame): graphs.append(self.get_mi_ri_fi_graph(mi, ri, frame)) return graphs, end_frame - start_frame def __getitem__(self, idx): """ Get a complete motion sequence with text. Args: idx: Tuple (mi, src_ri, tgt_ri) where: - mi: motion index - src_ri: source rig index - tgt_ri: target rig index Returns: Tuple (src_graphs, tgt_graphs, text, m_lens) where: - src_graphs: List of source skeleton pose graphs - tgt_graphs: List of target skeleton pose graphs - text: List of text descriptions for this motion - m_lens: Length of the motion sequence (number of frames) """ mi, src_ri, tgt_ri = idx # Get frame count for this motion fi = self.mi_ri_2_fi[mi][src_ri] frame_cnt = self.frame_cnts[fi] # Handle sequence length if self.seq_length is not None: # Use specified sequence length actual_length = min(self.seq_length, frame_cnt) start_frame = 0 else: # Use full sequence actual_length = frame_cnt start_frame = 0 # Get sequence of graphs src_graphs, _ = self.get_sequence_graphs(mi, src_ri, start_frame, actual_length) tgt_graphs, _ = self.get_sequence_graphs(mi, tgt_ri, start_frame, actual_length) # Get text descriptions text = self.mi_2_text[mi] if mi < len(self.mi_2_text) else [] # Actual sequence length m_lens = len(src_graphs) return src_graphs, tgt_graphs, text, m_lens def __len__(self): """Return the number of motion sequences in the dataset.""" return len(self.mi_ri_2_fi) class SequenceSampler(Sampler): """ Sampler that yields motion sequence indices. Unlike frame-based samplers, this sampler yields complete motion sequences. """ def __init__(self, dataset, batch_size, shuffle=True): """ Initialize the sequence sampler. Args: dataset: SequenceDataset instance batch_size: Number of sequences per batch shuffle: Whether to shuffle motion indices """ self.dataset = dataset self.batch_size = batch_size self.shuffle = shuffle # Motion indices self.motion_indices = list(range(len(self.dataset.mi_ri_2_fi))) # For specifying fixed src/tgt rig indices self.src_ri = None self.tgt_ri = None def __iter__(self): # Shuffle motion indices if required indices = self.motion_indices.copy() if self.shuffle: random.shuffle(indices) batch = [] for mi in indices: # Random src/tgt skeletons (including the original ones) R = len(self.dataset.mi_ri_2_fi[mi]) if R >= 2: src_ri, tgt_ri = random.sample(range(R), 2) else: # Only one rig available src_ri = tgt_ri = 0 # Override if src/tgt ri is specified if self.src_ri is not None: src_ri = self.src_ri if self.tgt_ri is not None: tgt_ri = self.tgt_ri batch.append((mi, src_ri, tgt_ri)) # Yield batch when full if len(batch) == self.batch_size: yield batch batch = [] # Yield remaining items if len(batch) > 0: yield batch def __len__(self): """Return the number of batches.""" return (len(self.motion_indices) + self.batch_size - 1) // self.batch_size def sequence_collate_fn(batch, mask_option=[], device="cpu", pad_to_max=False, return_list=False): """ Collate function for sequence-based batches. Args: batch: List of tuples (src_graphs, tgt_graphs, text, m_lens) from __getitem__ mask_option: List of strings indicating which graphs to apply random masking to device: Device to move tensors to pad_to_max: Default false. If True, pad all sequences to the maximum length in the batch return_list: If False (default), return aggregated Batch objects. If True, return list of Batch objects (one per sequence). Returns: If return_list=False (default): Tuple (src_batch, tgt_batch, text_list, m_lens, masks, batch_info) where: - src_batch: Single aggregated Batch containing all sequences - tgt_batch: Single aggregated Batch containing all sequences - text_list: List of text descriptions - m_lens: Numpy array of sequence lengths [T_0, T_1, ..., T_{B-1}] - masks: Tensor of shape (batch_size, max_len) indicating valid frames (if pad_to_max) - batch_info: Dict with keys: - 'cumsum_lens': Cumulative sum of lengths for splitting [0, T_0, T_0+T_1, ...] - 'num_sequences': Number of sequences in batch (batch_size) - 'total_frames': Total number of frames across all sequences If return_list=True: Tuple (src_batch_list, tgt_batch_list, text_list, m_lens, masks, batch_info) where: - src_batch_list: List of batched source graphs (one per sequence) - tgt_batch_list: List of batched target graphs (one per sequence) - text_list: List of text descriptions - m_lens: Numpy array of sequence lengths - masks: Tensor of shape (batch_size, max_len) indicating valid frames (if pad_to_max) - batch_info: Dict with metadata """ # Unpack batch src_graphs_list = [item[0] for item in batch] tgt_graphs_list = [item[1] for item in batch] text_list = [item[2] for item in batch] m_lens = [item[3] for item in batch] if pad_to_max and len(set(m_lens)) > 1: # Sequences have different lengths, need to pad max_len = max(m_lens) # Create masks (True for valid frames, False for padding) masks = [] padded_src_graphs_list = [] padded_tgt_graphs_list = [] for src_graphs, tgt_graphs, seq_len in zip(src_graphs_list, tgt_graphs_list, m_lens): # Create mask for this sequence mask = [True] * seq_len + [False] * (max_len - seq_len) masks.append(mask) # Pad sequences by repeating the last frame if seq_len < max_len: padding_needed = max_len - seq_len # Repeat last frame for padding src_graphs_padded = src_graphs + [src_graphs[-1]] * padding_needed tgt_graphs_padded = tgt_graphs + [tgt_graphs[-1]] * padding_needed else: src_graphs_padded = src_graphs tgt_graphs_padded = tgt_graphs padded_src_graphs_list.append(src_graphs_padded) padded_tgt_graphs_list.append(tgt_graphs_padded) src_graphs_list = padded_src_graphs_list tgt_graphs_list = padded_tgt_graphs_list m_lens = [max_len] * len(m_lens) masks = torch.BoolTensor(masks).to(device) else: # All sequences have the same length, no padding needed masks = None # Batch graphs for each sequence src_batch_list = [] tgt_batch_list = [] for src_graphs, tgt_graphs in zip(src_graphs_list, tgt_graphs_list): src_batch = Batch.from_data_list(src_graphs).to(device) tgt_batch = Batch.from_data_list(tgt_graphs).to(device) # Apply masking if requested if "src" in mask_option: src_batch.mask = rnd_mask(src_batch, consq_n=len(src_graphs)) if "tgt" in mask_option: tgt_batch.mask = rnd_mask(tgt_batch, consq_n=len(tgt_graphs)) src_batch_list.append(src_batch) tgt_batch_list.append(tgt_batch) # Random select one text for each sequence if multiple texts are available for i, texts in enumerate(text_list): text_list[i] = random.choice(texts) if texts else {} # Create batch info for splitting later cumsum_lens = np.concatenate([[0], np.cumsum(m_lens)]) batch_info = { 'cumsum_lens': cumsum_lens, # [0, T_0, T_0+T_1, ..., sum(T_i)] 'num_sequences': len(m_lens), # Batch size 'total_frames': int(np.sum(m_lens)), # Total frames } # Return list format if requested if return_list: return src_batch_list, tgt_batch_list, text_list, m_lens, masks, batch_info # Aggregate all sequences into a single Batch (default behavior) # Flatten all graphs from all sequences all_src_graphs = [graph for graphs in src_graphs_list for graph in graphs] all_tgt_graphs = [graph for graphs in tgt_graphs_list for graph in graphs] # Create single aggregated batch src_batch = Batch.from_data_list(all_src_graphs).to(device) tgt_batch = Batch.from_data_list(all_tgt_graphs).to(device) # Apply masking if requested (on the aggregated batch) if "src" in mask_option: src_batch.mask = rnd_mask(src_batch, consq_n=sum(m_lens)) if "tgt" in mask_option: tgt_batch.mask = rnd_mask(tgt_batch, consq_n=sum(m_lens)) return src_batch, tgt_batch, text_list, m_lens, masks, batch_info def get_sequence_data_loader( data_dir, batch_size, shuffle=True, mask_option=[], device="cpu", text_dir=None, seq_length=None, pad_to_max=True, return_list=False ): """ Create a DataLoader for sequence-based motion data. Args: data_dir: Directory containing motion data batch_size: Number of sequences per batch shuffle: Whether to shuffle the data mask_option: List of strings for masking ('src', 'tgt') device: Device to load data to text_dir: Directory containing text descriptions seq_length: Fixed sequence length (None for variable length) pad_to_max: Whether to pad sequences to max length in batch return_list: If False (default), return aggregated Batch objects. If True, return list of Batch objects (one per sequence). Returns: DataLoader instance """ ds = SequenceDataset(text_dir=text_dir, seq_length=seq_length) ds.load_data_dir_pairs(data_dir) sampler = SequenceSampler(ds, batch_size=batch_size, shuffle=shuffle) dl = DataLoader( ds, batch_sampler=sampler, collate_fn=partial( sequence_collate_fn, mask_option=mask_option, device=device, pad_to_max=pad_to_max, return_list=return_list, ), pin_memory=True, ) return dl