()
| 62 | |
| 63 | |
| 64 | def main(): |
| 65 | parser = argparse.ArgumentParser(description="Prepare ViMoGen-228K training data") |
| 66 | parser.add_argument("--input_json", type=str, required=True, |
| 67 | help="Path to original ViMoGen-228K.json") |
| 68 | parser.add_argument("--motion_root", type=str, default="./data/ViMoGen-228K", |
| 69 | help="Root directory for motion files") |
| 70 | parser.add_argument("--output_dir", type=str, default="./data/meta_info", |
| 71 | help="Output directory for processed files") |
| 72 | parser.add_argument("--num_workers", type=int, default=8, |
| 73 | help="Number of workers for parallel loading") |
| 74 | parser.add_argument("--skip_stats", action="store_true", |
| 75 | help="Skip statistics computation (use pre-computed mean/std in ./data/meta_info/)") |
| 76 | args = parser.parse_args() |
| 77 | |
| 78 | os.makedirs(args.output_dir, exist_ok=True) |
| 79 | |
| 80 | # Load original JSON |
| 81 | print(f"Loading {args.input_json}...") |
| 82 | with open(args.input_json, 'r') as f: |
| 83 | data_list = json.load(f) |
| 84 | print(f"Loaded {len(data_list)} entries") |
| 85 | |
| 86 | # Process each entry |
| 87 | print("Processing entries...") |
| 88 | for entry in tqdm(data_list, desc="Adding sample_id and motion_root"): |
| 89 | # Add sample_id from id field |
| 90 | entry['sample_id'] = str(entry['id']) |
| 91 | |
| 92 | # Prefix motion_path with motion_root |
| 93 | original_path = entry['motion_path'] |
| 94 | entry['motion_path'] = os.path.join(args.motion_root, original_path) |
| 95 | |
| 96 | # Compute statistics unless skipped |
| 97 | if not args.skip_stats: |
| 98 | print("Computing motion statistics from full dataset...") |
| 99 | mean, std = compute_statistics(data_list, args.num_workers) |
| 100 | |
| 101 | # Save statistics |
| 102 | mean_path = os.path.join(args.output_dir, "mean.npy") |
| 103 | std_path = os.path.join(args.output_dir, "std.npy") |
| 104 | np.save(mean_path, mean) |
| 105 | np.save(std_path, std) |
| 106 | print(f"Saved mean to {mean_path}") |
| 107 | print(f"Saved std to {std_path}") |
| 108 | print(f"Mean shape: {mean.shape}, Std shape: {std.shape}") |
| 109 | else: |
| 110 | print("Skipping statistics computation (--skip_stats)") |
| 111 | |
| 112 | # Save processed JSON |
| 113 | output_json = os.path.join(args.output_dir, "ViMoGen-228K_train.json") |
| 114 | print(f"Saving processed JSON to {output_json}...") |
| 115 | with open(output_json, 'w') as f: |
| 116 | json.dump(data_list, f, indent=2) |
| 117 | print(f"Done! Processed {len(data_list)} entries.") |
| 118 | |
| 119 | |
| 120 | if __name__ == "__main__": |
no test coverage detected