pbrt/src/core/image/mod.rs

use anyhow::Result;
use half::f16;
use shared::Float;
use shared::core::color::ColorEncoding;
use shared::core::color::LINEAR;
use shared::core::geometry::{Bounds2f, Point2f, Point2i};
use shared::core::image::{DeviceImage, ImageBase, PixelFormat, Pixels, WrapMode, WrapMode2D};
use shared::utils::containers::Array2D;
use shared::utils::math::square;
use smallvec::{SmallVec, smallvec};
use std::ops::{Deref, DerefMut};

pub mod io;
pub mod metadata;
pub mod ops;
pub mod pixel;

pub use io::ImageIO;
pub use metadata::*;

#[derive(Clone, Debug, Default)]
pub struct ImageChannelValues(pub SmallVec<[Float; 4]>);

impl ImageChannelValues {
    pub fn average(&self) -> Float {
        if self.0.is_empty() {
            return 0.0;
        }
        let sum: Float = self.0.iter().sum();
        sum / (self.0.len() as Float)
    }

    pub fn max_value(&self) -> Float {
        self.0.iter().fold(Float::MIN, |a, &b| a.max(b))
    }
}

impl From<&[Float]> for ImageChannelValues {
    fn from(slice: &[Float]) -> Self {
        Self(SmallVec::from_slice(slice))
    }
}

impl From<Vec<Float>> for ImageChannelValues {
    fn from(vec: Vec<Float>) -> Self {
        Self(SmallVec::from_vec(vec))
    }
}

impl<const N: usize> From<[Float; N]> for ImageChannelValues {
    fn from(arr: [Float; N]) -> Self {
        Self(SmallVec::from_slice(&arr))
    }
}

