glint-mask-tools 0.1.0

/// Memory-efficient loader for large single-file images (CIR, etc.).
///
/// This loader processes large images in chunks to avoid memory limits,
/// similar to the Python BigTiffLoader implementation.
use image::{DynamicImage, GenericImageView};
use ndarray::{Array2, Array3};
use std::path::Path;

use crate::core::{
    image_loader::{normalize_image, ImageCapture},
    GlintAlgorithm, ImageLoader, PostProcessor,
};
use crate::error::{GlintError, Result};

/// Image loader for large single-file captures that require chunked processing
///
/// This loader handles very large image files (like CIR imagery) by processing
/// them in chunks to avoid memory limitations. It implements a tiled processing
/// approach where the image is loaded and processed in small sections.
#[derive(Debug, Clone)]
pub struct BigTiffLoader {
    /// Supported file extensions
    extensions: Vec<String>,
    /// Number of expected bands/channels
    band_count: usize,
    /// Bit depth of the sensor
    bit_depth: u8,
    /// Chunk size for tiled processing (pixels)
    chunk_size: usize,
}

impl BigTiffLoader {
    /// Create a new big TIFF loader
    ///
    /// # Arguments
    ///
    /// * `extensions` - Supported file extensions (without dots, e.g., ["tif", "tiff"])
    /// * `band_count` - Expected number of bands/channels in each image
    /// * `bit_depth` - Bit depth of the sensor (8, 16, or 32)
    /// * `chunk_size` - Size of processing chunks in pixels (default: 256)
    pub fn new(
        extensions: Vec<String>,
        band_count: usize,
        bit_depth: u8,
        chunk_size: Option<usize>,
    ) -> Result<Self> {
        if extensions.is_empty() {
            return Err(GlintError::validation(
                "At least one file extension must be supported",
            ));
        }

        if band_count == 0 {
            return Err(GlintError::validation("Band count must be greater than 0"));
        }

        if !matches!(bit_depth, 8 | 16 | 32) {
            return Err(GlintError::InvalidBitDepth { bit_depth });
        }

        Ok(Self {
            extensions,
            band_count,
            bit_depth,
            chunk_size: chunk_size.unwrap_or(256),
        })
    }

    /// Create a loader for large CIR images
    pub fn cir() -> Result<Self> {
        Self::new(vec!["tif".to_string(), "tiff".to_string()], 4, 8, Some(256))
    }

    /// Get the base filename without extension
    fn get_base_name(&self, path: &Path) -> String {
        path.file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("unknown")
            .to_string()
    }

    /// Process a large image in chunks using the provided algorithm and post-processor
    pub fn process_chunked_image(
        &self,
        capture: &ImageCapture,
        algorithm: &dyn GlintAlgorithm,
        postprocessor: &dyn PostProcessor,
        bit_depth: u8,
        pixel_buffer: usize,
    ) -> Result<()> {
        if capture.paths.len() != 1 {
            return Err(GlintError::validation(format!(
                "Big TIFF loader expects exactly 1 file, got {}",
                capture.paths.len()
            )));
        }

        let path = &capture.paths[0];
        if !path.exists() {
            return Err(GlintError::MissingFiles {
                files: capture.paths.clone(),
            });
        }

        // We need to read the image metadata without loading the full image
        // For now, we'll implement a fallback strategy using smaller chunks
        self.process_with_fallback(capture, algorithm, postprocessor, bit_depth, pixel_buffer)
    }

