glint_mask_tools/core/
image_loader.rs

1/// Image loading abstraction for different sensor formats.
2///
3/// This module defines the [`ImageLoader`] trait which provides a uniform
4/// interface for loading images from various sensor formats, handling both
5/// single-file and multi-file captures.
6use crate::error::{GlintError, Result};
7use ndarray::{Array2, Array3};
8use std::path::{Path, PathBuf};
9
10/// Represents a single image capture, which may consist of multiple files
11/// (e.g., one file per band for multispectral sensors)
12#[derive(Debug, Clone)]
13pub struct ImageCapture {
14    /// Primary identifier for this capture (usually the first file path)
15    pub id: String,
16    /// List of file paths that make up this capture
17    pub paths: Vec<PathBuf>,
18    /// Expected output mask paths (one per input file for compatibility with Agisoft Metashape)
19    pub mask_paths: Vec<PathBuf>,
20}
21
22/// Trait for loading images from different sensor formats.
23///
24/// This trait abstracts the process of discovering, loading, and organizing
25/// image files from various sensor types. Each sensor type has its own
26/// file organization pattern (single files, multi-band files, etc.).
27pub trait ImageLoader: Send + Sync {
28    /// Support for downcasting to concrete types
29    fn as_any(&self) -> &dyn std::any::Any;
30    /// Discover all image captures in the given directory
31    ///
32    /// This method scans the input directory and identifies all valid
33    /// image captures based on the sensor's file pattern. Each capture
34    /// may consist of one or more files.
35    fn discover_captures(&self, input_dir: &Path, output_dir: &Path) -> Result<Vec<ImageCapture>>;
36
37    /// Load a single image capture into a 3D array (height, width, channels)
38    ///
39    /// The returned array should have:
40    /// - Axis 0: Height (Y)
41    /// - Axis 1: Width (X)  
42    /// - Axis 2: Channels/Bands
43    ///
44    /// Values should be in the sensor's native range (e.g., 0-255 for 8-bit, 0-65535 for 16-bit)
45    fn load_image(&self, capture: &ImageCapture) -> Result<Array3<f64>>;
46
47    /// Get the expected number of bands/channels for this loader
48    fn band_count(&self) -> usize;
49
50    /// Get the bit depth for this sensor type
51    fn bit_depth(&self) -> u8;
52
53    /// Get supported file extensions for this loader
54    fn supported_extensions(&self) -> Vec<String>;
55
56    /// Validate that a capture has the expected file structure
57    fn validate_capture(&self, capture: &ImageCapture) -> Result<()> {
58        // Check that all files exist
59        for path in &capture.paths {
60            if !path.exists() {
61                return Err(GlintError::MissingFiles {
62                    files: vec![path.clone()],
63                });
64            }
65        }
66
67        // Check that we have the expected number of files
68        if capture.paths.len() != self.expected_file_count() {
69            return Err(GlintError::validation(format!(
70                "Expected {} files, got {}",
71                self.expected_file_count(),
72                capture.paths.len()
73            )));
74        }
75
76        Ok(())
77    }
78
79    /// Get the expected number of files per capture for this sensor
80    fn expected_file_count(&self) -> usize {
81        1 // Default to single file, override for multi-file sensors
82    }
83
84    /// Save a mask image to the specified path
85    fn save_mask(&self, mask: &Array2<u8>, path: &Path) -> Result<()> {
86        let (height, width) = mask.dim();
87        let mut img_buffer = vec![0u8; height * width];
88
89        for (i, row) in mask.rows().into_iter().enumerate() {
90            for (j, &pixel) in row.iter().enumerate() {
91                img_buffer[i * width + j] = pixel;
92            }
93        }
94
95        let img = image::GrayImage::from_raw(width as u32, height as u32, img_buffer)
96            .ok_or_else(|| GlintError::processing("Failed to create image from mask data"))?;
97
98        // Ensure output directory exists
99        if let Some(parent) = path.parent() {
100            std::fs::create_dir_all(parent)?;
101        }
102
103        img.save(path)?;
104        Ok(())
105    }
106
107    /// Save masks to multiple paths (one per input file for compatibility with Agisoft Metashape)
108    fn save_masks(&self, mask: &Array2<u8>, capture: &ImageCapture) -> Result<()> {
109        for mask_path in &capture.mask_paths {
110            self.save_mask(mask, mask_path)?;
111        }
112        Ok(())
113    }
114
115    /// Generate the output mask path for a given capture
116    fn generate_mask_path(&self, capture: &ImageCapture, output_dir: &Path) -> PathBuf {
117        let filename = format!("{}_mask.png", capture.id);
118        output_dir.join(filename)
119    }
120
121    /// Generate output mask paths for a capture (one per input file for Agisoft Metashape compatibility)
122    fn generate_mask_paths(&self, capture: &ImageCapture, output_dir: &Path) -> Vec<PathBuf> {
123        capture
124            .paths
125            .iter()
126            .map(|input_path| {
127                let filename = if let Some(stem) = input_path.file_stem() {
128                    format!("{}_mask.png", stem.to_string_lossy())
129                } else {
130                    format!("{}_mask.png", capture.id)
131                };
132                output_dir.join(filename)
133            })
134            .collect()
135    }
136}
137
138/// Utility function to normalize image data from arbitrary bit depth to [0, 1] range
139pub fn normalize_image(img: &Array3<f64>, bit_depth: u8) -> Result<Array3<f64>> {
140    let max_value = match bit_depth {
141        8 => 255.0,
142        16 => 65535.0,
143        32 => 4294967295.0,
144        _ => return Err(GlintError::InvalidBitDepth { bit_depth }),
145    };
146
147    Ok(img / max_value)
148}
149
150/// Utility function to list image files in a directory with given extensions
151pub fn list_image_files(dir: &Path, extensions: &[String]) -> Result<Vec<PathBuf>> {
152    let mut files = Vec::new();
153
154    for entry in std::fs::read_dir(dir)? {
155        let entry = entry?;
156        let path = entry.path();
157
158        if path.is_file() {
159            if let Some(ext) = path.extension() {
160                let ext_str = ext.to_string_lossy().to_lowercase();
161                if extensions.iter().any(|e| e.to_lowercase() == ext_str) {
162                    files.push(path);
163                }
164            }
165        }
166    }
167
168    files.sort();
169    Ok(files)
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175    use tempfile::tempdir;
176
177    #[test]
178    fn test_normalize_image() {
179        let img = Array3::from_shape_vec((2, 2, 1), vec![0.0, 127.0, 255.0, 128.0]).unwrap();
180        let normalized = normalize_image(&img, 8).unwrap();
181
182        assert_eq!(normalized[[0, 0, 0]], 0.0);
183        assert!((normalized[[0, 1, 0]] - 0.498).abs() < 0.001);
184        assert_eq!(normalized[[1, 0, 0]], 1.0);
185    }
186
187    #[test]
188    fn test_list_image_files() {
189        let dir = tempdir().unwrap();
190        let img_path = dir.path().join("test.jpg");
191        std::fs::write(&img_path, b"fake image data").unwrap();
192
193        let files = list_image_files(dir.path(), &["jpg".to_string()]).unwrap();
194        assert_eq!(files.len(), 1);
195        assert_eq!(files[0], img_path);
196    }
197}
glint_mask_tools/core/image_loader.rs

glint_mask_tools/core/
image_loader.rs
