pbrt/src/core/image/mod.rs

use crate::utils::containers::Array2D;
use anyhow::{Result, anyhow};
use half::f16;
use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSliceMut};
use shared::Float;
use shared::Ptr;
use shared::core::color::{ColorEncoding, ColorEncodingTrait, LINEAR};
use shared::core::geometry::{Bounds2f, Point2f, Point2i};
use shared::core::image::{Image, ImageBase, PixelFormat, Pixels, WrapMode, WrapMode2D};
use shared::utils::math::square;
use smallvec::{SmallVec, smallvec};
use std::ops::{Deref, DerefMut};
use std::sync::Arc;

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 struct HostImage {
    pub inner: Image,
    pub channel_names: Vec<String>,
}

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

impl HostImage {
    pub fn from_u8(
        data: &[u8],
        resolution: Point2i,
        channel_names: &[impl AsRef<str>],
        encoding: ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        Self {
            inner: Image::from_u8(data, resolution, n_channels, encoding),
            channel_names: channel_names.iter().map(|s| s.as_ref().to_string()).collect(),
        }
    }

    pub fn from_f32(
        data: &[f32],
        resolution: Point2i,
        channel_names: &[impl AsRef<str>],
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        Self {
            inner: Image::from_f32(data, resolution, n_channels),
            channel_names: channel_names.iter().map(|s| s.as_ref().to_string()).collect(),
        }
    }

    pub fn from_f16(
        data: &[u16],
        resolution: Point2i,
        channel_names: &[impl AsRef<str>],
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        Self {
            inner: Image::from_f16(data, resolution, n_channels),
            channel_names: channel_names.iter().map(|s| s.as_ref().to_string()).collect(),
        }
    }

    pub fn new(
        format: PixelFormat,
        resolution: Point2i,
        channel_names: &[impl AsRef<str>],
        encoding: ColorEncoding,
    ) -> Self {
        let n_channels = channel_names.len() as i32;
        Self {
            inner: Image::new(format, resolution, n_channels, encoding),
            channel_names: channel_names.iter().map(|s| s.as_ref().to_string()).collect(),
        }
    }

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

        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);
        }

        Self::from_f32(&data, resolution, channel_names)
    }

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

    pub fn get_channel_desc(
        &self,
        requested: &[impl AsRef<str> + std::fmt::Display],
    ) -> Result<ImageChannelDesc> {
        let mut offset = Vec::with_capacity(requested.len());
        for req in requested {
            match self.channel_names.iter().position(|n| n == req.as_ref()) {
                Some(idx) => offset.push(idx),
                None => {
                    return Err(anyhow!(
                        "Missing channel '{}'. Available: {:?}",
                        req,
                        self.channel_names
                    ));
                }
            }
        }
        Ok(ImageChannelDesc { offset })
    }

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

    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_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.inner.remap_pixel_coords(&mut p, wrap_mode) {
            return ImageChannelValues(SmallVec::from_elem(0.0, self.inner.n_channels() as usize));
        }
        let offset = self.inner.pixel_offset(p);
        let nc = self.inner.n_channels() as usize;
        let mut values = SmallVec::with_capacity(nc);
        for i in 0..nc {
            values.push(unsafe { self.inner.pixels.read(offset + i, &self.inner.encoding) });
        }
        ImageChannelValues(values)
    }

    pub fn get_channels_with_desc(
        &self,
        p: Point2i,
        desc: &ImageChannelDesc,
        wrap_mode: WrapMode2D,
    ) -> ImageChannelValues {
        let mut pp = p;
        if !self.inner.remap_pixel_coords(&mut pp, wrap_mode) {
            return ImageChannelValues(SmallVec::from_elem(0.0, desc.offset.len()));
        }
        let pixel_offset = self.inner.pixel_offset(pp);
        let mut values = SmallVec::with_capacity(desc.offset.len());
        for &c in &desc.offset {
            values.push(unsafe { self.inner.pixels.read(pixel_offset + c, &self.inner.encoding) });
        }
        ImageChannelValues(values)
    }

    pub fn set_channels(&mut self, p: Point2i, values: &ImageChannelValues) {
        for i in 0..values.len() {
            self.inner.set_channel(p, i as i32, values[i]);
        }
    }

    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.inner.resolution();
        let pixel_count = (res.x() * res.y()) as usize;
        let src_nc = self.inner.n_channels() as usize;
        let dst_nc = desc.offset.len();

        // Always produce f32 output for simplicity
        let mut dst = vec![0.0f32; pixel_count * dst_nc];
        for i in 0..pixel_count {
            let src_offset = i * src_nc;
            for (out_idx, &in_c) in desc.offset.iter().enumerate() {
                dst[i * dst_nc + out_idx] = unsafe {
                    self.inner.pixels.read(src_offset + in_c, &self.inner.encoding)
                };
            }
        }

        Self::from_f32(&dst, res, &new_names)
    }

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

        let mut dist = Array2D::new_dims(width, height);

        dist.as_mut_slice()
            .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(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 domain = Bounds2f::from_points(Point2f::new(0.0, 0.0), Point2f::new(1.0, 1.0));
        self.get_sampling_distribution(|_| 1.0, domain)
    }

    pub fn has_any_infinite_pixels(&self) -> bool {
        self.inner.has_any_infinite_pixels()
    }

    pub fn has_any_nan_pixels(&self) -> bool {
        self.inner.has_any_nan_pixels()
    }

    pub fn mse(
        &self,
        desc: &ImageChannelDesc,
        ref_img: &HostImage,
        generate_mse_image: bool,
    ) -> (ImageChannelValues, Option<HostImage>) {
        let res = self.inner.resolution();
        assert_eq!(res, ref_img.inner.resolution());

        let ref_desc = ref_img
            .get_channel_desc(&self.channel_names_from_desc(desc))
            .expect("Channels not found in reference image");

        let width = res.x() as usize;
        let height = res.y() as usize;
        let n_channels = desc.size();

        let mut sum_se: Vec<f64> = vec![0.0; n_channels];
        let mut mse_pixels = if generate_mse_image {
            vec![0.0f32; width * height * n_channels]
        } else {
            Vec::new()
        };

        for y in 0..res.y() {
            for x in 0..res.x() {
                let v = self.get_channels_with_desc(
                    Point2i::new(x, y), desc, WrapMode::Clamp.into(),
                );
                let v_ref = ref_img.get_channels_with_desc(
                    Point2i::new(x, y), &ref_desc, WrapMode::Clamp.into(),
                );
                for c in 0..n_channels {
                    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 = (res.x() * res.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 {
            let names = self.channel_names_from_desc(desc);
            Some(HostImage::from_f32(&mse_pixels, res, &names))
        } else {
            None
        };

        (ImageChannelValues(mse_values), mse_image)
    }

    pub fn image(&self) -> &Image {
        &self.inner
    }

    pub fn into_image(self) -> Image {
        self.inner
    }
}