    /// Fallback processing method that attempts smaller chunks if memory limits are exceeded
    fn process_with_fallback(
        &self,
        capture: &ImageCapture,
        algorithm: &dyn GlintAlgorithm,
        postprocessor: &dyn PostProcessor,
        bit_depth: u8,
        _pixel_buffer: usize,
    ) -> Result<()> {
        let path = &capture.paths[0];

        // Try to load the image with larger memory limits
        // We'll set a very high pixel limit to allow large images
        let img = self.load_large_image(path)?;

        // Get image dimensions
        let (width, height) = img.dimensions();
        let (img_height, img_width, bands) = (height as usize, width as usize, self.band_count);

        // Validate band count
        if bands != self.band_count {
            return Err(GlintError::BandCountMismatch {
                expected: self.band_count,
                actual: bands,
            });
        }

        // Create output mask buffer
        let mut output_mask = Array2::<u8>::zeros((img_height, img_width));

        // Process the image in chunks
        let chunk_size = self.chunk_size;

        for y in (0..img_height).step_by(chunk_size) {
            for x in (0..img_width).step_by(chunk_size) {
                let chunk_width = std::cmp::min(chunk_size, img_width - x);
                let chunk_height = std::cmp::min(chunk_size, img_height - y);

                // Extract chunk from the image
                let chunk = self.extract_chunk(&img, x, y, chunk_width, chunk_height)?;

                // Normalize the chunk
                let normalized_chunk = normalize_image(&chunk, bit_depth)?;

                // Apply algorithm
                let chunk_mask = algorithm.detect_glint(&normalized_chunk)?;

                // Apply post-processing
                let processed_chunk_mask = postprocessor.process_mask(&chunk_mask)?;

                // Copy the processed chunk back to the output mask
                for (i, row) in processed_chunk_mask.rows().into_iter().enumerate() {
                    for (j, &pixel) in row.iter().enumerate() {
                        let global_y = y + i;
                        let global_x = x + j;
                        if global_y < img_height && global_x < img_width {
                            output_mask[[global_y, global_x]] = pixel;
                        }
                    }
                }
            }
        }

        // Save the final mask to all mask paths
        self.save_masks(&output_mask, capture)?;

        Ok(())
    }

    /// Load a large image with increased memory limits
    fn load_large_image(&self, path: &Path) -> Result<DynamicImage> {
        // For very large images, we need to disable the memory limits in the image crate
        // This is similar to setting Image.MAX_IMAGE_PIXELS = None in Python PIL

        // Create a custom image reader with relaxed limits
        use image::{ImageReader, Limits};
        use std::fs::File;

        let file = File::open(path)?;
        let mut reader = ImageReader::new(std::io::BufReader::new(file))
            .with_guessed_format()
            .map_err(|e| GlintError::processing(format!("Failed to create image reader: {}", e)))?;

        // Set very high limits to allow large images
        let mut limits = Limits::default();
        limits.max_image_width = Some(50000);
        limits.max_image_height = Some(50000);
        limits.max_alloc = Some(4_000_000_000); // 4GB
        reader.limits(limits);

        match reader.decode() {
            Ok(img) => Ok(img),
            Err(image::ImageError::Limits(ref _limit_error)) => {
                // If we still hit a memory limit, we need true streaming processing
                Err(GlintError::processing(format!(
                    "Image too large to load even with relaxed limits: {}. True streaming processing needed.",
                    path.display()
                )))
            }
            Err(e) => Err(GlintError::Image(e)),
        }
    }

    /// Extract a chunk from the image as an ndarray
    fn extract_chunk(
        &self,
        img: &DynamicImage,
        x: usize,
        y: usize,
        width: usize,
        height: usize,
    ) -> Result<Array3<f64>> {
        // Create a subimage (crop)
        let cropped = img.crop_imm(x as u32, y as u32, width as u32, height as u32);

        // Convert to the expected format based on band count
        let processed_img = match self.band_count {
            1 => cropped.to_luma8().into(),
            3 => cropped.to_rgb8().into(),
            4 => cropped.to_rgba8().into(),
            _ => cropped,
        };

        // Convert to ndarray
        self.dynamic_image_to_array(processed_img, width, height)
    }

