""" DataLoader for loading data from preprocessed posterior files. Adapted to this repository's training flow and compatible with train_mdm.py. Based on the external preprocessed_posterior_loader.py. The output format matches the original t2m_collate and can be used directly by existing training code. """ import os import torch import numpy as np from torch.utils.data import Dataset, DataLoader from pathlib import Path from os.path import join as pjoin from data_loaders.tensors import collate as mdm_collate class PreprocessedPosteriorDataset(Dataset): """ Dataset for loading data from preprocessed posterior files. Adapts the data format to the MDM training flow. The output format matches datasets such as HumanML3D and can use the existing collate function directly. """ def __init__(self, posterior_dir, max_samples=None, resample=False, mean=None, std=None): """ Args: posterior_dir: directory containing preprocessed posterior files max_samples: maximum number of samples; None loads all samples resample: whether to resample the latent while loading; False uses the preprocessed sample mean: mean used for normalization (optional, for compatibility) std: standard deviation used for normalization (optional, for compatibility) """ self.posterior_dir = posterior_dir self.resample = resample # Scan the directory and collect posterior file paths, excluding meta_info.pt and preprocess_config.pt. all_files = os.listdir(posterior_dir) posterior_files = sorted([ f for f in all_files if f.endswith('.pt') and f not in ['meta_info.pt', 'preprocess_config.pt'] ]) # Save the file path list. self.file_paths = [pjoin(posterior_dir, f) for f in posterior_files] self.total_samples = len(self.file_paths) print(f"Scanned directory: found {self.total_samples} posterior files") print(f"Scanned directory: found {self.total_samples} posterior files") # Limit the number of samples. if max_samples is not None: self.total_samples = min(self.total_samples, max_samples) self.file_paths = self.file_paths[:max_samples] print(f"Limited to the first {self.total_samples} samples") # Try to load metadata if available. meta_path = pjoin(posterior_dir, 'meta_info.pt') if os.path.exists(meta_path): meta_info = torch.load(meta_path) self.config = meta_info.get('config', {}) print("Loaded metadata configuration") else: self.config = {} # Store mean and std for compatibility with dataset interfaces such as HumanML3D. self.mean = mean self.std = std print(f"PreprocessedPosteriorDataset initialized with {self.total_samples} samples") if self.resample: print("Resampling mode enabled: resampling latent from posterior on each load") def sample_from_posterior(self, mean, logvar): """ Sample from the posterior distribution. Uses the reparameterization trick: z = mean + std * eps. Args: mean: posterior mean logvar: posterior log variance Returns: sampled latent """ std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) return mean + eps * std def inv_transform(self, data): """Inverse transform for compatibility.""" if self.mean is not None and self.std is not None: return data * self.std + self.mean return data def __len__(self): return self.total_samples def __getitem__(self, idx): """ Load posterior data for sample idx. Returns: dict: matches the format expected by the original t2m_collate and contains: - inp: motion data with shape (C, 1, T) - text: text description (string) - tokens: token list if available - lengths: motion length (integer) """ posterior_path = self.file_paths[idx] if not os.path.exists(posterior_path): raise FileNotFoundError(f"Posterior file not found: {posterior_path}") posterior_data = torch.load(posterior_path) # Get data. text_info = posterior_data['text'] m_len = posterior_data['m_len'] # If resampling is enabled, resample from mean and logvar. if self.resample: mean = posterior_data['mean'] # (T, C) logvar = posterior_data['logvar'] # (T, C) sample = self.sample_from_posterior(mean, logvar) else: if 'sample' in posterior_data: sample = posterior_data['sample'] # (T, C) elif 'z' in posterior_data: sample = posterior_data['z'] # (T, C) else: raise KeyError(f"Posterior file is missing the 'sample' or 'z' key: {posterior_path}") # Convert to the format expected by MDM: (T, C) -> (C, 1, T). # Align with the conversion in t2m_collate: torch.tensor(b[4].T).float().unsqueeze(1). inp = sample.T.unsqueeze(1).float() # (C, 1, T) # Handle length. if isinstance(m_len, (int, np.integer)): length = int(m_len) elif hasattr(m_len, 'item'): length = int(m_len.item()) else: length = int(m_len) # Handle text information across multiple formats. # text_info may be a string, dict, or another format. if isinstance(text_info, dict): text = text_info.get('caption', str(text_info)) tokens = text_info.get('tokens', None) elif isinstance(text_info, str): text = text_info tokens = None else: text = str(text_info) tokens = None # Build the return dictionary to match the format expected by t2m_collate. result = { 'inp': inp, 'text': text, 'lengths': length, } # Add tokens if available. if tokens is not None: result['tokens'] = tokens return result def collate_preprocessed_posterior(batch): """ Use the original collate function directly to keep the output format identical. Args: batch: list of dicts, each containing inp, text, lengths, etc. Returns: motion: (B, C, 1, T_max) tensor cond: dict containing a 'y' key with mask, lengths, text, etc. Output format exactly matches t2m_collate: - motion: (B, C, 1, T_max) - cond['y']['mask']: (B, 1, 1, T_max) - cond['y']['lengths']: (B,) - cond['y']['text']: list of str """ # Use the original collate function directly. return mdm_collate(batch) class PreprocessedPosteriorDataLoader: """ DataLoader wrapper for preprocessed posterior data. Provides an interface compatible with existing data loaders. This class emulates the original DataLoader interface, including dataset attribute access. """ def __init__(self, posterior_dir, batch_size=32, shuffle=True, num_workers=4, max_samples=None, resample=False, mean=None, std=None, drop_last=True): """ Create a DataLoader for preprocessed posterior data. Args: posterior_dir: directory containing preprocessed posterior files batch_size: batch size shuffle: whether to shuffle num_workers: number of data loading workers max_samples: maximum number of samples resample: whether to resample the latent while loading mean: mean used for normalization std: standard deviation used for normalization drop_last: whether to drop the last incomplete batch """ self.dataset = PreprocessedPosteriorDataset( posterior_dir, max_samples=max_samples, resample=resample, mean=mean, std=std ) self.dataloader = DataLoader( self.dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=collate_preprocessed_posterior, pin_memory=True, drop_last=drop_last, ) # Store attributes for external access and training_loop.py compatibility. self.batch_size = batch_size # These attributes are used for compatibility with target_cond_modifier in training_loop. # preprocessed_posterior may not need these, but keep the interface. self.mean = mean self.std = std def __iter__(self): return iter(self.dataloader) def __len__(self): return len(self.dataloader) def get_preprocessed_posterior_loader(posterior_dir, batch_size=32, shuffle=True, num_workers=4, max_samples=None, resample=False, mean=None, std=None, drop_last=True): """ Convenience function for creating a DataLoader for preprocessed posterior data. Args: posterior_dir: directory containing preprocessed posterior files batch_size: batch size shuffle: whether to shuffle num_workers: number of data loading workers max_samples: maximum number of samples resample: whether to resample the latent while loading mean: mean used for normalization std: standard deviation used for normalization drop_last: whether to drop the last incomplete batch Returns: PreprocessedPosteriorDataLoader object """ return PreprocessedPosteriorDataLoader( posterior_dir=posterior_dir, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, max_samples=max_samples, resample=resample, mean=mean, std=std, drop_last=drop_last ) # Test code if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description='Test the preprocessed posterior data loader (MDM-compatible version)') parser.add_argument('--posterior_dir', type=str, required=True, help='Directory containing preprocessed posterior files') parser.add_argument('--batch_size', type=int, default=4, help='Batch size') parser.add_argument('--max_samples', type=int, default=None, help='Maximum number of samples for testing') parser.add_argument('--resample', action='store_true', help='Enable resampling mode') args = parser.parse_args() print("="*50) print("Testing preprocessed posterior data loader (MDM-compatible version)") print("="*50) # Create dataloader. dataloader = get_preprocessed_posterior_loader( args.posterior_dir, batch_size=args.batch_size, shuffle=False, num_workers=0, # Use a single process for testing. max_samples=args.max_samples, resample=args.resample ) print(f"\nDataLoader created with {len(dataloader)} batches") print(f"Resampling mode: {'enabled' if args.resample else 'disabled'}") # Test loading the first batch. print("\nTesting first batch load...") for batch_idx, (motion, cond) in enumerate(dataloader): print(f"\nBatch {batch_idx}:") print(f" motion shape: {motion.shape}") # Expected: (B, C, 1, T_max) print(f" mask shape: {cond['y']['mask'].shape}") # Expected: (B, 1, 1, T_max) print(f" lengths: {cond['y']['lengths']}") print(f" texts: {len(cond['y']['text'])} text entries") # Only test the first batch. if batch_idx == 0: print("\nText information for the first sample:") print(f" {cond['y']['text'][0]}") # Show motion statistics. print("\nMotion statistics for the first sample:") first_len = cond['y']['lengths'][0].item() first_motion_valid = motion[0, :, :, :first_len] # Only inspect the valid part. print(f" valid length: {first_len}") print(f" mean: {first_motion_valid.mean().item():.6f}") print(f" std: {first_motion_valid.std().item():.6f}") print(f" min: {first_motion_valid.min().item():.6f}") print(f" max: {first_motion_valid.max().item():.6f}") break print("\nTest completed!") print("\nUsage example:") print(" from data_loaders.preprocessed_posterior_loader import get_preprocessed_posterior_loader") print(" data = get_preprocessed_posterior_loader(posterior_dir, batch_size=32)") print(" for motion, cond in data:") print(" # motion: (B, C, 1, T)") print(" # cond['y']['mask']: (B, 1, 1, T)") print(" # cond['y']['lengths']: (B,)") print(" # cond['y']['text']: list of texts")