glint_mask_tools/loaders/
single_file.rs

1/// Single-file image loader for RGB and similar sensors.
2///
3/// This loader handles sensors where each capture is stored in a single file,
4/// typically RGB images from standard cameras.
5use image::DynamicImage;
6use ndarray::Array3;
7use std::path::Path;
8
9use crate::core::{
10    image_loader::{list_image_files, ImageCapture},
11    ImageLoader,
12};
13use crate::error::{GlintError, Result};
14
15/// Image loader for single-file captures (RGB images, etc.)
16///
17/// This loader handles standard image files where each capture consists
18/// of a single file containing all bands. Common for RGB cameras and
19/// some single-band or composite multispectral formats.
20#[derive(Debug, Clone)]
21pub struct SingleFileLoader {
22    /// Supported file extensions
23    extensions: Vec<String>,
24    /// Number of expected bands/channels
25    band_count: usize,
26    /// Bit depth of the sensor
27    bit_depth: u8,
28}
29
30impl SingleFileLoader {
31    /// Create a new single-file loader
32    ///
33    /// # Arguments
34    ///
35    /// * `extensions` - Supported file extensions (without dots, e.g., ["jpg", "png"])
36    /// * `band_count` - Expected number of bands/channels in each image
37    /// * `bit_depth` - Bit depth of the sensor (8, 16, or 32)
38    pub fn new(extensions: Vec<String>, band_count: usize, bit_depth: u8) -> Result<Self> {
39        if extensions.is_empty() {
40            return Err(GlintError::validation(
41                "At least one file extension must be supported",
42            ));
43        }
44
45        if band_count == 0 {
46            return Err(GlintError::validation("Band count must be greater than 0"));
47        }
48
49        if !matches!(bit_depth, 8 | 16 | 32) {
50            return Err(GlintError::InvalidBitDepth { bit_depth });
51        }
52
53        Ok(Self {
54            extensions,
55            band_count,
56            bit_depth,
57        })
58    }
59
60    /// Create a loader for standard RGB images
61    pub fn rgb() -> Result<Self> {
62        Self::new(
63            vec![
64                "jpg".to_string(),
65                "jpeg".to_string(),
66                "png".to_string(),
67                "tif".to_string(),
68                "tiff".to_string(),
69            ],
70            3,
71            8,
72        )
73    }
74
75    /// Create a loader for grayscale images
76    pub fn grayscale() -> Result<Self> {
77        Self::new(
78            vec![
79                "jpg".to_string(),
80                "jpeg".to_string(),
81                "png".to_string(),
82                "tif".to_string(),
83                "tiff".to_string(),
84            ],
85            1,
86            8,
87        )
88    }
89
90    /// Create a loader for 16-bit TIFF images
91    pub fn tiff_16bit(band_count: usize) -> Result<Self> {
92        Self::new(vec!["tif".to_string(), "tiff".to_string()], band_count, 16)
93    }
94
95    /// Get the base filename without extension
96    fn get_base_name(&self, path: &Path) -> String {
97        path.file_stem()
98            .and_then(|s| s.to_str())
99            .unwrap_or("unknown")
100            .to_string()
101    }
102
103    /// Load an image file and convert to the expected format
104    fn load_image_file(&self, path: &Path) -> Result<Array3<f64>> {
105        // Load the image
106        let img = image::open(path)?;
107
108        // Convert to the expected format based on band count
109        let processed_img = match self.band_count {
110            1 => img.to_luma8().into(),
111            3 => img.to_rgb8().into(),
112            4 => img.to_rgba8().into(),
113            _ => {
114                // For other band counts, try to use the image as-is and validate later
115                img
116            }
117        };
118
119        // Convert to ndarray
120        self.dynamic_image_to_array(processed_img)
121    }
122
123    /// Convert a DynamicImage to an ndarray
124    fn dynamic_image_to_array(&self, img: DynamicImage) -> Result<Array3<f64>> {
125        match img {
126            DynamicImage::ImageLuma8(img) => {
127                let (width, height) = img.dimensions();
128                let data: Vec<f64> = img.into_raw().into_iter().map(|x| x as f64).collect();
129                let array = Array3::from_shape_vec((height as usize, width as usize, 1), data)
130                    .map_err(|_| GlintError::processing("Failed to reshape image data"))?;
131                Ok(array)
132            }
133            DynamicImage::ImageLuma16(img) => {
134                let (width, height) = img.dimensions();
135                let data: Vec<f64> = img.into_raw().into_iter().map(|x| x as f64).collect();
136                let array = Array3::from_shape_vec((height as usize, width as usize, 1), data)
137                    .map_err(|_| GlintError::processing("Failed to reshape image data"))?;
138                Ok(array)
139            }
140            DynamicImage::ImageRgb8(img) => {
141                let (width, height) = img.dimensions();
142                let raw_data = img.into_raw();
143                let mut data = Vec::with_capacity(raw_data.len());
144
145                // Convert RGB interleaved to f64
146                for &pixel in &raw_data {
147                    data.push(pixel as f64);
148                }
149
150                // Reshape to (height, width, 3)
151                let array = Array3::from_shape_vec((height as usize, width as usize, 3), data)
152                    .map_err(|_| GlintError::processing("Failed to reshape RGB image data"))?;
153                Ok(array)
154            }
155            DynamicImage::ImageRgb16(img) => {
156                let (width, height) = img.dimensions();
157                let raw_data = img.into_raw();
158                let mut data = Vec::with_capacity(raw_data.len());
159
160                for &pixel in &raw_data {
161                    data.push(pixel as f64);
162                }
163
164                let array = Array3::from_shape_vec((height as usize, width as usize, 3), data)
165                    .map_err(|_| GlintError::processing("Failed to reshape RGB16 image data"))?;
166                Ok(array)
167            }
168            DynamicImage::ImageRgba8(img) => {
169                let (width, height) = img.dimensions();
170                let raw_data = img.into_raw();
171                let mut data = Vec::with_capacity(raw_data.len());
172
173                for &pixel in &raw_data {
174                    data.push(pixel as f64);
175                }
176
177                let array = Array3::from_shape_vec((height as usize, width as usize, 4), data)
178                    .map_err(|_| GlintError::processing("Failed to reshape RGBA image data"))?;
179                Ok(array)
180            }
181            DynamicImage::ImageRgba16(img) => {
182                let (width, height) = img.dimensions();
183                let raw_data = img.into_raw();
184                let mut data = Vec::with_capacity(raw_data.len());
185
186                for &pixel in &raw_data {
187                    data.push(pixel as f64);
188                }
189
190                let array = Array3::from_shape_vec((height as usize, width as usize, 4), data)
191                    .map_err(|_| GlintError::processing("Failed to reshape RGBA16 image data"))?;
192                Ok(array)
193            }
194            _ => Err(GlintError::processing("Unsupported image format")),
195        }
196    }
197}
198
199impl ImageLoader for SingleFileLoader {
200    fn as_any(&self) -> &dyn std::any::Any {
201        self
202    }
203
204    fn discover_captures(&self, input_dir: &Path, output_dir: &Path) -> Result<Vec<ImageCapture>> {
205        let image_files = list_image_files(input_dir, &self.extensions)?;
206
207        let mut captures = Vec::new();
208
209        for file_path in image_files {
210            let base_name = self.get_base_name(&file_path);
211            let mask_paths = self.generate_mask_paths(
212                &ImageCapture {
213                    id: base_name.clone(),
214                    paths: vec![file_path.clone()],
215                    mask_paths: Vec::new(), // Temporary, will be overwritten
216                },
217                output_dir,
218            );
219
220            captures.push(ImageCapture {
221                id: base_name,
222                paths: vec![file_path],
223                mask_paths,
224            });
225        }
226
227        // Sort captures by ID for consistent ordering
228        captures.sort_by(|a, b| a.id.cmp(&b.id));
229
230        Ok(captures)
231    }
232
233    fn load_image(&self, capture: &ImageCapture) -> Result<Array3<f64>> {
234        if capture.paths.len() != 1 {
235            return Err(GlintError::validation(format!(
236                "Single-file loader expects exactly 1 file, got {}",
237                capture.paths.len()
238            )));
239        }
240
241        let path = &capture.paths[0];
242        if !path.exists() {
243            return Err(GlintError::MissingFiles {
244                files: capture.paths.clone(),
245            });
246        }
247
248        let array = self.load_image_file(path)?;
249
250        // Validate band count
251        let (_, _, actual_bands) = array.dim();
252        if actual_bands != self.band_count {
253            return Err(GlintError::BandCountMismatch {
254                expected: self.band_count,
255                actual: actual_bands,
256            });
257        }
258
259        Ok(array)
260    }
261
262    fn band_count(&self) -> usize {
263        self.band_count
264    }
265
266    fn bit_depth(&self) -> u8 {
267        self.bit_depth
268    }
269
270    fn supported_extensions(&self) -> Vec<String> {
271        self.extensions.clone()
272    }
273
274    fn expected_file_count(&self) -> usize {
275        1
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use std::fs;
283    use std::path::PathBuf;
284    use tempfile::tempdir;
285
286    #[test]
287    fn test_single_file_loader_creation() {
288        let loader = SingleFileLoader::new(vec!["jpg".to_string()], 3, 8).unwrap();
289        assert_eq!(loader.band_count(), 3);
290        assert_eq!(loader.bit_depth(), 8);
291        assert_eq!(loader.supported_extensions(), vec!["jpg"]);
292
293        // Invalid parameters should fail
294        assert!(SingleFileLoader::new(vec![], 3, 8).is_err());
295        assert!(SingleFileLoader::new(vec!["jpg".to_string()], 0, 8).is_err());
296        assert!(SingleFileLoader::new(vec!["jpg".to_string()], 3, 7).is_err());
297    }
298
299    #[test]
300    fn test_rgb_loader() {
301        let loader = SingleFileLoader::rgb().unwrap();
302        assert_eq!(loader.band_count(), 3);
303        assert_eq!(loader.bit_depth(), 8);
304        assert!(loader.supported_extensions().contains(&"jpg".to_string()));
305        assert!(loader.supported_extensions().contains(&"png".to_string()));
306    }
307
308    #[test]
309    fn test_grayscale_loader() {
310        let loader = SingleFileLoader::grayscale().unwrap();
311        assert_eq!(loader.band_count(), 1);
312        assert_eq!(loader.bit_depth(), 8);
313    }
314
315    #[test]
316    fn test_discover_captures() {
317        let temp_dir = tempdir().unwrap();
318        let input_dir = temp_dir.path().join("input");
319        let output_dir = temp_dir.path().join("output");
320        fs::create_dir_all(&input_dir).unwrap();
321        fs::create_dir_all(&output_dir).unwrap();
322
323        // Create test files
324        let file1 = input_dir.join("image001.jpg");
325        let file2 = input_dir.join("image002.jpg");
326        let file3 = input_dir.join("document.txt"); // Should be ignored
327
328        fs::write(&file1, b"fake image data").unwrap();
329        fs::write(&file2, b"fake image data").unwrap();
330        fs::write(&file3, b"text file").unwrap();
331
332        let loader = SingleFileLoader::rgb().unwrap();
333        let captures = loader.discover_captures(&input_dir, &output_dir).unwrap();
334
335        assert_eq!(captures.len(), 2);
336        assert_eq!(captures[0].id, "image001");
337        assert_eq!(captures[1].id, "image002");
338        assert_eq!(captures[0].paths.len(), 1);
339        assert_eq!(captures[0].paths[0], file1);
340        assert_eq!(captures[0].mask_paths.len(), 1);
341        assert!(captures[0].mask_paths[0]
342            .to_string_lossy()
343            .contains("image001_mask.png"));
344    }
345
346    #[test]
347    fn test_base_name_extraction() {
348        let loader = SingleFileLoader::rgb().unwrap();
349
350        let path = Path::new("/path/to/image.jpg");
351        assert_eq!(loader.get_base_name(path), "image");
352
353        let path = Path::new("image.jpeg");
354        assert_eq!(loader.get_base_name(path), "image");
355
356        let path = Path::new("/path/to/image.with.dots.png");
357        assert_eq!(loader.get_base_name(path), "image.with.dots");
358    }
359
360    #[test]
361    fn test_validation() {
362        let loader = SingleFileLoader::rgb().unwrap();
363
364        // Valid capture should pass
365        let capture = ImageCapture {
366            id: "test".to_string(),
367            paths: vec![PathBuf::from("test.jpg")],
368            mask_paths: vec![PathBuf::from("test_mask.png")],
369        };
370
371        // Multiple files should fail for single-file loader
372        let _invalid_capture = ImageCapture {
373            id: "test".to_string(),
374            paths: vec![PathBuf::from("test1.jpg"), PathBuf::from("test2.jpg")],
375            mask_paths: vec![PathBuf::from("test_mask.png")],
376        };
377
378        assert!(loader.validate_capture(&capture).is_err()); // File doesn't exist
379                                                             // Note: We can't easily test successful validation without creating actual files
380    }
381}
glint_mask_tools/loaders/single_file.rs

glint_mask_tools/loaders/
single_file.rs
