stable_diffusion_trainer/data_set/image_data_set/
mod.rs

1//! ImageDataSet is a simple struct that holds the path to a directory containing images. It is used to represent a dataset of images that can be used for training a model.
2
3use crate::prelude::*;
4use std::path::{Path, PathBuf};
5
6/// A data set of images.
7#[derive(Debug, Serialize, Deserialize)]
8pub struct ImageDataSet(PathBuf);
9
10impl ImageDataSet {
11    /// Create a new image data set from a directory.
12    pub fn from_dir<Path: AsRef<std::path::Path>>(path: Path) -> Self {
13        let path = path.as_ref().to_path_buf();
14        ImageDataSet(path)
15    }
16
17    /// Set the path to the image data set.
18    pub fn set_path(&mut self, path: PathBuf) {
19        self.0 = path;
20    }
21
22    /// Get the path to the image data set.
23    pub fn path(&self) -> &Path {
24        &self.0
25    }
26}