environment/training-source/source-snapshot/examples/common/mod.rs

//! Dataset handling shared by the decoder trainer and evaluator.

#![allow(dead_code)] // Each example uses a different subset of these helpers.

use std::collections::HashMap;
use std::io::Read;
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetManifest {
    pub schema_version: u32,
    pub name: String,
    /// Relative to the manifest file unless absolute. May be overridden by
    /// `DINOVISION_DATA_ROOT` on another machine.
    pub root: PathBuf,
    pub images: Vec<ManifestImage>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ManifestImage {
    pub path: PathBuf,
    pub split: String,
    pub source: String,
    /// Capture session, photo collection, or other leakage boundary. A group
    /// must occur in exactly one split.
    pub group: String,
    pub sha256: String,
    pub bytes: u64,
}

pub struct LoadedManifest {
    pub manifest: DatasetManifest,
    pub root: PathBuf,
}

pub fn load_manifest(path: &Path) -> Result<LoadedManifest, Box<dyn std::error::Error>> {
    let bytes = std::fs::read(path)?;
    let manifest: DatasetManifest = serde_json::from_slice(&bytes)?;
    if manifest.schema_version != 1 {
        return Err(format!(
            "{} uses unsupported manifest schema {}",
            path.display(),
            manifest.schema_version
        )
        .into());
    }
    let root = match std::env::var_os("DINOVISION_DATA_ROOT") {
        Some(value) => PathBuf::from(value),
        None if manifest.root.is_absolute() => manifest.root.clone(),
        None => path
            .parent()
            .unwrap_or_else(|| Path::new("."))
            .join(&manifest.root),
    };

    validate_manifest(&manifest)?;
    Ok(LoadedManifest { manifest, root })
}

fn validate_manifest(manifest: &DatasetManifest) -> Result<(), Box<dyn std::error::Error>> {
    if manifest.images.is_empty() {
        return Err("dataset manifest contains no images".into());
    }

    let mut group_splits: HashMap<(&str, &str), &str> = HashMap::new();
    let mut hashes: HashMap<&str, (&str, &Path)> = HashMap::new();
    let mut paths: HashMap<&Path, &str> = HashMap::new();
    for image in &manifest.images {
        if image.split.is_empty()
            || image.source.is_empty()
            || image.group.is_empty()
            || image.sha256.len() != 64
        {
            return Err(format!("invalid manifest entry for {}", image.path.display()).into());
        }
        let key = (image.source.as_str(), image.group.as_str());
        if let Some(old) = group_splits.insert(key, image.split.as_str())
            && old != image.split
        {
            return Err(format!(
                "group {}/{} leaks across splits {old:?} and {:?}",
                image.source, image.group, image.split
            )
            .into());
        }
        if let Some(old_split) = paths.insert(&image.path, image.split.as_str()) {
            return Err(format!(
                "path {} is listed more than once ({old_split} and {})",
                image.path.display(),
                image.split
            )
            .into());
        }
        if let Some((old_split, old_path)) =
            hashes.insert(image.sha256.as_str(), (image.split.as_str(), &image.path))
        {
            return Err(format!(
                "identical files occur more than once: {} ({old_split}) and {} ({})",
                old_path.display(),
                image.path.display(),
                image.split
            )
            .into());
        }
    }
    Ok(())
}

pub fn images_for_split(
    path: &Path,
    split: &str,
    limit: usize,
) -> Result<Vec<(PathBuf, ManifestImage)>, Box<dyn std::error::Error>> {
    let loaded = load_manifest(path)?;
    let mut result = Vec::new();
    for image in loaded
        .manifest
        .images
        .into_iter()
        .filter(|image| image.split == split)
        .take(limit)
    {
        result.push((loaded.root.join(&image.path), image));
    }
    if result.is_empty() {
        return Err(format!("manifest has no images in split {split:?}").into());
    }
    Ok(result)
}

/// Collect image paths, round-robin across subdirectories.
pub fn find_images(root: &Path, limit: usize) -> Vec<PathBuf> {
    let mut groups: Vec<Vec<PathBuf>> = Vec::new();
    let mut stack = vec![root.to_path_buf()];
    while let Some(dir) = stack.pop() {
        let Ok(entries) = std::fs::read_dir(&dir) else {
            continue;
        };
        let mut here = Vec::new();
        let mut items: Vec<_> = entries.filter_map(|e| e.ok()).map(|e| e.path()).collect();
        items.sort();
        for path in items {
            if path.is_dir() {
                stack.push(path);
            } else if matches!(
                path.extension().and_then(|e| e.to_str()),
                Some("JPEG" | "jpeg" | "jpg" | "png" | "rgb")
            ) {
                here.push(path);
            }
        }
        if !here.is_empty() {
            groups.push(here);
        }
    }
    groups.sort();

    let mut out = Vec::new();
    let deepest = groups.iter().map(|g| g.len()).max().unwrap_or(0);
    'outer: for i in 0..deepest {
        for group in &groups {
            if let Some(path) = group.get(i) {
                out.push(path.clone());
                if out.len() >= limit {
                    break 'outer;
                }
            }
        }
    }
    out
}

