Compute mean and std from motion data using the full dataset.
(data_list, num_workers=8)
| 22 | |
| 23 | |
| 24 | def compute_statistics(data_list, num_workers=8): |
| 25 | """Compute mean and std from motion data using the full dataset.""" |
| 26 | # Collect all motion tensors |
| 27 | all_motions = [] |
| 28 | failed_count = 0 |
| 29 | |
| 30 | def load_single_motion(entry): |
| 31 | motion_path = entry['motion_path'] |
| 32 | try: |
| 33 | motion = load_motion_tensor(motion_path) |
| 34 | return motion |
| 35 | except Exception as e: |
| 36 | return None |
| 37 | |
| 38 | # Use thread pool for parallel loading |
| 39 | with ThreadPoolExecutor(max_workers=num_workers) as executor: |
| 40 | futures = {executor.submit(load_single_motion, entry): entry for entry in data_list} |
| 41 | for future in tqdm(as_completed(futures), total=len(futures), desc="Loading motions for stats"): |
| 42 | result = future.result() |
| 43 | if result is not None: |
| 44 | all_motions.append(result) |
| 45 | else: |
| 46 | failed_count += 1 |
| 47 | |
| 48 | print(f"Loaded {len(all_motions)} motions, {failed_count} failed") |
| 49 | |
| 50 | # Concatenate all motions along frame dimension |
| 51 | all_frames = torch.cat(all_motions, dim=0) |
| 52 | print(f"Total frames: {all_frames.shape[0]}, Motion dim: {all_frames.shape[1]}") |
| 53 | |
| 54 | # Compute mean and std |
| 55 | mean = all_frames.mean(dim=0).numpy() |
| 56 | std = all_frames.std(dim=0).numpy() |
| 57 | |
| 58 | # Avoid division by zero |
| 59 | std = np.clip(std, a_min=1e-8, a_max=None) |
| 60 | |
| 61 | return mean, std |
| 62 | |
| 63 | |
| 64 | def main(): |