impl Deref for ImageChannelValues {
    type Target = SmallVec<[Float; 4]>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for ImageChannelValues {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

#[derive(Debug, Clone)]
pub enum PixelStorage {
    U8(Box<[u8]>),
    F16(Box<[f16]>),
    F32(Box<[f32]>),
}

impl PixelStorage {
    pub fn as_pixels(&self) -> Pixels {
        match self {
            PixelStorage::U8(data) => Pixels::U8(data.as_ptr()),
            PixelStorage::F16(data) => Pixels::F16(data.as_ptr() as *const u16),
            PixelStorage::F32(data) => Pixels::F32(data.as_ptr()),
        }
    }

    pub fn format(&self) -> PixelFormat {
        match self {
            PixelStorage::U8(_) => PixelFormat::U8,
            PixelStorage::F16(_) => PixelFormat::F16,
            PixelStorage::F32(_) => PixelFormat::F32,
        }
    }

    pub fn len(&self) -> usize {
        match self {
            PixelStorage::U8(d) => d.len(),
            PixelStorage::F16(d) => d.len(),
            PixelStorage::F32(d) => d.len(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Image {
    storage: PixelStorage,
    channel_names: Vec<String>,
    pub device: DeviceImage,
}

// impl Deref for Image {
//     type Target = DeviceImage;
//     #[inline]
//     fn deref(&self) -> &Self::Target {
//         &self.device
//     }
// }

#[derive(Debug, Clone)]
pub struct ImageAndMetadata {
    pub image: Image,
    pub metadata: ImageMetadata,
}

impl Image {
    // Constructors
    fn from_storage(
        storage: PixelStorage,
        resolution: Point2i,
        channel_names: &[&str],
        encoding: ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        let expected = (resolution.x() * resolution.y()) as usize * n_channels as usize;
        assert_eq!(storage.len(), expected, "Pixel data size mismatch");

        let device = DeviceImage {
            base: ImageBase {
                format: storage.format(),
                encoding,
                resolution,
                n_channels,
            },
            pixels: storage.as_pixels(),
        };

        Self {
            storage,
            channel_names,
            device,
        }
    }

    pub fn from_u8(
        data: Vec<u8>,
        resolution: Point2i,
        channel_names: &[&str],
        encoding: ColorEncoding,
    ) -> Self {
        Self::from_storage(
            PixelStorage::U8(data.into_boxed_slice()),
            resolution,
            channel_names,
            encoding,
        )
    }

    pub fn from_u8(
        data: Vec<u8>,
        resolution: Point2i,
        channel_names: &[&str],
        encoding: ColorEncoding,
    ) -> Self {
        Self::from_storage(
            PixelStorage::U8(data.into_boxed_slice()),
            resolution,
            channel_names,
            encoding,
        )
    }

    pub fn from_f16(data: Vec<half::f16>, resolution: Point2i, channel_names: &[&str]) -> Self {
        Self::from_storage(
            PixelStorage::F16(data.into_boxed_slice()),
            resolution,
            channel_names,
            ColorEncoding::Linear,
        )
    }

    pub fn from_f32(data: Vec<f32>, resolution: Point2i, channel_names: &[&str]) -> Self {
        Self::from_storage(
            PixelStorage::F32(data.into_boxed_slice()),
            resolution,
            channel_names,
            ColorEncoding::Linear,
        )
    }

    pub fn new(
        format: PixelFormat,
        resolution: Point2i,
        channel_names: &[&str],
        encoding: Arc<ColorEncoding>,
    ) -> Self {
        let n_channels = channel_names.len();
        let pixel_count = (resolution.x * resolution.y) as usize * n_channels;
        let owned_names: Vec<String> = channel_names.iter().map(|s| s.to_string()).collect();

        let storage = match format {
            PixelFormat::U8 => PixelStorage::U8(vec![0; pixel_count]),
            PixelFormat::F16 => PixelStorage::F16(vec![0; pixel_count]),
            PixelFormat::F32 => PixelStorage::F32(vec![0.0; pixel_count]),
        };

        Self::from_storage(storage, resolution, owned_names, encoding)
    }

    pub fn new_constant(resolution: Point2i, channel_names: &[&str], values: &[f32]) -> Self {
        let n_channels = channel_names.len();
        if values.len() != n_channels {
            panic!(
                "Image::new_constant: values length ({}) must match channel count ({})",
                values.len(),
                n_channels
            );
        }

        let n_pixels = (resolution.x * resolution.y) as usize;

        let mut data = Vec::with_capacity(n_pixels * n_channels);

        for _ in 0..n_pixels {
            data.extend_from_slice(values);
        }

        let owned_names: Vec<String> = channel_names.iter().map(|s| s.to_string()).collect();

        Self::from_f32(data, resolution, owned_names)
    }

    // Access
    pub fn device_image(&self) -> &DeviceImage {
        &self.device
    }

    pub fn resolution(&self) -> Point2i {
        self.base.resolution
    }

    fn n_channels(&self) -> i32 {
        self.base.n_channels
    }

    pub fn format(&self) -> PixelFormat {
        self.device.base.format
    }

    pub fn channel_names(&self) -> Vec<&str> {
        self.channel_names.iter().map(|s| s.as_str()).collect()
    }

    pub fn encoding(&self) -> ColorEncoding {
        self.view.encoding
    }

    fn pixel_offset(&self, p: Point2i) -> usize {
        let width = self.resolution().x() as usize;
        let idx = p.y() as usize * width + p.x() as usize;
        idx * self.n_channels() as usize
    }

    // Read
    pub fn get_channel(&self, p: Point2i, c: i32) -> Float {
        self.get_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn get_channel_with_wrap(&self, mut p: Point2i, c: i32, wrap_mode: WrapMode2D) -> Float {
        if !self.device.base.remap_pixel_coords(&mut p, wrap_mode) {
            return 0.0;
        }

        let offset = self.pixel_offset(p) + c as usize;

        match &self.storage {
            PixelStorage::U8(data) => self.device.base.encoding.to_linear_scalar(data[offset]),
            PixelStorage::F16(data) => data[offset].to_f32(),
            PixelStorage::F32(data) => data[offset],
        }
    }

    pub fn get_channels(&self, p: Point2i) -> ImageChannelValues {
        self.get_channels_with_wrap(p, WrapMode::Clamp.into())
    }

    pub fn get_channels_with_wrap(
        &self,
        mut p: Point2i,
        wrap_mode: WrapMode2D,
    ) -> ImageChannelValues {
        if !self.device.base.remap_pixel_coords(&mut p, wrap_mode) {
            return ImageChannelValues(smallvec![0.0; self.n_channels() as usize]);
        }

        let offset = self.pixel_offset(p);
        let nc = self.n_channels() as usize;
        let mut values = SmallVec::with_capacity(nc);

        match &self.storage {
            PixelStorage::U8(data) => {
                for i in 0..nc {
                    values.push(self.device.base.encoding.to_linear_scalar(data[offset + i]));
                }
            }
            PixelStorage::F16(data) => {
                for i in 0..nc {
                    values.push(data[offset + i].to_f32());
                }
            }
            PixelStorage::F32(data) => {
                for i in 0..nc {
                    values.push(data[offset + i]);
                }
            }
        }

        ImageChannelValues(values)
    }

    // Write
    pub fn set_channel(&mut self, p: Point2i, c: i32, mut value: Float) {
        if value.is_nan() {
            value = 0.0;
        }

        let res = self.resolution();
        if p.x() < 0 || p.x() >= res.x() || p.y() < 0 || p.y() >= res.y() {
            return;
        }

        let offset = self.pixel_offset(p) + c as usize;

        match &mut self.storage {
            PixelStorage::U8(data) => {
                let data = Box::as_mut(data);
                data[offset] = self.device.base.encoding.from_linear_scalar(value);
            }
            PixelStorage::F16(data) => {
                let data = Box::as_mut(data);
                data[offset] = f16::from_f32(value);
            }
            PixelStorage::F32(data) => {
                let data = Box::as_mut(data);
                data[offset] = value;
            }
        }
    }

    // Descriptions
    pub fn get_channels_with_desc(
        &self,
        p: Point2i,
        desc: &ImageChannelDesc,
        wrap_mode: WrapMode2D,
    ) -> ImageChannelValues {
        let mut pp = p;
        if !self.device.base.remap_pixel_coords(&mut pp, wrap_mode) {
            return ImageChannelValues(smallvec![0.0; desc.offset.len()]);
        }

        let pixel_offset = self.pixel_offset(pp);
        let mut values = SmallVec::with_capacity(desc.offset.len());

        match &self.storage {
            PixelStorage::U8(data) => {
                for &c in &desc.offset {
                    let raw = data[pixel_offset + c];
                    values.push(self.device.base.encoding.to_linear_scalar(raw));
                }
            }
            PixelStorage::F16(data) => {
                for &c in &desc.offset {
                    values.push(data[pixel_offset + c].to_f32());
                }
            }
            PixelStorage::F32(data) => {
                for &c in &desc.offset {
                    values.push(data[pixel_offset + c]);
                }
            }
        }

        ImageChannelValues(values)
    }

    pub fn channel_names_from_desc(&self, desc: &ImageChannelDesc) -> Vec<&str> {
        desc.offset
            .iter()
            .map(|&i| self.channel_names[i].as_str())
            .collect()
    }

    pub fn get_channel_desc(&self, requested_channels: &[&str]) -> Result<ImageChannelDesc> {
        let mut offset = Vec::with_capacity(requested_channels.len());

        for &req in requested_channels.iter() {
            match self.channel_names.iter().position(|n| n == req) {
                Some(idx) => {
                    offset.push(idx);
                }
                None => {
                    return Err(format!(
                        "Missing channel '{}'. Available: {:?}",
                        req, self.channel_names
                    ));
                }
            }
        }

        Ok(ImageChannelDesc { offset })
    }

    pub fn all_channels_desc(&self) -> ImageChannelDesc {
        ImageChannelDesc {
            offset: (0..self.n_channels() as usize).collect(),
        }
    }

    pub fn select_channels(&self, desc: &ImageChannelDesc) -> Self {
        let new_names: Vec<String> = desc
            .offset
            .iter()
            .map(|&i| self.channel_names[i].clone())
            .collect();

        let res = self.resolution();
        let pixel_count = (res.x() * res.y()) as usize;
        let src_nc = self.n_channels() as usize;
        let dst_nc = desc.offset.len();

        let new_storage = match &self.storage {
            PixelStorage::U8(src) => {
                let mut dst = vec![0u8; pixel_count * dst_nc];
                for i in 0..pixel_count {
                    for (out_idx, &in_c) in desc.offset.iter().enumerate() {
                        dst[i * dst_nc + out_idx] = src[i * src_nc + in_c];
                    }
                }
                PixelStorage::U8(dst.into_boxed_slice())
            }
            PixelStorage::F16(src) => {
                let mut dst = vec![f16::ZERO; pixel_count * dst_nc];
                for i in 0..pixel_count {
                    for (out_idx, &in_c) in desc.offset.iter().enumerate() {
                        dst[i * dst_nc + out_idx] = src[i * src_nc + in_c];
                    }
                }
                PixelStorage::F16(dst.into_boxed_slice())
            }
            PixelStorage::F32(src) => {
                let mut dst = vec![0.0f32; pixel_count * dst_nc];
                for i in 0..pixel_count {
                    for (out_idx, &in_c) in desc.offset.iter().enumerate() {
                        dst[i * dst_nc + out_idx] = src[i * src_nc + in_c];
                    }
                }
                PixelStorage::F32(dst.into_boxed_slice())
            }
        };

        Self::from_storage(new_storage, res, new_names, self.encoding())
    }

    pub fn get_sampling_distribution<F>(&self, dxd_a: F, domain: Bounds2f) -> Array2D<Float>
    where
        F: Fn(Point2f) -> Float + Sync + Send,
    {
        let width = self.resolution().x();
        let height = self.resolution().y();

        let mut dist = Array2D::new_with_dims(width as usize, height as usize);

        dist.values
            .par_chunks_mut(width as usize)
            .enumerate()
            .for_each(|(y, row)| {
                let y = y as i32;

                for (x, out_val) in row.iter_mut().enumerate() {
                    let x = x as i32;

                    let value = self.get_channels_default(Point2i::new(x, y)).average();

                    let u = (x as Float + 0.5) / width as Float;
                    let v = (y as Float + 0.5) / height as Float;
                    let p = domain.lerp(Point2f::new(u, v));
                    *out_val = value * dxd_a(p);
                }
            });

        dist
    }

    pub fn get_sampling_distribution_uniform(&self) -> Array2D<Float> {
        let default_domain = Bounds2f::from_points(Point2f::new(0.0, 0.0), Point2f::new(1.0, 1.0));

        self.get_sampling_distribution(|_| 1.0, default_domain)
    }

    pub fn mse(
        &self,
        desc: ImageChannelDesc,
        ref_img: &Image,
        generate_mse_image: bool,
    ) -> (ImageChannelValues, Option<Image>) {
        let mut sum_se: Vec<f64> = vec![0.; desc.size()];
        let names_ref = self.channel_names_from_desc(&desc);
        let ref_desc = ref_img
            .get_channel_desc(&self.channel_names_from_desc(&desc))
            .expect("Channels not found in image");
        assert_eq!(self.resolution(), ref_img.resolution());

        let width = self.resolution.x() as usize;
        let height = self.resolution.y() as usize;
        let n_channels = desc.offset.len();
        let mut mse_pixels = if generate_mse_image {
            vec![0.0f32; width * height * n_channels]
        } else {
            Vec::new()
        };

        for y in 0..self.resolution().y() {
            for x in 0..self.resolution().x() {
                let v =
                    self.get_channel_with_desc(Point2i::new(x, y), &desc, WrapMode::Clamp.into());
                let v_ref = self.get_channel_with_desc(
                    Point2i::new(x, y),
                    &ref_desc,
                    WrapMode::Clamp.into(),
                );
                for c in 0..desc.size() {
                    let se = square(v[c] as f64 - v_ref[c] as f64);
                    if se.is_infinite() {
                        continue;
                    }
                    sum_se[c] += se;
                    if generate_mse_image {
                        let idx = (y as usize * width + x as usize) * n_channels + c;
                        mse_pixels[idx] = se as f32;
                    }
                }
            }
        }

        let pixel_count = (self.resolution().x() * self.resolution.y()) as f64;
        let mse_values: SmallVec<[Float; 4]> =
            sum_se.iter().map(|&s| (s / pixel_count) as Float).collect();

        let mse_image = if generate_mse_image {
            Some(Image::new(
                PixelFormat::F32,
                self.resolution,
                &names_ref,
                LINEAR,
            ))
        } else {
            None
        };

        (ImageChannelValues(mse_values), mse_image)
    }

    pub fn update_view_pointers(&mut self) {
        self.view.pixels = match &self._storage {
            PixelStorage::U8(vec) => Pixels::U8(vec.as_ptr()),
            PixelStorage::F16(vec) => Pixels::F16(vec.as_ptr() as *const u16),
            PixelStorage::F32(vec) => Pixels::F32(vec.as_ptr()),
        };
    }

    pub fn has_any_infinite_pixels(&self) -> bool {
        if self.format() == PixelFormat::Float {
            return false;
        }

        for y in 0..self.resolution().y() {
            for x in 0..self.resolution().x() {
                for c in 0..self.n_channels() {
                    if self.get_channel(Point2i::new(x, y), c).is_infinite() {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    pub fn has_any_nan_pixels(&self) -> bool {
        if self.format() == PixelFormat::Float {
            return false;
        }

        for y in 0..self.resolution().y() {
            for x in 0..self.resolution().x() {
                for c in 0..self.n_channels() {
                    if self.get_channel(Point2i::new(x, y), c).is_nan() {
                        return true;
                    }
                }
            }
        }
        return false;
    }
}

impl std::ops::Deref for Image {
    type Target = DeviceImage;

    fn deref(&self) -> &DeviceImage {
        &self.device
    }
}