/// Load a photograph or a raw headset capture as interleaved RGB8.
pub fn load_frame(path: &Path, size: u32) -> Option<Vec<u8>> {
    if path.extension().and_then(|e| e.to_str()) == Some("rgb") {
        let bytes = std::fs::read(path).ok()?;
        let expected = (size * size * 3) as usize;
        if bytes.len() != expected {
            log::warn!(
                "{}: {} bytes, expected {expected} — captured at a different resolution?",
                path.display(),
                bytes.len()
            );
            return None;
        }
        return Some(bytes);
    }
    load_rgb(path, size)
}

fn load_rgb(path: &Path, size: u32) -> Option<Vec<u8>> {
    let img = image::open(path).ok()?.to_rgb8();
    let (w, h) = (img.width(), img.height());
    let side = w.min(h);
    let cropped = image::imageops::crop_imm(&img, (w - side) / 2, (h - side) / 2, side, side);
    let resized = image::imageops::resize(
        &cropped.to_image(),
        size,
        size,
        image::imageops::FilterType::CatmullRom,
    );
    Some(resized.into_raw())
}

pub fn sha256(path: &Path) -> Result<String, std::io::Error> {
    let mut file = std::fs::File::open(path)?;
    let mut hasher = Sha256::new();
    let mut buffer = vec![0u8; 1024 * 1024];
    loop {
        let n = file.read(&mut buffer)?;
        if n == 0 {
            break;
        }
        hasher.update(&buffer[..n]);
    }
    Ok(format!("{:x}", hasher.finalize()))
}

pub fn verify_image(path: &Path, image: &ManifestImage) -> Result<(), Box<dyn std::error::Error>> {
    let metadata = std::fs::metadata(path)?;
    if metadata.len() != image.bytes {
        return Err(format!(
            "{} is {} bytes; manifest records {}",
            path.display(),
            metadata.len(),
            image.bytes
        )
        .into());
    }
    let digest = sha256(path)?;
    if digest != image.sha256 {
        return Err(format!(
            "{} has SHA-256 {digest}; manifest records {}",
            path.display(),
            image.sha256
        )
        .into());
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn image(path: &str, split: &str, group: &str, hash: char) -> ManifestImage {
        ManifestImage {
            path: PathBuf::from(path),
            split: split.to_string(),
            source: "camera".to_string(),
            group: group.to_string(),
            sha256: std::iter::repeat_n(hash, 64).collect(),
            bytes: 1,
        }
    }

    fn manifest(images: Vec<ManifestImage>) -> DatasetManifest {
        DatasetManifest {
            schema_version: 1,
            name: "test".to_string(),
            root: PathBuf::from("."),
            images,
        }
    }

    #[test]
    fn independent_groups_and_splits_are_valid() {
        validate_manifest(&manifest(vec![
            image("a.rgb", "train", "session-a", 'a'),
            image("b.rgb", "test", "session-b", 'b'),
        ]))
        .unwrap();
    }

    #[test]
    fn group_cannot_cross_a_split() {
        let error = validate_manifest(&manifest(vec![
            image("a.rgb", "train", "session-a", 'a'),
            image("b.rgb", "test", "session-a", 'b'),
        ]))
        .unwrap_err()
        .to_string();
        assert!(error.contains("leaks across splits"), "{error}");
    }

    #[test]
    fn duplicate_content_is_rejected_even_within_one_split() {
        let error = validate_manifest(&manifest(vec![
            image("a.rgb", "train", "session-a", 'a'),
            image("b.rgb", "train", "session-a", 'a'),
        ]))
        .unwrap_err()
        .to_string();
        assert!(error.contains("identical files"), "{error}");
    }

    #[test]
    fn duplicate_path_is_rejected() {
        let error = validate_manifest(&manifest(vec![
            image("a.rgb", "train", "session-a", 'a'),
            image("a.rgb", "train", "session-a", 'b'),
        ]))
        .unwrap_err()
        .to_string();
        assert!(error.contains("listed more than once"), "{error}");
    }
}