Skip to main content

shrew_data/
lib.rs

1//! # shrew-data
2//!
3//! Data loading, datasets, and batching for Shrew.
4//!
5//! This crate provides:
6//! - [`Dataset`] trait — unified interface for any dataset
7//! - [`DataLoader`] — batching, shuffling, parallel iteration over a Dataset
8//! - [`AsyncDataLoader`] — prefetching data loader with background workers
9//   - Dataset combinators — SubsetDataset, ConcatDataset, MapDataset, VecDataset
10//   - Image augmentation transforms — RandomFlip, RandomCrop, etc.
11//   - CSV dataset loader
12//   - Built-in datasets: MNIST (IDX format parser)
13//   - ImageFolder (directory-based image classification dataset)
14//   - Train/test splitting with reproducible seeding
15
16pub mod async_loader;
17pub mod augment;
18pub mod combinators;
19pub mod csv_dataset;
20pub mod dataset;
21pub mod image_folder;
22pub mod loader;
23pub mod mnist;
24pub mod transform;
25
26pub use async_loader::{AsyncDataLoader, AsyncDataLoaderConfig, Batch, PrefetchIterator};
27pub use augment::{
28    ColorJitter, RandomCrop, RandomErasing, RandomHorizontalFlip, RandomNoise, RandomVerticalFlip,
29};
30pub use combinators::{train_test_split, ConcatDataset, MapDataset, SubsetDataset, VecDataset};
31pub use csv_dataset::{CsvConfig, CsvDataset};
32pub use dataset::{Dataset, Sample};
33pub use loader::{DataLoader, DataLoaderConfig};
34pub use mnist::MnistDataset;
35pub use transform::Transform;
36
37#[cfg(feature = "image-folder")]
38pub use image_folder::{ImageFolder, ImageFolderBuilder};
39
40// Error types for ImageFolder
41
42/// Errors from the ImageFolder dataset.
43#[derive(Debug)]
44pub enum ImageFolderError {
45    /// The root path is not a directory.
46    NotADirectory(String),
47    /// No class subdirectories found.
48    NoClasses(String),
49    /// No image files found.
50    NoImages(String),
51    /// Image decoding failed.
52    ImageDecode(String, String),
53    /// I/O error.
54    Io(std::io::Error),
55}
56
57impl std::fmt::Display for ImageFolderError {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        match self {
60            ImageFolderError::NotADirectory(p) => write!(f, "Not a directory: {p}"),
61            ImageFolderError::NoClasses(p) => write!(f, "No class subdirectories in {p}"),
62            ImageFolderError::NoImages(p) => write!(f, "No image files found in {p}"),
63            ImageFolderError::ImageDecode(p, e) => write!(f, "Failed to decode {p}: {e}"),
64            ImageFolderError::Io(e) => write!(f, "I/O error: {e}"),
65        }
66    }
67}
68
69impl std::error::Error for ImageFolderError {}