tx2_iff/

decoder.rs

//! Image decoder (Layer reconstruction)
//!
//! The decoder reconstructs an image from the three layers:
//! 1. Inverse wavelet transform (skeleton)
//! 2. Texture synthesis (deterministic noise)
//! 3. Warp field application (geometric deformation)
//! 4. Residual addition (correction)

use crate::error::Result;
use crate::format::IffImage;
use crate::texture::TextureSynthesizer;
use crate::warp::BicubicSampler;
use crate::wavelet::Cdf53Transform;
use crate::color::{ycocg_to_rgb, upsample_420, YCoCgImage, Channel};
use serde::{Deserialize, Serialize};

/// Decoder configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DecoderConfig {
    /// Enable Layer 2 (texture synthesis)
    pub enable_layer2: bool,
    /// Enable Layer 3 (warp fields)
    pub enable_layer3: bool,
    /// Enable residual correction
    pub enable_residual: bool,
}

impl Default for DecoderConfig {
    fn default() -> Self {
        DecoderConfig {
            enable_layer2: true,
            enable_layer3: true,
            enable_residual: true,
        }
    }
}

/// Image decoder
pub struct Decoder {
    config: DecoderConfig,
}

impl Decoder {
    /// Create a new decoder with given configuration
    pub fn new(config: DecoderConfig) -> Self {
        Decoder { config }
    }

    /// Decode an IFF image to RGB buffer
    pub fn decode(&self, iff_image: &IffImage) -> Result<Vec<[u8; 3]>> {
        let width = iff_image.header.width as usize;
        let height = iff_image.header.height as usize;

        // Layer 1: Inverse wavelet transform
        let transform = Cdf53Transform::new(iff_image.header.wavelet_levels as usize);
        
        let mut image = if iff_image.header.flags.ycocg_420 {
            // Decode Y (Luma)
            let y_coeffs = iff_image.layer1.y.to_dense()?;
            let y_data = transform.inverse(&y_coeffs[0], width, height)?;
            
            // Decode Co (Chroma Orange) - Subsampled
            let co_width = iff_image.layer1.co.width as usize;
            let co_height = iff_image.layer1.co.height as usize;
            let co_coeffs = iff_image.layer1.co.to_dense()?;
            let co_data = transform.inverse(&co_coeffs[0], co_width, co_height)?;
            
            // Decode Cg (Chroma Green) - Subsampled
            let cg_width = iff_image.layer1.cg.width as usize;
            let cg_height = iff_image.layer1.cg.height as usize;
            let cg_coeffs = iff_image.layer1.cg.to_dense()?;
            let cg_data = transform.inverse(&cg_coeffs[0], cg_width, cg_height)?;
            
            // Upsample Chroma
            let co_channel = Channel { width: co_width, height: co_height, data: co_data };
            let cg_channel = Channel { width: cg_width, height: cg_height, data: cg_data };
            
            let co_up = upsample_420(&co_channel, width, height);
            let cg_up = upsample_420(&cg_channel, width, height);
            
            // Convert YCoCg to RGB
            let ycocg = YCoCgImage {
                width,
                height,
                y: Channel { width, height, data: y_data },
                co: co_up,
                cg: cg_up,
            };
            
            ycocg_to_rgb(&ycocg)?
        } else {
            // Legacy/RGB Mode
            // In this mode, we stored R in Y, G in Co, B in Cg
            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 pixels = Vec::with_capacity(width * height);
            for i in 0..(width * height) {
                pixels.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,
                ]);
            }
            pixels
        };

        // Layer 2: Texture synthesis (if enabled)
        if self.config.enable_layer2 && !iff_image.layer2.regions.is_empty() {
            let synthesizer = TextureSynthesizer::new();
            for region in &iff_image.layer2.regions {
                // Synthesize texture for each region
                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 < image.len() {
                            let color = synthesizer.synthesize_region_pixel(region, x, y);
                            image[idx] = color;
                        }
                    }
                }
            }
        }

        // Layer 3: Warp field (if enabled)
        if self.config.enable_layer3 && !iff_image.layer3.vortices.is_empty() {
            // Create warped buffer
            let mut warped = image.clone();

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

                    // Sample with bicubic interpolation
                    let color = BicubicSampler::sample(&image, width, height, src_x, src_y);

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

            image = warped;
        }

        // Layer 4: Residual correction (if enabled)
        if self.config.enable_residual && !iff_image.residual.data.is_empty() {
            let residual_pixels = iff_image.residual.to_dense()?;
            
            if residual_pixels.len() != image.len() {
                 log::warn!("Residual size mismatch: expected {}, got {}", image.len(), residual_pixels.len());
            }

            for (i, rgb) in residual_pixels.iter().enumerate() {
                if i < image.len() {
                    for c in 0..3 {
                        let corrected = (image[i][c] as i32 + rgb[c] as i32).clamp(0, 255);
                        image[i][c] = corrected as u8;
                    }
                }
            }
        }

        Ok(image)
    }

    /// Decode to a dynamic image (convenience method)
    #[cfg(feature = "decoder")]
    pub fn decode_to_image(&self, iff_image: &IffImage) -> Result<image::DynamicImage> {
        let buffer = self.decode(iff_image)?;
        let width = iff_image.header.width;
        let height = iff_image.header.height;

        // Convert to ImageBuffer
        let mut img_buffer = image::RgbImage::new(width, height);

        for (i, pixel) in buffer.iter().enumerate() {
            let x = (i % width as usize) as u32;
            let y = (i / width as usize) as u32;
            img_buffer.put_pixel(x, y, image::Rgb(*pixel));
        }

        Ok(image::DynamicImage::ImageRgb8(img_buffer))
    }
}

/// Create a decoder with default configuration
impl Default for Decoder {
    fn default() -> Self {
        Decoder::new(DecoderConfig::default())
    }
}

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

    #[test]
    fn test_decoder_config_default() {
        let config = DecoderConfig::default();
        assert!(config.enable_layer2);
        assert!(config.enable_layer3);
        assert!(config.enable_residual);
    }

    #[test]
    fn test_decoder_creation() {
        let config = DecoderConfig::default();
        let decoder = Decoder::new(config);
        assert!(decoder.config.enable_layer2);
    }

    #[test]
    fn test_decode_empty_image() {
        let decoder = Decoder::default();
        let iff_image = IffImage::new(64, 64, 3);

        let result = decoder.decode(&iff_image);
        assert!(result.is_ok());

        let buffer = result.unwrap();
        assert_eq!(buffer.len(), 64 * 64);
    }
}