    /// Convert a DynamicImage to an ndarray with specific dimensions
    fn dynamic_image_to_array(
        &self,
        img: DynamicImage,
        width: usize,
        height: usize,
    ) -> Result<Array3<f64>> {
        match img {
            DynamicImage::ImageLuma8(img) => {
                let data: Vec<f64> = img.into_raw().into_iter().map(|x| x as f64).collect();
                let array = Array3::from_shape_vec((height, width, 1), data)
                    .map_err(|_| GlintError::processing("Failed to reshape image data"))?;
                Ok(array)
            }
            DynamicImage::ImageLuma16(img) => {
                let data: Vec<f64> = img.into_raw().into_iter().map(|x| x as f64).collect();
                let array = Array3::from_shape_vec((height, width, 1), data)
                    .map_err(|_| GlintError::processing("Failed to reshape image data"))?;
                Ok(array)
            }
            DynamicImage::ImageRgb8(img) => {
                let raw_data = img.into_raw();
                let mut data = Vec::with_capacity(raw_data.len());

                for &pixel in &raw_data {
                    data.push(pixel as f64);
                }

                let array = Array3::from_shape_vec((height, width, 3), data)
                    .map_err(|_| GlintError::processing("Failed to reshape RGB image data"))?;
                Ok(array)
            }
            DynamicImage::ImageRgb16(img) => {
                let raw_data = img.into_raw();
                let mut data = Vec::with_capacity(raw_data.len());

                for &pixel in &raw_data {
                    data.push(pixel as f64);
                }

                let array = Array3::from_shape_vec((height, width, 3), data)
                    .map_err(|_| GlintError::processing("Failed to reshape RGB16 image data"))?;
                Ok(array)
            }
            DynamicImage::ImageRgba8(img) => {
                let raw_data = img.into_raw();
                let mut data = Vec::with_capacity(raw_data.len());

                for &pixel in &raw_data {
                    data.push(pixel as f64);
                }

                let array = Array3::from_shape_vec((height, width, 4), data)
                    .map_err(|_| GlintError::processing("Failed to reshape RGBA image data"))?;
                Ok(array)
            }
            DynamicImage::ImageRgba16(img) => {
                let raw_data = img.into_raw();
                let mut data = Vec::with_capacity(raw_data.len());

                for &pixel in &raw_data {
                    data.push(pixel as f64);
                }

                let array = Array3::from_shape_vec((height, width, 4), data)
                    .map_err(|_| GlintError::processing("Failed to reshape RGBA16 image data"))?;
                Ok(array)
            }
            _ => Err(GlintError::processing("Unsupported image format")),
        }
    }
}

impl ImageLoader for BigTiffLoader {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn discover_captures(&self, input_dir: &Path, output_dir: &Path) -> Result<Vec<ImageCapture>> {
        let image_files = crate::core::image_loader::list_image_files(input_dir, &self.extensions)?;

        let mut captures = Vec::new();

        for file_path in image_files {
            let base_name = self.get_base_name(&file_path);
            let mask_paths = self.generate_mask_paths(
                &ImageCapture {
                    id: base_name.clone(),
                    paths: vec![file_path.clone()],
                    mask_paths: Vec::new(), // Temporary, will be overwritten
                },
                output_dir,
            );

            captures.push(ImageCapture {
                id: base_name,
                paths: vec![file_path],
                mask_paths,
            });
        }

        // Sort captures by ID for consistent ordering
        captures.sort_by(|a, b| a.id.cmp(&b.id));

        Ok(captures)
    }

    fn load_image(&self, _capture: &ImageCapture) -> Result<Array3<f64>> {
        // For BigTiffLoader, we don't use the standard load_image method
        // Instead, we use process_chunked_image for memory-efficient processing
        Err(GlintError::processing(
            "BigTiffLoader requires chunked processing. Use process_chunked_image instead.",
        ))
    }

    fn band_count(&self) -> usize {
        self.band_count
    }

    fn bit_depth(&self) -> u8 {
        self.bit_depth
    }

    fn supported_extensions(&self) -> Vec<String> {
        self.extensions.clone()
    }

    fn expected_file_count(&self) -> usize {
        1
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_big_tiff_loader_creation() {
        let loader = BigTiffLoader::new(vec!["tif".to_string()], 4, 8, Some(256)).unwrap();
        assert_eq!(loader.band_count(), 4);
        assert_eq!(loader.bit_depth(), 8);
        assert_eq!(loader.chunk_size, 256);
        assert_eq!(loader.supported_extensions(), vec!["tif"]);

        // Invalid parameters should fail
        assert!(BigTiffLoader::new(vec![], 4, 8, Some(256)).is_err());
        assert!(BigTiffLoader::new(vec!["tif".to_string()], 0, 8, Some(256)).is_err());
        assert!(BigTiffLoader::new(vec!["tif".to_string()], 4, 7, Some(256)).is_err());
    }

    #[test]
    fn test_cir_loader() {
        let loader = BigTiffLoader::cir().unwrap();
        assert_eq!(loader.band_count(), 4);
        assert_eq!(loader.bit_depth(), 8);
        assert_eq!(loader.chunk_size, 256);
        assert!(loader.supported_extensions().contains(&"tif".to_string()));
        assert!(loader.supported_extensions().contains(&"tiff".to_string()));
    }
}