Commit Graph

2 Commits

Author SHA1 Message Date
CarolinePascal 7035ecf9b2 fix(datasets): dequantize depth video frames when recomputing stats
Depth video stats were computed on raw 12-bit codec values, leaving them in
codec space instead of the recorded depth unit. Dequantize decoded frames via
the feature's depth encoder config (matching DatasetReader) so recomputed stats
match record-time stats.

Also fix the SLURM example: the --skip-image-video flag was inverted (0 skipped
visual stats), and add a --video-backend option so pyav can be used when
torchcodec fails to load locally.
2026-07-10 16:25:40 +02:00
CarolinePascal 7bee7fb9e3 feat(datasets): distribute stats recomputation across SLURM workers
Expose the shardable unit of work behind recompute_stats: compute_dataset_episode_stats
computes per-episode stats for an episode subset, and aggregate_episode_stats merges the
concatenated shards (count-weighted) and writes stats.json. recompute_stats now composes
these, so single-process behavior is unchanged.

Add examples/dataset/slurm_recompute_stats.py, a datatrove compute/aggregate driver that
shards episodes across workers and is read-only safe (reference-copies the source when
--new-root is given). Most useful for the expensive image/video stats path.
2026-07-10 14:47:18 +02:00