osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
2223 lines
76 KiB
Rust
2223 lines
76 KiB
Rust
//! JEPA Data Pipeline
|
||
//!
|
||
//! Complete data pipeline for I-JEPA / V-JEPA training:
|
||
//! - ImageRecord: decoded image storage
|
||
//! - MultiScaleRandomCrop: bilinear-resampled random crop augmentation
|
||
//! - RandomHorizontalFlip: stochastic horizontal flip
|
||
//! - JepaAugmentationPipeline: chained augmentation with ImageNet normalization
|
||
//! - JepaDataConfig: pipeline configuration
|
||
//! - InMemoryShard: in-memory image shard (CPU-testable without filesystem)
|
||
//! - JepaBatch: one ready-to-train batch with masks
|
||
//! - JepaDataPipeline: streaming batched iterator over shards
|
||
//! - DatasetStats: aggregate statistics
|
||
//! - WebDatasetShard: filesystem shard descriptor (stub)
|
||
|
||
// ============================================================================
|
||
// ImageRecord
|
||
// ============================================================================
|
||
|
||
/// A single decoded image stored in host memory.
|
||
#[derive(Debug, Clone)]
|
||
pub struct ImageRecord {
|
||
/// Pixel data in HWC layout, values in [0, 1].
|
||
pub pixels: Vec<f32>,
|
||
pub width: usize,
|
||
pub height: usize,
|
||
pub channels: usize,
|
||
/// Optional class label.
|
||
pub label: Option<usize>,
|
||
/// Dataset key / file name.
|
||
pub key: String,
|
||
}
|
||
|
||
impl ImageRecord {
|
||
/// Pixel value at (row, col, channel). Returns 0.0 if out of range.
|
||
#[inline]
|
||
fn get(&self, row: usize, col: usize, c: usize) -> f32 {
|
||
if row < self.height && col < self.width && c < self.channels {
|
||
self.pixels[row * self.width * self.channels + col * self.channels + c]
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// LCG PRNG helper
|
||
// ============================================================================
|
||
|
||
#[inline(always)]
|
||
fn lcg_step(state: &mut u64) -> u64 {
|
||
*state = state
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
*state
|
||
}
|
||
|
||
/// Sample a f32 in [0, 1) from the LCG state.
|
||
#[inline(always)]
|
||
fn lcg_f32(state: &mut u64) -> f32 {
|
||
(lcg_step(state) >> 11) as f32 / (1u64 << 53) as f32
|
||
}
|
||
|
||
/// Sample a usize in [0, n).
|
||
#[inline(always)]
|
||
fn lcg_usize(state: &mut u64, n: usize) -> usize {
|
||
if n == 0 {
|
||
return 0;
|
||
}
|
||
(lcg_step(state) % n as u64) as usize
|
||
}
|
||
|
||
// ============================================================================
|
||
// MultiScaleRandomCrop
|
||
// ============================================================================
|
||
|
||
/// Random crop followed by bilinear resize to a fixed target size.
|
||
///
|
||
/// Mimics the torchvision `RandomResizedCrop` transform used in I-JEPA:
|
||
/// - Sample crop area as `scale * total_area` of the source image.
|
||
/// - Sample aspect ratio from `ratio_range`.
|
||
/// - Bilinear-resample to `(target_size, target_size, channels)`.
|
||
#[derive(Debug, Clone)]
|
||
pub struct MultiScaleRandomCrop {
|
||
pub target_size: usize,
|
||
pub scale_range: (f32, f32),
|
||
pub ratio_range: (f32, f32),
|
||
}
|
||
|
||
impl MultiScaleRandomCrop {
|
||
pub fn new(target_size: usize, scale_range: (f32, f32), ratio_range: (f32, f32)) -> Self {
|
||
Self {
|
||
target_size,
|
||
scale_range,
|
||
ratio_range,
|
||
}
|
||
}
|
||
|
||
/// Crop and bilinear-resize a single `ImageRecord`.
|
||
///
|
||
/// `seed` seeds the LCG PRNG so results are deterministic per sample.
|
||
/// Returns a flat `[target_size, target_size, channels]` buffer.
|
||
pub fn crop_and_resize(&self, image: &ImageRecord, seed: u64) -> Vec<f32> {
|
||
let mut lcg = seed.wrapping_mul(2654435761).wrapping_add(1);
|
||
|
||
let src_h = image.height;
|
||
let src_w = image.width;
|
||
let c = image.channels;
|
||
let ts = self.target_size;
|
||
|
||
// Clamp scale range to (0, 1]
|
||
let scale_lo = self.scale_range.0.max(1e-6_f32).min(1.0);
|
||
let scale_hi = self.scale_range.1.max(scale_lo).min(1.0);
|
||
let ratio_lo = self.ratio_range.0.max(1e-4_f32);
|
||
let ratio_hi = self.ratio_range.1.max(ratio_lo);
|
||
|
||
// Try up to 10 attempts to find a valid crop; fall back to centre crop.
|
||
let mut crop_top = 0usize;
|
||
let mut crop_left = 0usize;
|
||
let mut crop_h = src_h;
|
||
let mut crop_w = src_w;
|
||
|
||
for _ in 0..10 {
|
||
let scale = scale_lo + lcg_f32(&mut lcg) * (scale_hi - scale_lo);
|
||
let ratio = ratio_lo + lcg_f32(&mut lcg) * (ratio_hi - ratio_lo);
|
||
|
||
let area = (src_h * src_w) as f32 * scale;
|
||
let h_f = (area / ratio).sqrt();
|
||
let w_f = (area * ratio).sqrt();
|
||
let h = h_f.round() as usize;
|
||
let w = w_f.round() as usize;
|
||
|
||
if h == 0 || w == 0 || h > src_h || w > src_w {
|
||
continue;
|
||
}
|
||
|
||
let top = lcg_usize(&mut lcg, src_h - h + 1);
|
||
let left = lcg_usize(&mut lcg, src_w - w + 1);
|
||
|
||
crop_top = top;
|
||
crop_left = left;
|
||
crop_h = h;
|
||
crop_w = w;
|
||
break;
|
||
}
|
||
|
||
// Bilinear resample crop → target_size × target_size
|
||
let mut out = vec![0.0f32; ts * ts * c];
|
||
for oy in 0..ts {
|
||
for ox in 0..ts {
|
||
// Map output pixel to source pixel in the crop region
|
||
// Using half-pixel convention for better edge behaviour
|
||
let src_y_f = (oy as f32 + 0.5) / ts as f32 * crop_h as f32 - 0.5 + crop_top as f32;
|
||
let src_x_f =
|
||
(ox as f32 + 0.5) / ts as f32 * crop_w as f32 - 0.5 + crop_left as f32;
|
||
|
||
let y0 = src_y_f.floor() as isize;
|
||
let x0 = src_x_f.floor() as isize;
|
||
let y1 = y0 + 1;
|
||
let x1 = x0 + 1;
|
||
|
||
let dy = src_y_f - y0 as f32;
|
||
let dx = src_x_f - x0 as f32;
|
||
|
||
// Clamp coordinates
|
||
let y0c = y0.clamp(0, src_h as isize - 1) as usize;
|
||
let y1c = y1.clamp(0, src_h as isize - 1) as usize;
|
||
let x0c = x0.clamp(0, src_w as isize - 1) as usize;
|
||
let x1c = x1.clamp(0, src_w as isize - 1) as usize;
|
||
|
||
let base = (oy * ts + ox) * c;
|
||
for ch in 0..c {
|
||
let p00 = image.get(y0c, x0c, ch);
|
||
let p01 = image.get(y0c, x1c, ch);
|
||
let p10 = image.get(y1c, x0c, ch);
|
||
let p11 = image.get(y1c, x1c, ch);
|
||
out[base + ch] = p00 * (1.0 - dy) * (1.0 - dx)
|
||
+ p01 * (1.0 - dy) * dx
|
||
+ p10 * dy * (1.0 - dx)
|
||
+ p11 * dy * dx;
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// RandomHorizontalFlip
|
||
// ============================================================================
|
||
|
||
/// Randomly flip an image left-right with probability `prob`.
|
||
#[derive(Debug, Clone)]
|
||
pub struct RandomHorizontalFlip {
|
||
pub prob: f32,
|
||
}
|
||
|
||
impl RandomHorizontalFlip {
|
||
pub fn new(prob: f32) -> Self {
|
||
Self { prob }
|
||
}
|
||
|
||
/// Apply (or not) horizontal flip.
|
||
///
|
||
/// `pixels` must be in HWC layout with `w * channels` stride per row.
|
||
pub fn apply(
|
||
&self,
|
||
pixels: &[f32],
|
||
w: usize,
|
||
h: usize,
|
||
channels: usize,
|
||
seed: u64,
|
||
) -> Vec<f32> {
|
||
let mut lcg = seed
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
let r = lcg_f32(&mut lcg);
|
||
|
||
if r >= self.prob {
|
||
return pixels.to_vec();
|
||
}
|
||
|
||
let mut out = pixels.to_vec();
|
||
for row in 0..h {
|
||
for col in 0..w / 2 {
|
||
let mirror = w - 1 - col;
|
||
for c in 0..channels {
|
||
let a = row * w * channels + col * channels + c;
|
||
let b = row * w * channels + mirror * channels + c;
|
||
out.swap(a, b);
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// JepaAugmentationPipeline
|
||
// ============================================================================
|
||
|
||
/// Chained augmentation: MultiScaleRandomCrop → RandomHorizontalFlip → normalize.
|
||
#[derive(Debug, Clone)]
|
||
pub struct JepaAugmentationPipeline {
|
||
pub crop: MultiScaleRandomCrop,
|
||
pub flip: RandomHorizontalFlip,
|
||
pub mean: [f32; 3],
|
||
pub std: [f32; 3],
|
||
}
|
||
|
||
impl JepaAugmentationPipeline {
|
||
/// Build from a `JepaDataConfig`.
|
||
pub fn from_config(config: &JepaDataConfig) -> Self {
|
||
let mean = if config.imagenet_normalize {
|
||
[0.485, 0.456, 0.406]
|
||
} else {
|
||
[0.0, 0.0, 0.0]
|
||
};
|
||
let std = if config.imagenet_normalize {
|
||
[0.229, 0.224, 0.225]
|
||
} else {
|
||
[1.0, 1.0, 1.0]
|
||
};
|
||
Self {
|
||
crop: MultiScaleRandomCrop::new(
|
||
config.image_size,
|
||
config.scale_range,
|
||
config.ratio_range,
|
||
),
|
||
flip: RandomHorizontalFlip::new(if config.use_horizontal_flip { 0.5 } else { 0.0 }),
|
||
mean,
|
||
std,
|
||
}
|
||
}
|
||
|
||
/// Per-channel ImageNet normalization in-place.
|
||
pub fn normalize(pixels: &mut Vec<f32>, mean: [f32; 3], std: [f32; 3]) {
|
||
let c = mean.len(); // 3
|
||
for chunk in pixels.chunks_mut(c) {
|
||
for (i, v) in chunk.iter_mut().enumerate() {
|
||
let ch = i % c;
|
||
let s = std[ch].max(1e-7_f32);
|
||
*v = (*v - mean[ch]) / s;
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Run the full augmentation pipeline on one `ImageRecord`.
|
||
///
|
||
/// Returns a flat `[H_out * W_out * C]` buffer (C=channels of the source image).
|
||
pub fn process(&self, image: &ImageRecord, seed: u64) -> Vec<f32> {
|
||
let ts = self.crop.target_size;
|
||
let c = image.channels;
|
||
|
||
// 1. Multi-scale random crop → [ts, ts, c]
|
||
let mut cropped = self.crop.crop_and_resize(image, seed);
|
||
|
||
// 2. Horizontal flip (use a derived seed so crop and flip are independent)
|
||
let flip_seed = seed
|
||
.wrapping_mul(2654435761)
|
||
.wrapping_add(1442695040888963407);
|
||
let flipped = self.flip.apply(&cropped, ts, ts, c, flip_seed);
|
||
cropped = flipped;
|
||
|
||
// 3. Normalize per-channel
|
||
Self::normalize(&mut cropped, self.mean, self.std);
|
||
|
||
cropped
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// JepaDataConfig
|
||
// ============================================================================
|
||
|
||
/// Top-level configuration for the JEPA data pipeline.
|
||
#[derive(Debug, Clone)]
|
||
pub struct JepaDataConfig {
|
||
/// Target image size after crop (e.g. 224).
|
||
pub image_size: usize,
|
||
/// Patch size used for masking (e.g. 16 → 14×14 grid on 224px images).
|
||
pub patch_size: usize,
|
||
/// Number of images per batch.
|
||
pub batch_size: usize,
|
||
/// Concurrency hint (not used in CPU reference implementation).
|
||
pub num_workers: usize,
|
||
/// Paths to WebDataset shards (.tar files).
|
||
pub shard_paths: Vec<String>,
|
||
/// Scale range for `MultiScaleRandomCrop` (area fraction).
|
||
pub scale_range: (f32, f32),
|
||
/// Aspect ratio range for `MultiScaleRandomCrop`.
|
||
pub ratio_range: (f32, f32),
|
||
/// Whether to apply random horizontal flip.
|
||
pub use_horizontal_flip: bool,
|
||
/// Whether to apply ImageNet mean/std normalization.
|
||
pub imagenet_normalize: bool,
|
||
}
|
||
|
||
impl Default for JepaDataConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
image_size: 224,
|
||
patch_size: 16,
|
||
batch_size: 1024,
|
||
num_workers: 8,
|
||
shard_paths: Vec::new(),
|
||
scale_range: (0.2, 1.0),
|
||
ratio_range: (0.75, 1.33),
|
||
use_horizontal_flip: true,
|
||
imagenet_normalize: true,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl JepaDataConfig {
|
||
/// Grid side length in patches.
|
||
pub fn grid_size(&self) -> usize {
|
||
self.image_size / self.patch_size
|
||
}
|
||
|
||
/// Total patch count per image.
|
||
pub fn num_patches(&self) -> usize {
|
||
let g = self.grid_size();
|
||
g * g
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// InMemoryShard
|
||
// ============================================================================
|
||
|
||
/// A shard of images held entirely in host memory.
|
||
///
|
||
/// Used for unit testing without a real filesystem.
|
||
#[derive(Debug, Clone)]
|
||
pub struct InMemoryShard {
|
||
pub records: Vec<ImageRecord>,
|
||
pub shard_id: usize,
|
||
}
|
||
|
||
impl InMemoryShard {
|
||
/// Create `n` synthetic images of size `image_size × image_size × 3`
|
||
/// with LCG-generated random pixels in `[0, 1]` and sequential labels.
|
||
pub fn synthetic(n: usize, image_size: usize, shard_id: usize) -> Self {
|
||
let mut records = Vec::with_capacity(n);
|
||
let mut lcg: u64 = (shard_id as u64)
|
||
.wrapping_mul(2654435761)
|
||
.wrapping_add(1442695040888963407);
|
||
|
||
for i in 0..n {
|
||
let pixel_count = image_size * image_size * 3;
|
||
let pixels: Vec<f32> = (0..pixel_count).map(|_| lcg_f32(&mut lcg)).collect();
|
||
|
||
records.push(ImageRecord {
|
||
pixels,
|
||
width: image_size,
|
||
height: image_size,
|
||
channels: 3,
|
||
label: Some(i % 1000),
|
||
key: format!("shard{shard_id}/img{i:06}"),
|
||
});
|
||
}
|
||
|
||
Self { records, shard_id }
|
||
}
|
||
|
||
/// Number of records in this shard.
|
||
pub fn len(&self) -> usize {
|
||
self.records.len()
|
||
}
|
||
|
||
/// True if the shard is empty.
|
||
pub fn is_empty(&self) -> bool {
|
||
self.records.is_empty()
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// JepaBatch
|
||
// ============================================================================
|
||
|
||
/// One training batch ready for JEPA forward + loss.
|
||
#[derive(Debug, Clone)]
|
||
pub struct JepaBatch {
|
||
/// Augmented image pixels for each sample, flattened to `[H*W*C]`.
|
||
pub images: Vec<Vec<f32>>,
|
||
/// Context patch indices per sample `[batch_size][n_context]`.
|
||
pub context_patch_indices: Vec<Vec<usize>>,
|
||
/// Target block patch indices per sample `[batch_size][num_blocks][n_patches]`.
|
||
pub target_patch_groups: Vec<Vec<Vec<usize>>>,
|
||
/// Union of all target patch indices per sample `[batch_size][n_total_target]`.
|
||
pub all_target_indices: Vec<Vec<usize>>,
|
||
/// Class labels (if available) per sample.
|
||
pub labels: Vec<Option<usize>>,
|
||
/// Monotonically increasing batch ID within the epoch.
|
||
pub batch_id: usize,
|
||
}
|
||
|
||
impl JepaBatch {
|
||
/// Actual number of samples in this batch (may be less than config.batch_size at epoch end).
|
||
pub fn batch_size(&self) -> usize {
|
||
self.images.len()
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// JepaDataPipeline
|
||
// ============================================================================
|
||
|
||
/// Streaming data pipeline: iterates over shards, applies augmentation, and
|
||
/// assembles `JepaBatch` instances with JEPA block masks.
|
||
pub struct JepaDataPipeline {
|
||
pub config: JepaDataConfig,
|
||
augmentation: JepaAugmentationPipeline,
|
||
mask_strategy: super::jepa::BlockMaskStrategy,
|
||
shards: Vec<InMemoryShard>,
|
||
/// Order in which shards are visited this epoch (shuffled each epoch).
|
||
shard_order: Vec<usize>,
|
||
/// Index into `shard_order` for the currently active shard.
|
||
current_shard: usize,
|
||
/// Index within the current shard's records.
|
||
current_idx: usize,
|
||
/// Current epoch (starts at 0, increments on `reset_epoch`).
|
||
epoch: usize,
|
||
/// Monotonic batch counter (resets on `reset_epoch`).
|
||
batch_counter: usize,
|
||
/// LCG seed for shard shuffling.
|
||
shuffle_lcg: u64,
|
||
}
|
||
|
||
impl JepaDataPipeline {
|
||
/// Build a new pipeline.
|
||
pub fn new(config: JepaDataConfig, shards: Vec<InMemoryShard>) -> Self {
|
||
let augmentation = JepaAugmentationPipeline::from_config(&config);
|
||
let mask_strategy = super::jepa::BlockMaskStrategy::default_ijepa();
|
||
let n = shards.len();
|
||
let shard_order: Vec<usize> = (0..n).collect();
|
||
Self {
|
||
config,
|
||
augmentation,
|
||
mask_strategy,
|
||
shards,
|
||
shard_order,
|
||
current_shard: 0,
|
||
current_idx: 0,
|
||
epoch: 0,
|
||
batch_counter: 0,
|
||
shuffle_lcg: 0xdeadbeef_cafebabe_u64,
|
||
}
|
||
}
|
||
|
||
/// Total images across all shards.
|
||
fn total_images(&self) -> usize {
|
||
self.shards.iter().map(|s| s.len()).sum()
|
||
}
|
||
|
||
/// Number of full (and potentially partial) batches per epoch.
|
||
pub fn num_batches_per_epoch(&self) -> usize {
|
||
let total = self.total_images();
|
||
if total == 0 || self.config.batch_size == 0 {
|
||
return 0;
|
||
}
|
||
(total + self.config.batch_size - 1) / self.config.batch_size
|
||
}
|
||
|
||
/// Current epoch number.
|
||
pub fn epoch(&self) -> usize {
|
||
self.epoch
|
||
}
|
||
|
||
/// Advance the shard pointer, skipping empty shards.
|
||
/// Returns false if all shards for this epoch have been consumed.
|
||
fn advance_shard(&mut self) -> bool {
|
||
self.current_shard += 1;
|
||
self.current_idx = 0;
|
||
self.current_shard < self.shard_order.len()
|
||
}
|
||
|
||
/// Reference to the currently active shard, or None if epoch is done.
|
||
fn active_shard(&self) -> Option<&InMemoryShard> {
|
||
let si = self.shard_order.get(self.current_shard)?;
|
||
self.shards.get(*si)
|
||
}
|
||
|
||
/// Pull the next `ImageRecord` from the stream, advancing shards as needed.
|
||
fn next_record(&mut self) -> Option<&ImageRecord> {
|
||
loop {
|
||
let shard = match self.active_shard() {
|
||
Some(s) => s,
|
||
None => return None,
|
||
};
|
||
|
||
if self.current_idx < shard.len() {
|
||
// This looks like a borrow issue but we borrow immutably here.
|
||
// We will return a raw pointer trick via direct indexing below.
|
||
break;
|
||
} else {
|
||
if !self.advance_shard() {
|
||
return None;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Re-resolve to get the record (needed because we can't hold shard ref through advance_shard).
|
||
let si = *self.shard_order.get(self.current_shard)?;
|
||
let shard = self.shards.get(si)?;
|
||
let rec = shard.records.get(self.current_idx)?;
|
||
self.current_idx += 1;
|
||
Some(rec)
|
||
}
|
||
|
||
/// Yield the next batch, or `None` when the epoch is exhausted.
|
||
pub fn next_batch(&mut self) -> Option<JepaBatch> {
|
||
let batch_size = self.config.batch_size;
|
||
let grid = self.config.grid_size();
|
||
let batch_id = self.batch_counter;
|
||
|
||
// Collect up to `batch_size` records.
|
||
// We need to collect references first, then process.
|
||
// Because next_record borrows &mut self, we collect indices instead.
|
||
let mut sample_indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); // (shard_order_idx, record_idx)
|
||
|
||
for _ in 0..batch_size {
|
||
loop {
|
||
let shard = match self.active_shard() {
|
||
Some(s) => s,
|
||
None => break,
|
||
};
|
||
if self.current_idx < shard.len() {
|
||
let si = self.shard_order[self.current_shard];
|
||
sample_indices.push((si, self.current_idx));
|
||
self.current_idx += 1;
|
||
break;
|
||
} else {
|
||
if !self.advance_shard() {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Check if we are done.
|
||
if self.active_shard().is_none()
|
||
&& sample_indices.len() < batch_size
|
||
&& !sample_indices.is_empty()
|
||
{
|
||
// Last partial batch — collect what we have.
|
||
break;
|
||
}
|
||
|
||
if sample_indices.len() == batch_size {
|
||
break;
|
||
}
|
||
}
|
||
|
||
if sample_indices.is_empty() {
|
||
return None;
|
||
}
|
||
|
||
let actual_bs = sample_indices.len();
|
||
let mut images = Vec::with_capacity(actual_bs);
|
||
let mut context_patch_indices = Vec::with_capacity(actual_bs);
|
||
let mut target_patch_groups = Vec::with_capacity(actual_bs);
|
||
let mut all_target_indices_out = Vec::with_capacity(actual_bs);
|
||
let mut labels = Vec::with_capacity(actual_bs);
|
||
|
||
for (sample_i, (shard_idx, record_idx)) in sample_indices.iter().enumerate() {
|
||
let record = &self.shards[*shard_idx].records[*record_idx];
|
||
|
||
// Build a per-sample seed from batch_id, sample_i, and shard/record position.
|
||
let seed: u64 = (batch_id as u64)
|
||
.wrapping_mul(1000003)
|
||
.wrapping_add(*shard_idx as u64)
|
||
.wrapping_mul(999983)
|
||
.wrapping_add(*record_idx as u64)
|
||
.wrapping_add(sample_i as u64);
|
||
|
||
// Augment image.
|
||
let aug = self.augmentation.process(record, seed);
|
||
images.push(aug);
|
||
labels.push(record.label);
|
||
|
||
// Generate JEPA block masks.
|
||
let mask = self.mask_strategy.generate(grid, grid, seed);
|
||
context_patch_indices.push(mask.context_indices);
|
||
target_patch_groups.push(mask.target_blocks);
|
||
all_target_indices_out.push(mask.all_target_indices);
|
||
}
|
||
|
||
self.batch_counter += 1;
|
||
|
||
Some(JepaBatch {
|
||
images,
|
||
context_patch_indices,
|
||
target_patch_groups,
|
||
all_target_indices: all_target_indices_out,
|
||
labels,
|
||
batch_id,
|
||
})
|
||
}
|
||
|
||
/// Reset for the next epoch: shuffle shard order and rewind position counters.
|
||
pub fn reset_epoch(&mut self) {
|
||
self.epoch += 1;
|
||
self.batch_counter = 0;
|
||
self.current_shard = 0;
|
||
self.current_idx = 0;
|
||
|
||
// Fisher-Yates shuffle of shard order using internal LCG.
|
||
let n = self.shard_order.len();
|
||
for i in 0..n {
|
||
self.shard_order[i] = i;
|
||
}
|
||
for i in 0..n {
|
||
let j = i + lcg_usize(&mut self.shuffle_lcg, n - i);
|
||
self.shard_order.swap(i, j);
|
||
}
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// DatasetStats
|
||
// ============================================================================
|
||
|
||
/// Aggregate statistics about the dataset and masking behaviour.
|
||
#[derive(Debug, Clone)]
|
||
pub struct DatasetStats {
|
||
pub total_images: usize,
|
||
pub num_shards: usize,
|
||
pub images_per_shard: Vec<usize>,
|
||
/// Mean context patch count across the first sampled batch.
|
||
pub avg_context_patches: f64,
|
||
/// Mean target patch count across the first sampled batch.
|
||
pub avg_target_patches: f64,
|
||
/// Fraction of total patches that are target patches (mask_efficiency).
|
||
pub mask_efficiency: f64,
|
||
}
|
||
|
||
impl DatasetStats {
|
||
/// Compute statistics by sampling one batch from a separate mask strategy instance.
|
||
pub fn compute(pipeline: &JepaDataPipeline) -> Self {
|
||
let total_images: usize = pipeline.shards.iter().map(|s| s.len()).sum();
|
||
let num_shards = pipeline.shards.len();
|
||
let images_per_shard: Vec<usize> = pipeline.shards.iter().map(|s| s.len()).collect();
|
||
|
||
let total_patches = pipeline.config.num_patches();
|
||
let grid = pipeline.config.grid_size();
|
||
|
||
// Sample a small number of masks to estimate average patch counts.
|
||
let sample_size = 32.min(total_images.max(1));
|
||
let mut mask_strategy = super::jepa::BlockMaskStrategy::default_ijepa();
|
||
let masks = mask_strategy.generate_batch(sample_size, grid, grid);
|
||
|
||
let avg_context = if masks.is_empty() {
|
||
0.0
|
||
} else {
|
||
masks
|
||
.iter()
|
||
.map(|m| m.context_indices.len() as f64)
|
||
.sum::<f64>()
|
||
/ masks.len() as f64
|
||
};
|
||
|
||
let avg_target = if masks.is_empty() {
|
||
0.0
|
||
} else {
|
||
masks
|
||
.iter()
|
||
.map(|m| m.all_target_indices.len() as f64)
|
||
.sum::<f64>()
|
||
/ masks.len() as f64
|
||
};
|
||
|
||
let mask_efficiency = if total_patches == 0 {
|
||
0.0
|
||
} else {
|
||
avg_target / total_patches as f64
|
||
};
|
||
|
||
Self {
|
||
total_images,
|
||
num_shards,
|
||
images_per_shard,
|
||
avg_context_patches: avg_context,
|
||
avg_target_patches: avg_target,
|
||
mask_efficiency,
|
||
}
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// WebDatasetShard
|
||
// ============================================================================
|
||
|
||
/// Filesystem shard descriptor for WebDataset-format `.tar` archives.
|
||
///
|
||
/// Use [`WebDatasetShard::load`] to read the actual archive from disk
|
||
/// (via [`read_webdataset_shard`]); [`WebDatasetShard::to_in_memory`]
|
||
/// generates synthetic data and exists for tests that need a shard
|
||
/// without touching the filesystem.
|
||
#[derive(Debug, Clone)]
|
||
pub struct WebDatasetShard {
|
||
pub path: String,
|
||
pub num_records: usize,
|
||
pub shard_id: usize,
|
||
pub compressed: bool,
|
||
}
|
||
|
||
impl WebDatasetShard {
|
||
/// Create a descriptor pointing at `path`.
|
||
///
|
||
/// `num_records` is used both for progress estimation and for synthetic
|
||
/// data generation in `to_in_memory`.
|
||
pub fn new(path: &str, num_records: usize, shard_id: usize) -> Self {
|
||
Self {
|
||
path: path.to_owned(),
|
||
num_records,
|
||
shard_id,
|
||
compressed: path.ends_with(".tar.gz") || path.ends_with(".tgz"),
|
||
}
|
||
}
|
||
|
||
/// Read the shard's `.tar` (or gzip-compressed `.tar.gz`/`.tgz`)
|
||
/// archive from disk and decode its records into an [`InMemoryShard`].
|
||
pub fn load(&self) -> Result<InMemoryShard, String> {
|
||
let (raw_records, _stats) = read_webdataset_shard(std::path::Path::new(&self.path))?;
|
||
let records: Vec<ImageRecord> = raw_records
|
||
.into_iter()
|
||
.map(webdataset_record_to_image)
|
||
.collect();
|
||
Ok(InMemoryShard {
|
||
records,
|
||
shard_id: self.shard_id,
|
||
})
|
||
}
|
||
|
||
/// Generate a synthetic `InMemoryShard` with `num_records` records of
|
||
/// size `image_size × image_size × 3`. Used for testing without a
|
||
/// real filesystem — use [`WebDatasetShard::load`] for real data.
|
||
pub fn to_in_memory(&self, image_size: usize) -> InMemoryShard {
|
||
InMemoryShard::synthetic(self.num_records, image_size, self.shard_id)
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// WebDatasetRecord
|
||
// ============================================================================
|
||
|
||
/// A record decoded from a WebDataset shard (one image + optional label).
|
||
#[derive(Debug, Clone)]
|
||
pub struct WebDatasetRecord {
|
||
/// Numeric prefix, e.g. "000042".
|
||
pub key: String,
|
||
/// Raw JPEG/PNG bytes (not decoded to pixels).
|
||
pub image_bytes: Vec<u8>,
|
||
/// Parsed from .cls file.
|
||
pub label: Option<usize>,
|
||
/// "jpg" or "png".
|
||
pub extension: String,
|
||
}
|
||
|
||
// ============================================================================
|
||
// ShardLoadStats
|
||
// ============================================================================
|
||
|
||
/// Statistics for a loaded shard.
|
||
#[derive(Debug, Clone)]
|
||
pub struct ShardLoadStats {
|
||
pub shard_id: usize,
|
||
pub records_loaded: usize,
|
||
pub bytes_read: u64,
|
||
pub load_duration_ms: u64,
|
||
}
|
||
|
||
// ============================================================================
|
||
// ShuffleBuffer
|
||
// ============================================================================
|
||
|
||
/// Fixed-size ring buffer for dataset shuffling.
|
||
/// Holds up to `capacity` records; reservoir sampling on insert.
|
||
pub struct ShuffleBuffer {
|
||
capacity: usize,
|
||
buffer: Vec<ImageRecord>,
|
||
rng_state: u64,
|
||
}
|
||
|
||
impl ShuffleBuffer {
|
||
/// Create a new shuffle buffer seeded from `capacity`.
|
||
pub fn new(capacity: usize) -> Self {
|
||
let rng_state = (capacity as u64).wrapping_mul(6364136223846793005);
|
||
Self {
|
||
capacity,
|
||
buffer: Vec::new(),
|
||
rng_state,
|
||
}
|
||
}
|
||
|
||
/// Insert a record. If the buffer is not full, push it; otherwise replace
|
||
/// a random slot using the internal LCG.
|
||
pub fn push(&mut self, record: ImageRecord) {
|
||
if self.buffer.len() < self.capacity {
|
||
self.buffer.push(record);
|
||
} else if self.capacity > 0 {
|
||
let slot = lcg_usize(&mut self.rng_state, self.capacity);
|
||
self.buffer[slot] = record;
|
||
}
|
||
}
|
||
|
||
/// Remove and return the first `n` items (or all items if fewer than `n` remain).
|
||
pub fn drain_batch(&mut self, n: usize) -> Vec<ImageRecord> {
|
||
let take = n.min(self.buffer.len());
|
||
self.buffer.drain(..take).collect()
|
||
}
|
||
|
||
/// Number of records currently in the buffer.
|
||
pub fn len(&self) -> usize {
|
||
self.buffer.len()
|
||
}
|
||
|
||
/// True when `len() >= capacity`.
|
||
pub fn is_full(&self) -> bool {
|
||
self.len() >= self.capacity
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Tar parsing helpers (stdlib only)
|
||
// ============================================================================
|
||
|
||
/// Parse a null-terminated octal ASCII string from a fixed-width tar field.
|
||
fn parse_octal(field: &[u8]) -> u64 {
|
||
// Trim leading/trailing null bytes and spaces
|
||
let s = field
|
||
.iter()
|
||
.take_while(|&&b| b != 0 && b != b' ')
|
||
.copied()
|
||
.collect::<Vec<u8>>();
|
||
let s = std::str::from_utf8(&s).unwrap_or("0");
|
||
u64::from_str_radix(s.trim(), 8).unwrap_or(0)
|
||
}
|
||
|
||
/// Parse all WebDataset records from an in-memory tar byte slice.
|
||
///
|
||
/// A tar archive is a sequence of 512-byte blocks:
|
||
/// - Header block: filename (bytes 0..100), size (bytes 124..136, octal),
|
||
/// typeflag (byte 156, '0' or '\0' = regular file, '5' = directory).
|
||
/// - Data blocks: ceil(size / 512) × 512 bytes follow the header.
|
||
/// - End-of-archive: two consecutive all-zero blocks.
|
||
pub fn parse_tar_bytes(data: &[u8]) -> Vec<WebDatasetRecord> {
|
||
// Accumulate raw file entries keyed by stem (numeric prefix).
|
||
// Values: (image_bytes, image_ext, label_bytes)
|
||
let mut image_map: std::collections::HashMap<String, (Vec<u8>, String)> =
|
||
std::collections::HashMap::new();
|
||
let mut label_map: std::collections::HashMap<String, Vec<u8>> =
|
||
std::collections::HashMap::new();
|
||
|
||
let mut offset = 0usize;
|
||
|
||
loop {
|
||
// Need at least one header block.
|
||
if offset + 512 > data.len() {
|
||
break;
|
||
}
|
||
|
||
let header = &data[offset..offset + 512];
|
||
offset += 512;
|
||
|
||
// End-of-archive: two consecutive zero blocks. A zero block starts here.
|
||
if header.iter().all(|&b| b == 0) {
|
||
break;
|
||
}
|
||
|
||
// Parse filename (bytes 0..100, null-terminated).
|
||
let fname_raw = &header[0..100];
|
||
let fname_len = fname_raw.iter().position(|&b| b == 0).unwrap_or(100);
|
||
let filename = match std::str::from_utf8(&fname_raw[..fname_len]) {
|
||
Ok(s) => s.trim_matches('/').to_owned(),
|
||
Err(_) => {
|
||
// Skip unreadable header.
|
||
continue;
|
||
}
|
||
};
|
||
|
||
// Parse typeflag (byte 156). '5' = directory, skip data blocks but
|
||
// don't advance data pointer here — size may still be non-zero.
|
||
let typeflag = header[156];
|
||
let is_directory = typeflag == b'5';
|
||
|
||
// Parse size (bytes 124..136, octal ASCII).
|
||
let size = parse_octal(&header[124..136]) as usize;
|
||
|
||
// Skip aligned data blocks for directories and zero-size files.
|
||
let data_blocks = (size + 511) / 512;
|
||
let data_bytes = data_blocks * 512;
|
||
|
||
if is_directory || size == 0 {
|
||
offset += data_bytes;
|
||
continue;
|
||
}
|
||
|
||
// Read `size` bytes of content.
|
||
if offset + data_bytes > data.len() {
|
||
// Truncated archive — stop.
|
||
break;
|
||
}
|
||
let content = data[offset..offset + size].to_vec();
|
||
offset += data_bytes;
|
||
|
||
// Derive extension and key from the filename.
|
||
// Filenames may include a directory prefix; take the basename.
|
||
let basename = filename.rsplit('/').next().unwrap_or(&filename);
|
||
|
||
let (stem, ext) = if let Some(dot) = basename.rfind('.') {
|
||
(&basename[..dot], &basename[dot + 1..])
|
||
} else {
|
||
(basename, "")
|
||
};
|
||
|
||
match ext {
|
||
"jpg" | "jpeg" | "png" => {
|
||
let image_ext = if ext == "jpeg" { "jpg" } else { ext };
|
||
image_map.insert(stem.to_owned(), (content, image_ext.to_owned()));
|
||
}
|
||
"cls" => {
|
||
label_map.insert(stem.to_owned(), content);
|
||
}
|
||
_ => {
|
||
// Unknown extension — ignore.
|
||
}
|
||
}
|
||
}
|
||
|
||
// Merge image + label by key.
|
||
let mut keys: Vec<String> = image_map.keys().cloned().collect();
|
||
keys.sort();
|
||
|
||
keys.into_iter()
|
||
.filter_map(|key| {
|
||
let (image_bytes, extension) = image_map.remove(&key)?;
|
||
let label = label_map.get(&key).and_then(|bytes| {
|
||
let s = std::str::from_utf8(bytes).ok()?.trim().to_owned();
|
||
s.parse::<usize>().ok()
|
||
});
|
||
Some(WebDatasetRecord {
|
||
key,
|
||
image_bytes,
|
||
label,
|
||
extension,
|
||
})
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Read all records from a WebDataset tar shard at `path`.
|
||
///
|
||
/// Returns records in shard order plus load statistics.
|
||
/// Gzip-compressed shards (`.tar.gz`/`.tgz`, detected by the 1F 8B magic
|
||
/// bytes rather than extension) are decompressed transparently.
|
||
pub fn read_webdataset_shard(
|
||
path: &std::path::Path,
|
||
) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
|
||
use std::io::Read;
|
||
use std::time::Instant;
|
||
|
||
let t0 = Instant::now();
|
||
|
||
let mut file =
|
||
std::fs::File::open(path).map_err(|e| format!("cannot open {}: {e}", path.display()))?;
|
||
|
||
let mut data = Vec::new();
|
||
file.read_to_end(&mut data)
|
||
.map_err(|e| format!("read error for {}: {e}", path.display()))?;
|
||
|
||
let bytes_read = data.len() as u64;
|
||
|
||
// Gzip magic: 0x1F 0x8B. Decompress before tar parsing.
|
||
if data.len() >= 2 && data[0] == 0x1f && data[1] == 0x8b {
|
||
let mut decoder = flate2::read::GzDecoder::new(&data[..]);
|
||
let mut decompressed = Vec::new();
|
||
decoder
|
||
.read_to_end(&mut decompressed)
|
||
.map_err(|e| format!("gzip decompression failed for {}: {e}", path.display()))?;
|
||
data = decompressed;
|
||
}
|
||
|
||
let records = parse_tar_bytes(&data);
|
||
let records_loaded = records.len();
|
||
let load_duration_ms = t0.elapsed().as_millis() as u64;
|
||
|
||
// shard_id is unknown at this level; callers should override if needed.
|
||
let stats = ShardLoadStats {
|
||
shard_id: 0,
|
||
records_loaded,
|
||
bytes_read,
|
||
load_duration_ms,
|
||
};
|
||
|
||
Ok((records, stats))
|
||
}
|
||
|
||
/// Decode a `WebDatasetRecord` into an `ImageRecord`.
|
||
///
|
||
/// Pixel values are placeholder 0.5f32 until rtx-vision image decoder is
|
||
/// wired. The function detects JPEG (`\xFF\xD8`) and PNG (`\x89PNG`) magic
|
||
/// bytes to validate the format; other byte sequences also receive the
|
||
/// placeholder.
|
||
fn webdataset_record_to_image(rec: WebDatasetRecord) -> ImageRecord {
|
||
const TARGET_W: usize = 224;
|
||
const TARGET_H: usize = 224;
|
||
const C: usize = 3;
|
||
|
||
#[cfg(feature = "image-decode")]
|
||
{
|
||
use image::io::Reader as ImageReader;
|
||
use std::io::Cursor;
|
||
|
||
if let Ok(reader) = ImageReader::new(Cursor::new(&rec.image_bytes)).with_guessed_format() {
|
||
if let Ok(img) = reader.decode() {
|
||
// Resize to 224×224 and convert to RGB
|
||
let rgb = img
|
||
.resize_exact(
|
||
TARGET_W as u32,
|
||
TARGET_H as u32,
|
||
image::imageops::FilterType::Triangle,
|
||
)
|
||
.into_rgb8();
|
||
|
||
let pixels: Vec<f32> = rgb
|
||
.into_raw()
|
||
.into_iter()
|
||
.map(|p| p as f32 / 255.0)
|
||
.collect();
|
||
|
||
return ImageRecord {
|
||
pixels,
|
||
width: TARGET_W,
|
||
height: TARGET_H,
|
||
channels: C,
|
||
label: rec.label,
|
||
key: rec.key,
|
||
};
|
||
}
|
||
}
|
||
}
|
||
|
||
// Fallback: placeholder (without image-decode feature, or on decode failure)
|
||
let _ = &rec.image_bytes; // suppress unused warning
|
||
ImageRecord {
|
||
pixels: vec![0.5f32; TARGET_W * TARGET_H * C],
|
||
width: TARGET_W,
|
||
height: TARGET_H,
|
||
channels: C,
|
||
label: rec.label,
|
||
key: rec.key,
|
||
}
|
||
}
|
||
|
||
impl JepaDataPipeline {
|
||
/// Load a pipeline from filesystem shards.
|
||
///
|
||
/// `shard_paths` is a list of `.tar` files (uncompressed WebDataset format).
|
||
/// Each file is read, parsed, and converted to an `InMemoryShard`.
|
||
///
|
||
/// Pixel values are placeholder 0.5f32 until rtx-vision image decoder is
|
||
/// wired.
|
||
pub fn from_filesystem(
|
||
config: JepaDataConfig,
|
||
shard_paths: Vec<std::path::PathBuf>,
|
||
) -> Result<Self, String> {
|
||
// Validate all paths exist before loading anything.
|
||
for p in &shard_paths {
|
||
if !p.exists() {
|
||
return Err(format!("shard path does not exist: {}", p.display()));
|
||
}
|
||
}
|
||
|
||
let mut shards: Vec<InMemoryShard> = Vec::with_capacity(shard_paths.len());
|
||
|
||
for (shard_id, path) in shard_paths.iter().enumerate() {
|
||
let (raw_records, _stats) = read_webdataset_shard(path)?;
|
||
|
||
let records: Vec<ImageRecord> = raw_records
|
||
.into_iter()
|
||
.map(webdataset_record_to_image)
|
||
.collect();
|
||
|
||
shards.push(InMemoryShard { records, shard_id });
|
||
}
|
||
|
||
Ok(Self::new(config, shards))
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// HTTP / WebDataset URL loading
|
||
// ============================================================================
|
||
|
||
/// Descriptor for a remote WebDataset shard at a URL.
|
||
#[derive(Debug, Clone)]
|
||
pub struct UrlShardDescriptor {
|
||
/// Full HTTP/HTTPS URL to the .tar shard file
|
||
pub url: String,
|
||
/// Expected number of records (0 = unknown)
|
||
pub expected_records: usize,
|
||
/// Shard index in the dataset
|
||
pub shard_id: usize,
|
||
}
|
||
|
||
impl UrlShardDescriptor {
|
||
pub fn new(url: impl Into<String>, expected_records: usize, shard_id: usize) -> Self {
|
||
Self {
|
||
url: url.into(),
|
||
expected_records,
|
||
shard_id,
|
||
}
|
||
}
|
||
|
||
/// Returns true if this descriptor points to an HTTPS URL (vs plain HTTP)
|
||
pub fn is_https(&self) -> bool {
|
||
self.url.starts_with("https://")
|
||
}
|
||
}
|
||
|
||
/// Like `download_webdataset_shard` but with a configurable timeout.
|
||
/// Used internally and for testing.
|
||
pub(crate) fn download_webdataset_shard_with_timeout(
|
||
url: &str,
|
||
shard_id: usize,
|
||
timeout_secs: u64,
|
||
) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
|
||
let start = std::time::Instant::now();
|
||
|
||
// Validate URL scheme
|
||
if !url.starts_with("http://") && !url.starts_with("https://") {
|
||
return Err(format!(
|
||
"unsupported URL scheme (expected http:// or https://): {url}"
|
||
));
|
||
}
|
||
|
||
// Download using a one-shot Tokio runtime
|
||
let bytes = {
|
||
let rt = tokio::runtime::Runtime::new()
|
||
.map_err(|e| format!("failed to create Tokio runtime: {e}"))?;
|
||
|
||
rt.block_on(async {
|
||
let client = reqwest::Client::builder()
|
||
.timeout(std::time::Duration::from_secs(timeout_secs))
|
||
.build()
|
||
.map_err(|e| format!("reqwest client build failed: {e}"))?;
|
||
|
||
let response = client
|
||
.get(url)
|
||
.send()
|
||
.await
|
||
.map_err(|e| format!("HTTP GET failed for {url}: {e}"))?;
|
||
|
||
if !response.status().is_success() {
|
||
return Err(format!("HTTP {} for {url}", response.status()));
|
||
}
|
||
|
||
response
|
||
.bytes()
|
||
.await
|
||
.map(|b| b.to_vec())
|
||
.map_err(|e| format!("failed to read response body: {e}"))
|
||
})?
|
||
};
|
||
|
||
let load_duration_ms = start.elapsed().as_millis() as u64;
|
||
let bytes_read = bytes.len() as u64;
|
||
|
||
let records = parse_tar_bytes(&bytes);
|
||
let records_loaded = records.len();
|
||
|
||
let stats = ShardLoadStats {
|
||
shard_id,
|
||
records_loaded,
|
||
bytes_read,
|
||
load_duration_ms,
|
||
};
|
||
|
||
Ok((records, stats))
|
||
}
|
||
|
||
/// Download a WebDataset `.tar` shard from an HTTP/HTTPS URL and parse it.
|
||
///
|
||
/// Returns `(records, stats)` on success, `Err(message)` on failure.
|
||
/// Requires `reqwest` (already a workspace dep).
|
||
///
|
||
/// The download is synchronous (blocks via `tokio::runtime::Handle`). For
|
||
/// async callers, use the async variant `download_shard_async` instead.
|
||
pub fn download_webdataset_shard(
|
||
url: &str,
|
||
shard_id: usize,
|
||
) -> Result<(Vec<WebDatasetRecord>, ShardLoadStats), String> {
|
||
download_webdataset_shard_with_timeout(url, shard_id, 300)
|
||
}
|
||
|
||
impl JepaDataPipeline {
|
||
/// Load a pipeline from HTTP/HTTPS WebDataset shard URLs.
|
||
///
|
||
/// Each URL must point to an uncompressed `.tar` file in WebDataset format.
|
||
/// Downloads are sequential (one shard at a time).
|
||
///
|
||
/// Returns `Err` if any URL fails to download or parse.
|
||
pub fn from_urls(config: JepaDataConfig, urls: Vec<String>) -> Result<Self, String> {
|
||
// Validate all URLs before downloading
|
||
for url in &urls {
|
||
if !url.starts_with("http://") && !url.starts_with("https://") {
|
||
return Err(format!(
|
||
"unsupported URL scheme in '{}' (expected http:// or https://)",
|
||
url
|
||
));
|
||
}
|
||
}
|
||
|
||
let mut shards: Vec<InMemoryShard> = Vec::with_capacity(urls.len());
|
||
|
||
for (shard_id, url) in urls.iter().enumerate() {
|
||
let (raw_records, _stats) = download_webdataset_shard(url, shard_id)?;
|
||
|
||
let records: Vec<ImageRecord> = raw_records
|
||
.into_iter()
|
||
.map(webdataset_record_to_image)
|
||
.collect();
|
||
|
||
shards.push(InMemoryShard { records, shard_id });
|
||
}
|
||
|
||
Ok(Self::new(config, shards))
|
||
}
|
||
|
||
/// Validate a list of URLs without downloading.
|
||
///
|
||
/// Returns `Ok(())` if all URLs have valid schemes, `Err` with the first invalid URL.
|
||
pub fn validate_urls(urls: &[String]) -> Result<(), String> {
|
||
for url in urls {
|
||
if !url.starts_with("http://") && !url.starts_with("https://") {
|
||
return Err(format!(
|
||
"invalid URL '{}': must start with http:// or https://",
|
||
url
|
||
));
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Tests
|
||
// ============================================================================
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
// ------------------------------------------------------------------
|
||
// Helper: build a small synthetic ImageRecord
|
||
// ------------------------------------------------------------------
|
||
|
||
fn make_record(w: usize, h: usize, label: Option<usize>) -> ImageRecord {
|
||
let pixels: Vec<f32> = (0..h * w * 3).map(|i| (i % 256) as f32 / 255.0).collect();
|
||
ImageRecord {
|
||
pixels,
|
||
width: w,
|
||
height: h,
|
||
channels: 3,
|
||
label,
|
||
key: "test".into(),
|
||
}
|
||
}
|
||
|
||
fn make_pipeline(
|
||
n_shards: usize,
|
||
imgs_per_shard: usize,
|
||
batch_size: usize,
|
||
) -> JepaDataPipeline {
|
||
let shards: Vec<InMemoryShard> = (0..n_shards)
|
||
.map(|id| InMemoryShard::synthetic(imgs_per_shard, 32, id))
|
||
.collect();
|
||
let config = JepaDataConfig {
|
||
image_size: 32,
|
||
patch_size: 8,
|
||
batch_size,
|
||
num_workers: 1,
|
||
shard_paths: Vec::new(),
|
||
scale_range: (0.5, 1.0),
|
||
ratio_range: (0.75, 1.33),
|
||
use_horizontal_flip: true,
|
||
imagenet_normalize: true,
|
||
};
|
||
JepaDataPipeline::new(config, shards)
|
||
}
|
||
|
||
// ── MultiScaleRandomCrop ──────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_crop_output_size() {
|
||
let crop = MultiScaleRandomCrop::new(64, (0.2, 1.0), (0.75, 1.33));
|
||
let img = make_record(128, 128, None);
|
||
let out = crop.crop_and_resize(&img, 42);
|
||
assert_eq!(out.len(), 64 * 64 * 3, "output must be [target, target, 3]");
|
||
}
|
||
|
||
#[test]
|
||
fn test_crop_output_in_range() {
|
||
// Input pixels are in [0,1], bilinear interp keeps them in [0,1]
|
||
let crop = MultiScaleRandomCrop::new(32, (0.2, 1.0), (0.75, 1.33));
|
||
let img = make_record(64, 64, None);
|
||
let out = crop.crop_and_resize(&img, 7);
|
||
for &v in &out {
|
||
assert!(v >= -1e-6 && v <= 1.0 + 1e-6, "pixel {v} out of [0,1]");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_crop_different_seeds_differ() {
|
||
let crop = MultiScaleRandomCrop::new(32, (0.2, 0.8), (0.75, 1.33));
|
||
let img = make_record(64, 64, None);
|
||
let a = crop.crop_and_resize(&img, 1);
|
||
let b = crop.crop_and_resize(&img, 999);
|
||
assert_ne!(a, b, "different seeds should produce different crops");
|
||
}
|
||
|
||
#[test]
|
||
fn test_crop_same_seed_reproducible() {
|
||
let crop = MultiScaleRandomCrop::new(32, (0.2, 1.0), (0.75, 1.33));
|
||
let img = make_record(64, 64, None);
|
||
let a = crop.crop_and_resize(&img, 42);
|
||
let b = crop.crop_and_resize(&img, 42);
|
||
assert_eq!(a, b, "same seed must be deterministic");
|
||
}
|
||
|
||
#[test]
|
||
fn test_crop_small_image() {
|
||
// Should not panic on a 1×1 image
|
||
let crop = MultiScaleRandomCrop::new(8, (0.5, 1.0), (1.0, 1.0));
|
||
let img = make_record(1, 1, None);
|
||
let out = crop.crop_and_resize(&img, 0);
|
||
assert_eq!(out.len(), 8 * 8 * 3);
|
||
}
|
||
|
||
// ── RandomHorizontalFlip ──────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_flip_is_own_inverse() {
|
||
let flip = RandomHorizontalFlip::new(1.0); // always flip
|
||
let pixels: Vec<f32> = (0..4 * 4 * 3).map(|i| i as f32 / 48.0).collect();
|
||
let flipped = flip.apply(&pixels, 4, 4, 3, 0);
|
||
let restored = flip.apply(&flipped, 4, 4, 3, 0);
|
||
for (a, b) in pixels.iter().zip(restored.iter()) {
|
||
assert!((a - b).abs() < 1e-6, "double flip should be identity");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_flip_prob_zero_never_flips() {
|
||
let flip = RandomHorizontalFlip::new(0.0);
|
||
let pixels: Vec<f32> = (0..6 * 6 * 3).map(|i| i as f32 / 100.0).collect();
|
||
for seed in 0..50u64 {
|
||
let out = flip.apply(&pixels, 6, 6, 3, seed);
|
||
assert_eq!(out, pixels, "prob=0 must never flip");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_flip_prob_one_always_flips() {
|
||
let flip = RandomHorizontalFlip::new(1.0);
|
||
// Use an asymmetric image so flipping is detectable.
|
||
let mut pixels = vec![0.0f32; 4 * 4 * 3];
|
||
pixels[0] = 1.0; // top-left pixel, channel 0
|
||
for seed in 0..10u64 {
|
||
let out = flip.apply(&pixels, 4, 4, 3, seed);
|
||
// After flip, top-left is now what was top-right (which is 0.0)
|
||
assert!((out[0] - 0.0).abs() < 1e-6, "prob=1 must always flip");
|
||
// The original top-left value (1.0) is now at the rightmost position of row 0
|
||
assert!((out[(4 - 1) * 3] - 1.0).abs() < 1e-6);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_flip_output_same_size() {
|
||
let flip = RandomHorizontalFlip::new(0.5);
|
||
let pixels: Vec<f32> = vec![0.5; 8 * 8 * 3];
|
||
let out = flip.apply(&pixels, 8, 8, 3, 12345);
|
||
assert_eq!(out.len(), pixels.len());
|
||
}
|
||
|
||
// ── JepaAugmentationPipeline ──────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_normalize_changes_values() {
|
||
let mut pixels = vec![0.5f32; 6];
|
||
let mean = [0.485, 0.456, 0.406];
|
||
let std = [0.229, 0.224, 0.225];
|
||
JepaAugmentationPipeline::normalize(&mut pixels, mean, std);
|
||
// 0.5 normalized should differ from 0.5
|
||
assert!(
|
||
(pixels[0] - 0.5).abs() > 1e-4,
|
||
"normalization must change values"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_augmentation_pipeline_output_shape() {
|
||
let config = JepaDataConfig {
|
||
image_size: 32,
|
||
patch_size: 8,
|
||
imagenet_normalize: true,
|
||
use_horizontal_flip: true,
|
||
scale_range: (0.5, 1.0),
|
||
ratio_range: (0.75, 1.33),
|
||
..Default::default()
|
||
};
|
||
let pipeline = JepaAugmentationPipeline::from_config(&config);
|
||
let img = make_record(64, 64, None);
|
||
let out = pipeline.process(&img, 0);
|
||
assert_eq!(out.len(), 32 * 32 * 3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_augmentation_consistent_shape_across_samples() {
|
||
let config = JepaDataConfig {
|
||
image_size: 32,
|
||
patch_size: 8,
|
||
imagenet_normalize: true,
|
||
use_horizontal_flip: true,
|
||
scale_range: (0.5, 1.0),
|
||
ratio_range: (0.75, 1.33),
|
||
..Default::default()
|
||
};
|
||
let pipeline = JepaAugmentationPipeline::from_config(&config);
|
||
for seed in 0..20u64 {
|
||
let img = make_record(64, 64, None);
|
||
let out = pipeline.process(&img, seed);
|
||
assert_eq!(out.len(), 32 * 32 * 3, "seed {seed}: wrong output size");
|
||
}
|
||
}
|
||
|
||
// ── InMemoryShard ─────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_in_memory_shard_record_count() {
|
||
let shard = InMemoryShard::synthetic(50, 32, 0);
|
||
assert_eq!(shard.len(), 50);
|
||
assert!(!shard.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_in_memory_shard_zero() {
|
||
let shard = InMemoryShard::synthetic(0, 32, 0);
|
||
assert!(shard.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_in_memory_shard_pixel_range() {
|
||
let shard = InMemoryShard::synthetic(5, 16, 1);
|
||
for rec in &shard.records {
|
||
for &px in &rec.pixels {
|
||
assert!(px >= 0.0 && px <= 1.0, "pixel {px} out of [0,1]");
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_in_memory_shard_sequential_labels() {
|
||
let shard = InMemoryShard::synthetic(10, 16, 0);
|
||
for (i, rec) in shard.records.iter().enumerate() {
|
||
assert_eq!(rec.label, Some(i % 1000));
|
||
}
|
||
}
|
||
|
||
// ── JepaDataPipeline ─────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_pipeline_num_batches_per_epoch() {
|
||
// 2 shards × 10 images = 20 total; batch_size=8 → ceil(20/8)=3
|
||
let pipeline = make_pipeline(2, 10, 8);
|
||
assert_eq!(pipeline.num_batches_per_epoch(), 3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_next_batch_exhausts_epoch() {
|
||
let mut pipeline = make_pipeline(1, 5, 2);
|
||
// 5 images, batch_size=2 → 3 batches (2+2+1)
|
||
let mut count = 0;
|
||
while pipeline.next_batch().is_some() {
|
||
count += 1;
|
||
}
|
||
assert_eq!(count, 3);
|
||
// After exhaustion, next_batch returns None again.
|
||
assert!(pipeline.next_batch().is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_batch_sizes_correct() {
|
||
let mut pipeline = make_pipeline(1, 7, 3);
|
||
// batch 0: 3, batch 1: 3, batch 2: 1
|
||
let b0 = pipeline.next_batch().expect("batch 0");
|
||
assert_eq!(b0.batch_size(), 3);
|
||
let b1 = pipeline.next_batch().expect("batch 1");
|
||
assert_eq!(b1.batch_size(), 3);
|
||
let b2 = pipeline.next_batch().expect("batch 2 (partial)");
|
||
assert_eq!(b2.batch_size(), 1);
|
||
assert!(pipeline.next_batch().is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_reset_epoch_allows_reiteration() {
|
||
let mut pipeline = make_pipeline(1, 4, 4);
|
||
assert!(pipeline.next_batch().is_some());
|
||
assert!(pipeline.next_batch().is_none());
|
||
pipeline.reset_epoch();
|
||
assert_eq!(pipeline.epoch(), 1);
|
||
assert!(pipeline.next_batch().is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_epoch_counter() {
|
||
let mut pipeline = make_pipeline(1, 2, 2);
|
||
assert_eq!(pipeline.epoch(), 0);
|
||
pipeline.reset_epoch();
|
||
assert_eq!(pipeline.epoch(), 1);
|
||
pipeline.reset_epoch();
|
||
assert_eq!(pipeline.epoch(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_multi_shard() {
|
||
// 3 shards × 4 images = 12 total; batch_size=5 → 3 batches
|
||
let mut pipeline = make_pipeline(3, 4, 5);
|
||
let mut total = 0;
|
||
while let Some(b) = pipeline.next_batch() {
|
||
total += b.batch_size();
|
||
}
|
||
assert_eq!(total, 12);
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_batch_size_gt_shard_size() {
|
||
// batch_size larger than any single shard — should still work
|
||
let mut pipeline = make_pipeline(2, 3, 100);
|
||
// 6 total images → 1 batch of 6
|
||
let batch = pipeline.next_batch().expect("should yield one batch");
|
||
assert_eq!(batch.batch_size(), 6);
|
||
assert!(pipeline.next_batch().is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_pipeline_single_shard() {
|
||
let mut pipeline = make_pipeline(1, 8, 4);
|
||
let b0 = pipeline.next_batch().unwrap();
|
||
let b1 = pipeline.next_batch().unwrap();
|
||
assert_eq!(b0.batch_size() + b1.batch_size(), 8);
|
||
assert!(pipeline.next_batch().is_none());
|
||
}
|
||
|
||
// ── JepaBatch mask indices ────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_batch_context_target_non_overlapping() {
|
||
let mut pipeline = make_pipeline(1, 8, 4);
|
||
let batch = pipeline.next_batch().unwrap();
|
||
for (ctx, tgt) in batch
|
||
.context_patch_indices
|
||
.iter()
|
||
.zip(batch.all_target_indices.iter())
|
||
{
|
||
for &ci in ctx {
|
||
assert!(!tgt.contains(&ci), "context index {ci} in target set");
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_batch_target_indices_in_range() {
|
||
let config = JepaDataConfig {
|
||
image_size: 32,
|
||
patch_size: 8,
|
||
batch_size: 4,
|
||
..JepaDataConfig::default()
|
||
};
|
||
let total_patches = config.num_patches(); // (32/8)^2 = 16
|
||
let shards = vec![InMemoryShard::synthetic(8, 32, 0)];
|
||
let mut pipeline = JepaDataPipeline::new(config, shards);
|
||
let batch = pipeline.next_batch().unwrap();
|
||
for tgt in &batch.all_target_indices {
|
||
for &idx in tgt {
|
||
assert!(idx < total_patches, "target index {idx} >= {total_patches}");
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_batch_context_sorted_deduped() {
|
||
let mut pipeline = make_pipeline(1, 8, 4);
|
||
let batch = pipeline.next_batch().unwrap();
|
||
for ctx in &batch.context_patch_indices {
|
||
for w in ctx.windows(2) {
|
||
assert!(
|
||
w[0] < w[1],
|
||
"context indices must be sorted and deduplicated"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_batch_labels_match_shard() {
|
||
let mut pipeline = make_pipeline(1, 6, 6);
|
||
let batch = pipeline.next_batch().unwrap();
|
||
// All labels must be Some (synthetic shards always set labels)
|
||
for label in &batch.labels {
|
||
assert!(label.is_some());
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_batch_id_increments() {
|
||
let mut pipeline = make_pipeline(1, 9, 3);
|
||
let b0 = pipeline.next_batch().unwrap();
|
||
let b1 = pipeline.next_batch().unwrap();
|
||
let b2 = pipeline.next_batch().unwrap();
|
||
assert_eq!(b0.batch_id, 0);
|
||
assert_eq!(b1.batch_id, 1);
|
||
assert_eq!(b2.batch_id, 2);
|
||
}
|
||
|
||
// ── DatasetStats ──────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_dataset_stats_totals() {
|
||
let pipeline = make_pipeline(3, 10, 8);
|
||
let stats = DatasetStats::compute(&pipeline);
|
||
assert_eq!(stats.total_images, 30);
|
||
assert_eq!(stats.num_shards, 3);
|
||
assert_eq!(stats.images_per_shard, vec![10, 10, 10]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_dataset_stats_mask_efficiency_positive() {
|
||
let pipeline = make_pipeline(1, 8, 4);
|
||
let stats = DatasetStats::compute(&pipeline);
|
||
assert!(stats.mask_efficiency > 0.0, "mask_efficiency should be > 0");
|
||
assert!(
|
||
stats.mask_efficiency <= 1.0,
|
||
"mask_efficiency should be <= 1"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_dataset_stats_avg_patches_positive() {
|
||
let pipeline = make_pipeline(1, 8, 4);
|
||
let stats = DatasetStats::compute(&pipeline);
|
||
assert!(stats.avg_context_patches > 0.0);
|
||
assert!(stats.avg_target_patches > 0.0);
|
||
}
|
||
|
||
// ── WebDatasetShard ───────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_web_dataset_shard_to_in_memory() {
|
||
let shard = WebDatasetShard::new("/data/imagenet/shard-0000.tar", 100, 0);
|
||
let mem = shard.to_in_memory(32);
|
||
assert_eq!(mem.len(), 100);
|
||
assert_eq!(mem.shard_id, 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_web_dataset_shard_compressed_flag() {
|
||
let gz = WebDatasetShard::new("data/shard.tar.gz", 10, 0);
|
||
assert!(gz.compressed);
|
||
let plain = WebDatasetShard::new("data/shard.tar", 10, 0);
|
||
assert!(!plain.compressed);
|
||
}
|
||
|
||
#[test]
|
||
fn test_web_dataset_shard_pixel_count() {
|
||
let shard = WebDatasetShard::new("test.tar", 5, 1);
|
||
let mem = shard.to_in_memory(16);
|
||
for rec in &mem.records {
|
||
assert_eq!(rec.pixels.len(), 16 * 16 * 3);
|
||
}
|
||
}
|
||
|
||
// ── Tar parsing helpers ───────────────────────────────────────────────────
|
||
|
||
/// Build a tar header block for a regular file with the given filename and size.
|
||
fn make_tar_header(filename: &str, size: u64) -> [u8; 512] {
|
||
let mut block = [0u8; 512];
|
||
// filename: bytes 0..100
|
||
let fname_bytes = filename.as_bytes();
|
||
block[..fname_bytes.len().min(99)]
|
||
.copy_from_slice(&fname_bytes[..fname_bytes.len().min(99)]);
|
||
// size: bytes 124..136, octal ASCII with trailing null
|
||
let octal = format!("{:011o}\0", size);
|
||
block[124..136].copy_from_slice(octal.as_bytes());
|
||
// typeflag: byte 156 = '0' for regular file
|
||
block[156] = b'0';
|
||
// checksum: bytes 148..156, sum of all header bytes with checksum field as spaces
|
||
for i in 148..156 {
|
||
block[i] = b' ';
|
||
}
|
||
let checksum: u32 = block.iter().map(|&b| b as u32).sum();
|
||
let chk_str = format!("{:06o}\0 ", checksum);
|
||
block[148..156].copy_from_slice(chk_str.as_bytes());
|
||
block
|
||
}
|
||
|
||
/// Build a tar header block for a directory entry.
|
||
fn make_tar_dir_header(dirname: &str) -> [u8; 512] {
|
||
let mut block = [0u8; 512];
|
||
let fname_bytes = dirname.as_bytes();
|
||
block[..fname_bytes.len().min(99)]
|
||
.copy_from_slice(&fname_bytes[..fname_bytes.len().min(99)]);
|
||
// size = 0
|
||
let octal = format!("{:011o}\0", 0u64);
|
||
block[124..136].copy_from_slice(octal.as_bytes());
|
||
// typeflag '5' = directory
|
||
block[156] = b'5';
|
||
for i in 148..156 {
|
||
block[i] = b' ';
|
||
}
|
||
let checksum: u32 = block.iter().map(|&b| b as u32).sum();
|
||
let chk_str = format!("{:06o}\0 ", checksum);
|
||
block[148..156].copy_from_slice(chk_str.as_bytes());
|
||
block
|
||
}
|
||
|
||
/// Build a minimal valid tar archive from (key, image_bytes, label) tuples.
|
||
///
|
||
/// Each entry emits a `{key}.jpg` file and, if label is Some, a `{key}.cls` file.
|
||
fn make_test_tar(records: &[(&str, &[u8], Option<usize>)]) -> Vec<u8> {
|
||
let mut out: Vec<u8> = Vec::new();
|
||
|
||
for (key, image_bytes, label) in records {
|
||
// Image file
|
||
let img_name = format!("{}.jpg", key);
|
||
let img_size = image_bytes.len() as u64;
|
||
out.extend_from_slice(&make_tar_header(&img_name, img_size));
|
||
out.extend_from_slice(image_bytes);
|
||
// Pad to next 512-byte boundary
|
||
let pad = (512 - (image_bytes.len() % 512)) % 512;
|
||
out.extend(std::iter::repeat(0u8).take(pad));
|
||
|
||
// Class label file (optional)
|
||
if let Some(lbl) = label {
|
||
let cls_content = format!("{}\n", lbl);
|
||
let cls_bytes = cls_content.as_bytes();
|
||
let cls_name = format!("{}.cls", key);
|
||
out.extend_from_slice(&make_tar_header(&cls_name, cls_bytes.len() as u64));
|
||
out.extend_from_slice(cls_bytes);
|
||
let cls_pad = (512 - (cls_bytes.len() % 512)) % 512;
|
||
out.extend(std::iter::repeat(0u8).take(cls_pad));
|
||
}
|
||
}
|
||
|
||
// End-of-archive: two zero blocks
|
||
out.extend(std::iter::repeat(0u8).take(1024));
|
||
out
|
||
}
|
||
|
||
// ── Tar parsing tests ─────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_webdataset_shard_load_reads_real_tar() {
|
||
// Write a real tar to a temp file and load it through WebDatasetShard.
|
||
let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
|
||
let tar = make_test_tar(&[("000000", fake_jpg, Some(7)), ("000001", fake_jpg, None)]);
|
||
let dir = std::env::temp_dir();
|
||
let path = dir.join(format!("jepa_shard_load_test_{}.tar", std::process::id()));
|
||
std::fs::write(&path, &tar).expect("write temp tar");
|
||
|
||
let shard = WebDatasetShard::new(path.to_str().unwrap(), 2, 3);
|
||
let mem = shard.load().expect("load real tar");
|
||
std::fs::remove_file(&path).ok();
|
||
|
||
assert_eq!(mem.shard_id, 3);
|
||
assert_eq!(mem.records.len(), 2);
|
||
assert_eq!(mem.records[0].label, Some(7));
|
||
assert_eq!(mem.records[1].label, None);
|
||
}
|
||
|
||
#[test]
|
||
fn test_webdataset_shard_load_reads_gzip_tar() {
|
||
use std::io::Write;
|
||
|
||
let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
|
||
let tar = make_test_tar(&[("000000", fake_jpg, Some(11))]);
|
||
|
||
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
|
||
encoder.write_all(&tar).expect("gzip write");
|
||
let gz = encoder.finish().expect("gzip finish");
|
||
|
||
let dir = std::env::temp_dir();
|
||
let path = dir.join(format!("jepa_shard_gz_test_{}.tar.gz", std::process::id()));
|
||
std::fs::write(&path, &gz).expect("write temp tar.gz");
|
||
|
||
let shard = WebDatasetShard::new(path.to_str().unwrap(), 1, 5);
|
||
let mem = shard.load().expect("load gzip tar");
|
||
std::fs::remove_file(&path).ok();
|
||
|
||
assert!(shard.compressed);
|
||
assert_eq!(mem.shard_id, 5);
|
||
assert_eq!(mem.records.len(), 1);
|
||
assert_eq!(mem.records[0].label, Some(11));
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_empty_tar() {
|
||
// Two zero blocks = empty archive
|
||
let data = vec![0u8; 1024];
|
||
let records = parse_tar_bytes(&data);
|
||
assert_eq!(records.len(), 0, "empty tar must yield 0 records");
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_single_record() {
|
||
let fake_jpg = b"\xFF\xD8\xFF\xE0fake jpeg content";
|
||
let tar = make_test_tar(&[("000000", fake_jpg, Some(42))]);
|
||
let records = parse_tar_bytes(&tar);
|
||
assert_eq!(records.len(), 1);
|
||
assert_eq!(records[0].key, "000000");
|
||
assert_eq!(records[0].label, Some(42));
|
||
assert_eq!(records[0].extension, "jpg");
|
||
assert_eq!(records[0].image_bytes, fake_jpg);
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_multiple_records() {
|
||
let fake_jpg = b"\xFF\xD8content";
|
||
let entries: Vec<(&str, &[u8], Option<usize>)> = vec![
|
||
("000000", fake_jpg, Some(0)),
|
||
("000001", fake_jpg, Some(1)),
|
||
("000002", fake_jpg, Some(2)),
|
||
("000003", fake_jpg, Some(3)),
|
||
("000004", fake_jpg, Some(4)),
|
||
];
|
||
let tar = make_test_tar(&entries);
|
||
let records = parse_tar_bytes(&tar);
|
||
assert_eq!(records.len(), 5, "expected 5 records");
|
||
for (i, rec) in records.iter().enumerate() {
|
||
assert_eq!(rec.label, Some(i));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_record_without_cls() {
|
||
let fake_jpg = b"\xFF\xD8no label";
|
||
let tar = make_test_tar(&[("000000", fake_jpg, None)]);
|
||
let records = parse_tar_bytes(&tar);
|
||
assert_eq!(records.len(), 1);
|
||
assert_eq!(
|
||
records[0].label, None,
|
||
"image with no .cls should have label=None"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_directory_entries_skipped() {
|
||
let mut out: Vec<u8> = Vec::new();
|
||
// Directory entry
|
||
out.extend_from_slice(&make_tar_dir_header("somedir/"));
|
||
// A regular image record after the directory
|
||
let fake_jpg = b"\xFF\xD8data";
|
||
let img_name = "000000.jpg";
|
||
out.extend_from_slice(&make_tar_header(img_name, fake_jpg.len() as u64));
|
||
out.extend_from_slice(fake_jpg);
|
||
let pad = (512 - (fake_jpg.len() % 512)) % 512;
|
||
out.extend(std::iter::repeat(0u8).take(pad));
|
||
// End-of-archive
|
||
out.extend(std::iter::repeat(0u8).take(1024));
|
||
|
||
let records = parse_tar_bytes(&out);
|
||
assert_eq!(records.len(), 1, "directory entries must be skipped");
|
||
assert_eq!(records[0].key, "000000");
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_large_file() {
|
||
// 50 000 bytes > 512*97 = 49664, spans many blocks
|
||
let big_image: Vec<u8> = (0..50_000u32).map(|i| (i % 251) as u8).collect();
|
||
let tar = make_test_tar(&[("000000", &big_image, Some(7))]);
|
||
let records = parse_tar_bytes(&tar);
|
||
assert_eq!(records.len(), 1);
|
||
assert_eq!(records[0].image_bytes.len(), 50_000);
|
||
assert_eq!(records[0].image_bytes, big_image);
|
||
}
|
||
|
||
// ── ShuffleBuffer tests ───────────────────────────────────────────────────
|
||
|
||
fn make_image_record(key: &str) -> ImageRecord {
|
||
ImageRecord {
|
||
pixels: vec![0.5f32; 8 * 8 * 3],
|
||
width: 8,
|
||
height: 8,
|
||
channels: 3,
|
||
label: None,
|
||
key: key.to_owned(),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_shuffle_buffer_fill() {
|
||
let mut buf = ShuffleBuffer::new(10);
|
||
for i in 0..10 {
|
||
buf.push(make_image_record(&format!("img{i}")));
|
||
}
|
||
assert_eq!(buf.len(), 10);
|
||
assert!(buf.is_full());
|
||
}
|
||
|
||
#[test]
|
||
fn test_shuffle_buffer_push_overflow() {
|
||
let mut buf = ShuffleBuffer::new(10);
|
||
for i in 0..15 {
|
||
buf.push(make_image_record(&format!("img{i}")));
|
||
}
|
||
assert_eq!(buf.len(), 10, "buffer must not grow beyond capacity");
|
||
}
|
||
|
||
#[test]
|
||
fn test_shuffle_buffer_drain() {
|
||
let mut buf = ShuffleBuffer::new(10);
|
||
for i in 0..10 {
|
||
buf.push(make_image_record(&format!("img{i}")));
|
||
}
|
||
let batch = buf.drain_batch(5);
|
||
assert_eq!(batch.len(), 5);
|
||
assert_eq!(buf.len(), 5);
|
||
}
|
||
|
||
#[test]
|
||
fn test_shuffle_buffer_empty_drain() {
|
||
let mut buf = ShuffleBuffer::new(10);
|
||
for i in 0..3 {
|
||
buf.push(make_image_record(&format!("img{i}")));
|
||
}
|
||
// Drain more than available
|
||
let batch = buf.drain_batch(100);
|
||
assert_eq!(
|
||
batch.len(),
|
||
3,
|
||
"drain should return at most what's in the buffer"
|
||
);
|
||
assert_eq!(buf.len(), 0);
|
||
|
||
// Drain zero
|
||
let empty = buf.drain_batch(0);
|
||
assert_eq!(empty.len(), 0);
|
||
}
|
||
|
||
// ── ShardLoadStats tests ──────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_shard_load_stats_bytes() {
|
||
use std::io::Write;
|
||
|
||
// Write a real tar to a temp file and read it back.
|
||
let fake_jpg = b"\xFF\xD8fake";
|
||
let entries: Vec<(&str, &[u8], Option<usize>)> = vec![
|
||
("000000", fake_jpg, Some(0)),
|
||
("000001", fake_jpg, Some(1)),
|
||
("000002", fake_jpg, Some(2)),
|
||
];
|
||
let tar_bytes = make_test_tar(&entries);
|
||
|
||
let mut tmp = std::env::temp_dir();
|
||
tmp.push("jepa_test_shard_stats.tar");
|
||
{
|
||
let mut f = std::fs::File::create(&tmp).expect("create temp file");
|
||
f.write_all(&tar_bytes).expect("write tar");
|
||
}
|
||
|
||
let (records, stats) = read_webdataset_shard(&tmp).expect("read shard");
|
||
std::fs::remove_file(&tmp).ok();
|
||
|
||
assert_eq!(records.len(), 3);
|
||
assert!(stats.bytes_read > 0, "bytes_read must be > 0");
|
||
assert_eq!(stats.records_loaded, 3);
|
||
}
|
||
|
||
// ── from_filesystem tests ─────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_from_filesystem_missing_path() {
|
||
let config = JepaDataConfig::default();
|
||
let missing = std::path::PathBuf::from("/nonexistent/path/shard-0000.tar");
|
||
let result = JepaDataPipeline::from_filesystem(config, vec![missing.clone()]);
|
||
assert!(result.is_err(), "missing path must return Err");
|
||
// Extract the error message without calling unwrap_err (JepaDataPipeline has no Debug).
|
||
let msg = match result {
|
||
Err(e) => e,
|
||
Ok(_) => panic!("expected Err"),
|
||
};
|
||
assert!(
|
||
msg.contains("nonexistent") || msg.contains("shard-0000.tar"),
|
||
"error message should mention the path: {msg}"
|
||
);
|
||
}
|
||
|
||
// ── webdataset_record_to_image ────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_webdataset_decode_placeholder_shape() {
|
||
let rec = WebDatasetRecord {
|
||
key: "test".to_string(),
|
||
image_bytes: vec![0u8; 100], // not a valid image
|
||
label: Some(1),
|
||
extension: "jpg".to_string(),
|
||
};
|
||
let img = webdataset_record_to_image(rec);
|
||
assert_eq!(img.pixels.len(), 224 * 224 * 3);
|
||
assert_eq!(img.width, 224);
|
||
assert_eq!(img.height, 224);
|
||
assert_eq!(img.channels, 3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_webdataset_decode_placeholder_value() {
|
||
let rec = WebDatasetRecord {
|
||
key: "k".to_string(),
|
||
image_bytes: vec![],
|
||
label: None,
|
||
extension: "jpg".to_string(),
|
||
};
|
||
let img = webdataset_record_to_image(rec);
|
||
for &p in &img.pixels {
|
||
assert!(
|
||
(p - 0.5f32).abs() < 1e-6 || (0.0..=1.0).contains(&p),
|
||
"pixel {} must be in [0, 1]",
|
||
p
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_webdataset_decode_label_preserved() {
|
||
let rec = WebDatasetRecord {
|
||
key: "sample_0000".to_string(),
|
||
image_bytes: vec![0xFF, 0xD8, 0xFF], // JPEG magic start
|
||
label: Some(42),
|
||
extension: "jpg".to_string(),
|
||
};
|
||
let img = webdataset_record_to_image(rec);
|
||
assert_eq!(img.label, Some(42));
|
||
}
|
||
|
||
#[test]
|
||
fn test_webdataset_decode_key_preserved() {
|
||
let rec = WebDatasetRecord {
|
||
key: "my_key_123".to_string(),
|
||
image_bytes: vec![],
|
||
label: None,
|
||
extension: "png".to_string(),
|
||
};
|
||
let img = webdataset_record_to_image(rec);
|
||
assert_eq!(img.key, "my_key_123");
|
||
}
|
||
|
||
// ── HTTP WebDataset ──────────────────────────────────────────────────────
|
||
|
||
// test: validate_urls accepts valid http URLs
|
||
#[test]
|
||
fn test_validate_urls_http_ok() {
|
||
let urls = vec![
|
||
"http://example.com/shard-0.tar".to_string(),
|
||
"https://storage.example.com/data/shard-1.tar".to_string(),
|
||
];
|
||
assert!(JepaDataPipeline::validate_urls(&urls).is_ok());
|
||
}
|
||
|
||
// test: validate_urls rejects non-http schemes
|
||
#[test]
|
||
fn test_validate_urls_bad_scheme() {
|
||
let urls = vec!["s3://bucket/shard.tar".to_string()];
|
||
assert!(JepaDataPipeline::validate_urls(&urls).is_err());
|
||
}
|
||
|
||
// test: validate_urls empty list is Ok
|
||
#[test]
|
||
fn test_validate_urls_empty() {
|
||
assert!(JepaDataPipeline::validate_urls(&[]).is_ok());
|
||
}
|
||
|
||
// test: from_urls with bad scheme returns Err immediately
|
||
#[test]
|
||
fn test_from_urls_bad_scheme_err() {
|
||
let config = JepaDataConfig::default();
|
||
let result =
|
||
JepaDataPipeline::from_urls(config, vec!["ftp://example.com/shard.tar".to_string()]);
|
||
assert!(result.is_err());
|
||
let msg = match result {
|
||
Err(e) => e,
|
||
Ok(_) => panic!("expected Err"),
|
||
};
|
||
assert!(msg.contains("unsupported URL scheme"), "error was: {msg}");
|
||
}
|
||
|
||
// test: download_webdataset_shard with invalid URL scheme returns Err
|
||
#[test]
|
||
fn test_download_bad_scheme() {
|
||
let result = download_webdataset_shard("file:///tmp/test.tar", 0);
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
// test: download_webdataset_shard to a non-existent host returns Err
|
||
// 192.0.2.0/24 is reserved "documentation" space — guaranteed unreachable.
|
||
// Uses a 5-second timeout so the test completes quickly.
|
||
#[test]
|
||
fn test_download_unreachable_host() {
|
||
let result =
|
||
download_webdataset_shard_with_timeout("http://192.0.2.1:9999/shard.tar", 0, 5);
|
||
assert!(result.is_err(), "unreachable host should return Err");
|
||
}
|
||
|
||
// test: UrlShardDescriptor creation
|
||
#[test]
|
||
fn test_url_shard_descriptor_new() {
|
||
let desc = UrlShardDescriptor::new("https://example.com/shard-0.tar", 1000, 0);
|
||
assert_eq!(desc.expected_records, 1000);
|
||
assert_eq!(desc.shard_id, 0);
|
||
assert!(desc.is_https());
|
||
}
|
||
|
||
// test: UrlShardDescriptor::is_https with http URL
|
||
#[test]
|
||
fn test_url_shard_descriptor_is_http() {
|
||
let desc = UrlShardDescriptor::new("http://example.com/shard.tar", 0, 1);
|
||
assert!(!desc.is_https());
|
||
}
|
||
|
||
// test: UrlShardDescriptor clone
|
||
#[test]
|
||
fn test_url_shard_descriptor_clone() {
|
||
let desc = UrlShardDescriptor::new("https://a.b/c.tar", 42, 3);
|
||
let desc2 = desc.clone();
|
||
assert_eq!(desc2.url, desc.url);
|
||
assert_eq!(desc2.shard_id, 3);
|
||
}
|
||
|
||
// test: from_urls with empty list succeeds with empty pipeline
|
||
#[test]
|
||
fn test_from_urls_empty_list() {
|
||
let config = JepaDataConfig::default();
|
||
// Empty URL list → empty pipeline (no shards loaded)
|
||
let result = JepaDataPipeline::from_urls(config, vec![]);
|
||
assert!(result.is_ok());
|
||
}
|
||
}
|
||
|
||
#[cfg(all(test, feature = "image-decode"))]
|
||
mod tests_image_decode {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_webdataset_decode_real_image() {
|
||
// Create a minimal 1×1 PNG using the image crate itself
|
||
use image::{ImageBuffer, Rgb};
|
||
let img_buf: ImageBuffer<Rgb<u8>, Vec<u8>> =
|
||
ImageBuffer::from_pixel(1, 1, Rgb([128u8, 64u8, 32u8]));
|
||
let mut bytes = Vec::new();
|
||
img_buf
|
||
.write_to(
|
||
&mut std::io::Cursor::new(&mut bytes),
|
||
image::ImageOutputFormat::Png,
|
||
)
|
||
.unwrap();
|
||
|
||
let rec = WebDatasetRecord {
|
||
key: "synthetic".to_string(),
|
||
image_bytes: bytes,
|
||
label: Some(0),
|
||
extension: "png".to_string(),
|
||
};
|
||
let decoded = webdataset_record_to_image(rec);
|
||
assert_eq!(decoded.pixels.len(), 224 * 224 * 3);
|
||
assert_eq!(decoded.width, 224);
|
||
assert_eq!(decoded.height, 224);
|
||
for &p in &decoded.pixels {
|
||
assert!((0.0..=1.0).contains(&p), "pixel {} out of range", p);
|
||
}
|
||
}
|
||
}
|