tx2_iff/

encoder.rs

//! Image encoder (Layer composition and optimization)
//!
//! The encoder analyzes an input image and decomposes it into three layers:
//! 1. Wavelet skeleton (sharp edges, text, structure)
//! 2. Texture synthesis regions (high-entropy content)
//! 3. Warp fields (geometric deformation)
//!
//! Plus a sparse residual for anything the layers can't represent.

use crate::error::Result;
use crate::format::{IffImage, Layer1, Residual};
use crate::prime::QuantizationTable;
use crate::wavelet::{Cdf53Transform, WaveletDecomposition};
use crate::color::{rgb_to_ycocg, subsample_420, ycocg_to_rgb, upsample_420, YCoCgImage};
use serde::{Deserialize, Serialize};

/// Encoder configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncoderConfig {
    /// Wavelet decomposition levels
    pub wavelet_levels: usize,
    /// Base quantization value
    pub base_quantization: u16,
    /// Texture region minimum size
    pub texture_min_size: usize,
    /// Texture synthesis iterations
    pub texture_iterations: usize,
    /// Residual threshold (RMSE)
    pub residual_threshold: f32,
    /// Texture entropy threshold (0.0-1.0)
    pub texture_entropy_threshold: f32,
    /// Enable Layer 2 (texture synthesis)
    pub enable_layer2: bool,
    /// Enable Layer 3 (warp fields)
    pub enable_layer3: bool,
    /// Use YCoCg-R color space with 4:2:0 subsampling
    pub use_ycocg_420: bool,
}

impl Default for EncoderConfig {
    fn default() -> Self {
        EncoderConfig {
            wavelet_levels: 5,
            base_quantization: 8,
            texture_min_size: 32,
            texture_iterations: 100,
            residual_threshold: 40.0,
            texture_entropy_threshold: 0.05,
            enable_layer2: true,
            enable_layer3: true,
            use_ycocg_420: true,
        }
    }
}

/// Image encoder
pub struct Encoder {
    config: EncoderConfig,
}

impl Encoder {
    /// Create a new encoder with given configuration
    pub fn new(config: EncoderConfig) -> Self {
        Encoder { config }
    }

    /// Encode an image to IFF format
    #[cfg(feature = "encoder")]
    pub fn encode(&self, image: &image::DynamicImage) -> Result<IffImage> {
        let rgb_image = image.to_rgb8();
        let width = rgb_image.width() as usize;
        let height = rgb_image.height() as usize;

        // Convert to working format
        let image_data: Vec<[u8; 3]> = rgb_image
            .pixels()
            .map(|p| [p[0], p[1], p[2]])
            .collect();

        // Create IFF image structure
        let mut iff_image = IffImage::new(width as u32, height as u32, self.config.wavelet_levels as u8);
        iff_image.header.flags.ycocg_420 = self.config.use_ycocg_420;

        // Layer 1: Wavelet transform
        log::info!("Encoding Layer 1: Wavelet skeleton");
        self.encode_layer1(&mut iff_image, &image_data, width, height)?;

        // Reconstruct from Layer 1 to calculate residual
        let mut reconstruction = self.decode_layer1(&iff_image, width, height)?;

        // Layer 2: Texture synthesis (if enabled)
        if self.config.enable_layer2 {
            log::info!("Encoding Layer 2: Texture synthesis");
            self.encode_layer2(&mut iff_image, &image_data, &mut reconstruction, width, height)?;
        }

        // Layer 3: Warp field (if enabled)
        if self.config.enable_layer3 {
            log::info!("Encoding Layer 3: Warp fields");
            self.encode_layer3(&mut iff_image, &image_data, &mut reconstruction, width, height)?;
        }

        // Calculate final residual
        log::info!("Calculating residual");
        self.calculate_residual(&mut iff_image, &image_data, &reconstruction, width, height)?;

        Ok(iff_image)
    }

