tenshift-core 0.1.2

//! Provides an `ImageFolderSource` that scans a directory layout for classification tasks.
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use crate::error::{Error, Result};
use crate::sample::{Sample, Tensor};
use crate::source::{Source, SourceIterator};

/// A source that loads files from a directory structure where each subdirectory
/// represents a class label. Just like `PyTorch`'s `ImageFolder`.
///
/// Example:
/// ```text
/// root/
/// ├── dog/
/// │   ├── 1.jpg
/// │   └── 2.jpg
/// └── cat/
///     ├── 1.jpg
///     └── 2.jpg
/// ```
///
/// Yields samples with three fields:
/// - `"data"`  -  Raw bytes of the file.
/// - `"label"`  -  `i64` index of the class (sorted alphabetically).
/// - `"path"`  -  Raw bytes of the file path.
pub struct ImageFolderSource {
    root: PathBuf,
    samples: Vec<(PathBuf, i64)>,
    classes: Vec<String>,
}

impl ImageFolderSource {
    /// Create a new `ImageFolder` source.
    ///
    /// # Errors
    /// Returns an error if the root directory cannot be read or contains no subdirectories.
    pub fn new(root: impl Into<PathBuf>) -> Result<Self> {
        let root = root.into();
        let mut class_dirs = BTreeMap::new();

        let entries = std::fs::read_dir(&root).map_err(|e| Error::ReadFailed {
            path: root.clone(),
            reason: e.to_string(),
        })?;

        for entry in entries {
            let entry = entry.map_err(|e| Error::ReadFailed {
                path: root.clone(),
                reason: e.to_string(),
            })?;

            let path = entry.path();
            if path.is_dir() {
                if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
                    class_dirs.insert(name.to_string(), path.clone());
                }
            }
        }

        if class_dirs.is_empty() {
            return Err(Error::EmptySource {
                pattern: root.to_string_lossy().to_string(),
            });
        }

        let mut classes = Vec::with_capacity(class_dirs.len());
        let mut samples = Vec::new();

        for (idx, (class_name, dir_path)) in class_dirs.into_iter().enumerate() {
            classes.push(class_name);

            // Note: Does not currently recurse into sub-subdirectories.
            // Fail loud on a read_dir or per-entry error rather than silently
            // skipping the class directory: the former `if let Ok(files)` +
            // `files.flatten()` dropped an unreadable/permission-denied class
            // (and individual unreadable entries) with no signal, training on a
            // silently incomplete, class-imbalanced dataset (Law 10).
            let read_failed = |e: std::io::Error| Error::ReadFailed {
                path: dir_path.clone(),
                reason: e.to_string(),
            };
            let files = std::fs::read_dir(&dir_path).map_err(&read_failed)?;
            for file in files {
                let file = file.map_err(&read_failed)?;
                let path = file.path();
                if path.is_file() {
                    let label = i64::try_from(idx).map_err(|_| Error::InvalidConfig {
                        reason: format!("class index {} exceeds i64::MAX", idx),
                    })?;
                    samples.push((path, label));
                }
            }
        }

        if samples.is_empty() {
            return Err(Error::EmptySource {
                pattern: root.to_string_lossy().to_string(),
            });
        }

        Ok(Self {
            root,
            samples,
            classes,
        })
    }

    /// Get the class names in index order.
    pub fn classes(&self) -> &[String] {
        &self.classes
    }
}

impl Source for ImageFolderSource {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        Ok(Box::new(ImageFolderIterator {
            samples: self.samples.clone(),
            index: 0,
        }))
    }

    fn len_hint(&self) -> Option<u64> {
        Some(self.samples.len() as u64)
    }

    fn name(&self) -> &str {
        self.root.to_str().map_or("ImageFolder", |s| s)
    }
}

struct ImageFolderIterator {
    samples: Vec<(PathBuf, i64)>,
    index: u64,
}

impl SourceIterator for ImageFolderIterator {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        let Ok(idx) = usize::try_from(self.index) else {
            return Some(Err(Error::InvalidConfig {
                reason: "source index exceeded addressable memory on this platform".to_string(),
            }));
        };