    /// Encode Layer 1: Wavelet skeleton
    fn encode_layer1(
        &self,
        iff_image: &mut IffImage,
        image_data: &[[u8; 3]],
        width: usize,
        height: usize,
    ) -> Result<()> {
        let transform = Cdf53Transform::new(self.config.wavelet_levels);
        let quantization_table = QuantizationTable::new(self.config.base_quantization);

        if self.config.use_ycocg_420 {
            // Convert to YCoCg
            let ycocg = rgb_to_ycocg(image_data, width, height);
            
            // Subsample chroma
            let co_sub = subsample_420(&ycocg.co);
            let cg_sub = subsample_420(&ycocg.cg);

            // Transform Y
            let mut y_coeffs = transform.forward(&ycocg.y.data, width, height)?;
            transform.quantize(&mut y_coeffs, width, height, &quantization_table);
            
            // Transform Co
            let mut co_coeffs = transform.forward(&co_sub.data, co_sub.width, co_sub.height)?;
            transform.quantize(&mut co_coeffs, co_sub.width, co_sub.height, &quantization_table);
            
            // Transform Cg
            let mut cg_coeffs = transform.forward(&cg_sub.data, cg_sub.width, cg_sub.height)?;
            transform.quantize(&mut cg_coeffs, cg_sub.width, cg_sub.height, &quantization_table);

            // Store
            iff_image.layer1 = Layer1 {
                y: WaveletDecomposition::from_dense(width as u32, height as u32, self.config.wavelet_levels, &[y_coeffs])?,
                co: WaveletDecomposition::from_dense(co_sub.width as u32, co_sub.height as u32, self.config.wavelet_levels, &[co_coeffs])?,
                cg: WaveletDecomposition::from_dense(cg_sub.width as u32, cg_sub.height as u32, self.config.wavelet_levels, &[cg_coeffs])?,
            };
        } else {
            // Standard RGB encoding (legacy/lossless mode)
            let mut channel_coeffs = Vec::with_capacity(3);
            for c in 0..3 {
                let channel: Vec<i32> = image_data
                    .iter()
                    .map(|p| p[c] as i32 - 128)
                    .collect();

                let mut coeffs = transform.forward(&channel, width, height)?;
                transform.quantize(&mut coeffs, width, height, &quantization_table);
                channel_coeffs.push(coeffs);
            }
            
            // We need to fit this into the new Layer1 structure
            // We'll put R in Y, G in Co, B in Cg (hacky but works for storage)
            iff_image.layer1 = Layer1 {
                y: WaveletDecomposition::from_dense(width as u32, height as u32, self.config.wavelet_levels, &[channel_coeffs[0].clone()])?,
                co: WaveletDecomposition::from_dense(width as u32, height as u32, self.config.wavelet_levels, &[channel_coeffs[1].clone()])?,
                cg: WaveletDecomposition::from_dense(width as u32, height as u32, self.config.wavelet_levels, &[channel_coeffs[2].clone()])?,
            };
        }

        Ok(())
    }