        if idx >= self.samples.len() {
            return None;
        }

        let (path, label) = &self.samples[idx];
        let index = self.index;
        self.index += 1;

        Some(load_image_folder_sample(path, *label, index))
    }
}

fn load_image_folder_sample(path: &Path, label: i64, index: u64) -> Result<Sample> {
    // Check file size before reading to prevent OOM on huge files
    let metadata = std::fs::metadata(path).map_err(|e| Error::ReadFailed {
        path: path.to_path_buf(),
        reason: e.to_string(),
    })?;
    if metadata.len() > crate::pipeline::MAX_LOAD_FILE_SIZE {
        return Err(Error::ReadFailed {
            path: path.to_path_buf(),
            reason: format!(
                "file size {} exceeds maximum {} bytes. Fix: check for oversized files or increase MAX_LOAD_FILE_SIZE",
                metadata.len(),
                crate::pipeline::MAX_LOAD_FILE_SIZE
            ),
        });
    }

    let data = std::fs::read(path).map_err(|e| Error::ReadFailed {
        path: path.to_path_buf(),
        reason: e.to_string(),
    })?;

    let filename = path.to_string_lossy().to_string();

    Ok(Sample::new()
        .with("data", Tensor::bytes(data))
        .with("label", Tensor::i64(&[label], vec![1]))
        .with("path", Tensor::bytes(filename.as_bytes().to_vec()))
        .with_metadata(path.to_string_lossy(), index))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::tempdir;

    #[test]
    fn image_folder_source() {
        let dir = tempdir().unwrap();

        let class_a = dir.path().join("cats");
        let class_b = dir.path().join("dogs");

        fs::create_dir(&class_a).unwrap();
        fs::create_dir(&class_b).unwrap();

        fs::write(class_a.join("1.jpg"), "cat1").unwrap();
        fs::write(class_a.join("2.jpg"), "cat2").unwrap();
        fs::write(class_b.join("1.jpg"), "dog1").unwrap();

        let source = ImageFolderSource::new(dir.path()).unwrap();
        assert_eq!(source.len_hint().unwrap(), 3);
        assert_eq!(source.classes(), &["cats".to_string(), "dogs".to_string()]);

        let mut iter = source.open().unwrap();

        // Iteration order depends on OS directory reading, but we can verify it yields correct format
        let mut count = 0;
        while let Some(res) = iter.next_sample() {
            let sample = res.unwrap();
            assert!(sample.contains("data"));
            assert!(sample.contains("label"));
            assert!(sample.contains("path"));
            count += 1;
        }
        assert_eq!(count, 3);
    }

    #[cfg(unix)]
    #[test]
    fn unreadable_class_dir_fails_loud_not_silently_skipped() {
        // Regression: the old `if let Ok(files) = read_dir(..)` + `.flatten()`
        // silently skipped an unreadable class directory (and unreadable
        // entries), training on a silently incomplete, class-imbalanced set.
        // A permission-denied class dir must now surface as Error::ReadFailed.
        use std::os::unix::fs::PermissionsExt;
        let dir = tempdir().unwrap();
        let class = dir.path().join("dog");
        fs::create_dir(&class).unwrap();
        fs::write(class.join("1.bin"), "x").unwrap();
        // A second readable class so class discovery does not early-out empty.
        let other = dir.path().join("cat");
        fs::create_dir(&other).unwrap();
        fs::write(other.join("1.bin"), "y").unwrap();

        fs::set_permissions(&class, fs::Permissions::from_mode(0o000)).unwrap();
        // Root bypasses permission bits; only assert when denial is real.
        let reproducible = fs::read_dir(&class).is_err();
        let result = ImageFolderSource::new(dir.path());
        // Restore perms so tempdir cleanup can remove the tree.
        fs::set_permissions(&class, fs::Permissions::from_mode(0o755)).unwrap();

        if reproducible {
            match result {
                Err(Error::ReadFailed { .. }) => {}
                Err(other) => panic!("expected ReadFailed, got a different error: {other:?}"),
                Ok(_) => panic!("unreadable class dir was silently skipped instead of failing loud"),
            }
        }
    }
}