    /// Decode Layer 1 for residual calculation
    fn decode_layer1(
        &self,
        iff_image: &IffImage,
        width: usize,
        height: usize,
    ) -> Result<Vec<[u8; 3]>> {
        let transform = Cdf53Transform::new(self.config.wavelet_levels);

        if iff_image.header.flags.ycocg_420 {
            // Decompress coefficients
            let y_coeffs = iff_image.layer1.y.to_dense()?;
            let co_coeffs = iff_image.layer1.co.to_dense()?;
            let cg_coeffs = iff_image.layer1.cg.to_dense()?;
            
            // Inverse transform
            let y_data = transform.inverse(&y_coeffs[0], width, height)?;
            let co_data = transform.inverse(&co_coeffs[0], iff_image.layer1.co.width as usize, iff_image.layer1.co.height as usize)?;
            let cg_data = transform.inverse(&cg_coeffs[0], iff_image.layer1.cg.width as usize, iff_image.layer1.cg.height as usize)?;
            
            // Upsample chroma
            let co_channel = crate::color::Channel { width: iff_image.layer1.co.width as usize, height: iff_image.layer1.co.height as usize, data: co_data };
            let cg_channel = crate::color::Channel { width: iff_image.layer1.cg.width as usize, height: iff_image.layer1.cg.height as usize, data: cg_data };
            
            let co_up = upsample_420(&co_channel, width, height);
            let cg_up = upsample_420(&cg_channel, width, height);
            
            // Convert to RGB
            let ycocg = YCoCgImage {
                width,
                height,
                y: crate::color::Channel { width, height, data: y_data },
                co: co_up,
                cg: cg_up,
            };
            
            ycocg_to_rgb(&ycocg)
        } else {
            // Legacy RGB mode
            let r_coeffs = iff_image.layer1.y.to_dense()?;
            let g_coeffs = iff_image.layer1.co.to_dense()?;
            let b_coeffs = iff_image.layer1.cg.to_dense()?;
            
            let r_data = transform.inverse(&r_coeffs[0], width, height)?;
            let g_data = transform.inverse(&g_coeffs[0], width, height)?;
            let b_data = transform.inverse(&b_coeffs[0], width, height)?;
            
            let mut result = Vec::with_capacity(width * height);
            for i in 0..(width * height) {
                result.push([
                    (r_data[i] + 128).clamp(0, 255) as u8,
                    (g_data[i] + 128).clamp(0, 255) as u8,
                    (b_data[i] + 128).clamp(0, 255) as u8,
                ]);
            }
            Ok(result)
        }
    }

    /// Encode Layer 2: Texture synthesis
    fn encode_layer2(
        &self,
        iff_image: &mut IffImage,
        original: &[[u8; 3]],
        reconstruction: &mut Vec<[u8; 3]>,
        width: usize,
        height: usize,
    ) -> Result<()> {
        use crate::texture::TextureAnalyzer;

        let analyzer = TextureAnalyzer::new(self.config.texture_entropy_threshold);

        // Detect high-entropy regions
        let regions = analyzer.detect_texture_regions(
            original,
            width,
            height,
            self.config.texture_min_size,
        );

        log::info!("Detected {} texture regions", regions.len());

        // Optimize each region
        for mut region in regions {
            let error = analyzer.optimize_region(
                original,
                &mut region,
                width,
                self.config.texture_iterations,
                self.config.residual_threshold,
            )?;

            // Only keep region if it improves compression
            if error < self.config.residual_threshold {
                // Apply synthesized texture to reconstruction
                self.apply_texture_region(&region, reconstruction, width, height);

                // Add to IFF
                iff_image.layer2.add_region(region);
            }
        }

        Ok(())
    }

    /// Apply a texture region to reconstruction
    fn apply_texture_region(
        &self,
        region: &crate::texture::Region,
        reconstruction: &mut [[u8; 3]],
        width: usize,
        _height: usize,
    ) {
        use crate::texture::TextureSynthesizer;

        let synthesizer = TextureSynthesizer::new();

        for y in 0..region.h {
            for x in 0..region.w {
                let global_x = region.x + x;
                let global_y = region.y + y;
                let idx = (global_y as usize) * width + (global_x as usize);

                if idx < reconstruction.len() {
                    reconstruction[idx] = synthesizer.synthesize_region_pixel(region, x, y);
                }
            }
        }
    }

    /// Encode Layer 3: Warp fields
    fn encode_layer3(
        &self,
        iff_image: &mut IffImage,
        original: &[[u8; 3]],
        reconstruction: &mut Vec<[u8; 3]>,
        width: usize,
        height: usize,
    ) -> Result<()> {
        

        // Simple warp field detection: find regions with self-similarity
        let vortices = self.detect_warp_regions(original, width, height);

        log::info!("Detected {} warp vortices", vortices.len());

        for vortex in vortices {
            iff_image.layer3.add_vortex(vortex);
        }

        // Apply warping to reconstruction
        if !iff_image.layer3.vortices.is_empty() {
            self.apply_warp_field(&iff_image.layer3, reconstruction, width, height);
        }

        Ok(())
    }

    /// Detect regions that would benefit from warping
    fn detect_warp_regions(&self, image: &[[u8; 3]], width: usize, height: usize) -> Vec<crate::warp::Vortex> {
        use crate::warp::Vortex;

        let mut vortices = Vec::new();

        // Simple detection: look for regions with high gradient magnitude
        // indicating deformation (wrinkles, folds, etc.)
        let step = 32; // Check every 32 pixels

        for y in (step..height - step).step_by(step) {
            for x in (step..width - step).step_by(step) {
                let gradient = self.calculate_gradient_magnitude(image, x, y, width);

                // High gradient suggests deformation
                if gradient > 50.0 {
                    // Create a small vortex at this location
                    let vortex = Vortex::new(
                        x as u16,
                        y as u16,
                        (gradient as i16).clamp(-1000, 1000),
                        16, // Small radius
                        128, // Medium decay
                    );
                    vortices.push(vortex);
                }
            }
        }

        // Limit number of vortices to avoid overhead
        vortices.truncate(50);

        vortices
    }

    /// Calculate gradient magnitude at a point
    fn calculate_gradient_magnitude(&self, image: &[[u8; 3]], x: usize, y: usize, width: usize) -> f32 {
        let idx = y * width + x;
        if idx >= image.len() || x == 0 || y == 0 {
            return 0.0;
        }

        let idx_left = y * width + (x - 1);
        let idx_up = (y - 1) * width + x;

        if idx_left >= image.len() || idx_up >= image.len() {
            return 0.0;
        }

        let mut grad_x = 0.0f32;
        let mut grad_y = 0.0f32;

        for c in 0..3 {
            grad_x += (image[idx][c] as f32 - image[idx_left][c] as f32).abs();
            grad_y += (image[idx][c] as f32 - image[idx_up][c] as f32).abs();
        }

        (grad_x * grad_x + grad_y * grad_y).sqrt()
    }

    /// Apply warp field to reconstruction
    fn apply_warp_field(
        &self,
        warp_field: &crate::warp::WarpField,
        reconstruction: &mut [[u8; 3]],
        width: usize,
        height: usize,
    ) {
        use crate::warp::BicubicSampler;

        let original = reconstruction.to_vec();

        for y in 0..height {
            for x in 0..width {
                let (src_x, src_y) = warp_field.warp_backwards(x as u16, y as u16);

                let color = BicubicSampler::sample(&original, width, height, src_x, src_y);

                let idx = y * width + x;
                reconstruction[idx] = color;
            }
        }
    }

    /// Calculate residual between original and reconstruction
    fn calculate_residual(
        &self,
        iff_image: &mut IffImage,
        original: &[[u8; 3]],
        reconstruction: &[[u8; 3]],
        width: usize,
        height: usize,
    ) -> Result<()> {
        let mut residual_pixels = Vec::with_capacity(original.len());

        for (orig, recon) in original.iter().zip(reconstruction.iter()) {
            let mut pixel = [0i16; 3];
            let mut has_residual = false;

            for c in 0..3 {
                let diff = orig[c] as i16 - recon[c] as i16;

                // Only store if very significant (threshold increased for better compression)
                if (diff.abs() as f32) > self.config.residual_threshold {
                    pixel[c] = diff;
                    has_residual = true;
                }
            }
            
            if !has_residual {
                pixel = [0, 0, 0];
            }
            
            residual_pixels.push(pixel);
        }

        iff_image.residual = Residual::from_dense(width as u32, height as u32, &residual_pixels)?;

        log::info!("Residual compressed size: {} bytes", iff_image.residual.data.len());

        Ok(())
    }
}

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

    #[test]
    fn test_encoder_config_default() {
        let config = EncoderConfig::default();
        assert_eq!(config.wavelet_levels, 5);
        assert_eq!(config.base_quantization, 10);
        assert!(config.use_ycocg_420);
    }
}