pbrt/shared/src/core/image.rs

use crate::core::color::{ColorEncoding, ColorEncodingTrait, LINEAR};
use crate::core::geometry::{Bounds2f, Point2f, Point2fi, Point2i};
use crate::utils::math::{f16_to_f32_software, lerp, square};
use crate::{gvec_with_capacity, Float, GVec, Ptr};
use anyhow::{bail, Result};
use core::hash;
use core::ops::{Deref, DerefMut};
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum WrapMode {
    Black,
    Clamp,
    Repeat,
    OctahedralSphere,
}

impl WrapMode {
    pub fn parse(name: &str) -> Result<WrapMode> {
        match name {
            "clamp" => Ok(WrapMode::Clamp),
            "black" => Ok(WrapMode::Black),
            "repeat" => Ok(WrapMode::Repeat),
            "octahedralsphere" => Ok(WrapMode::OctahedralSphere),
            _ => bail!("{:?}: wrap mode unknown", name),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct WrapMode2D {
    pub uv: [WrapMode; 2],
}

impl From<WrapMode> for WrapMode2D {
    fn from(w: WrapMode) -> Self {
        Self { uv: [w, w] }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum PixelFormat {
    U8,
    F16,
    F32,
}

impl core::fmt::Display for PixelFormat {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            PixelFormat::U8 => write!(f, "U256"),
            PixelFormat::F16 => write!(f, "Half"),
            PixelFormat::F32 => write!(f, "Float"),
        }
    }
}

impl PixelFormat {
    pub fn is_8bit(&self) -> bool {
        matches!(self, PixelFormat::U8)
    }

    pub fn is_16bit(&self) -> bool {
        matches!(self, PixelFormat::F16)
    }

    pub fn is_32bit(&self) -> bool {
        matches!(self, PixelFormat::F32)
    }

    pub fn texel_bytes(&self) -> usize {
        match self {
            PixelFormat::U8 => 1,
            PixelFormat::F16 => 2,
            PixelFormat::F32 => 4,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug)]
pub struct Pixels {
    data: GVec<u8>,
    format: PixelFormat,
}

impl Pixels {
    pub fn new(data: GVec<u8>, format: PixelFormat) -> Self {
        Self { data, format }
    }

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

    pub fn as_ptr(&self) -> *const u8 {
        self.data.as_ptr()
    }

    pub fn len(&self) -> usize {
        self.data.len()
    }

    pub fn texel_count(&self) -> usize {
        self.data.len() / self.format.texel_bytes()
    }

    pub unsafe fn read_u8(&self, texel_offset: usize) -> u8 {
        unsafe { *self.data.as_ptr().add(texel_offset) }
    }

    pub unsafe fn read_f16(&self, texel_offset: usize) -> u16 {
        let byte_offset = texel_offset * 2;
        unsafe { *(self.data.as_ptr().add(byte_offset) as *const u16) }
    }

    pub unsafe fn read_f32(&self, texel_offset: usize) -> f32 {
        let byte_offset = texel_offset * 4;
        unsafe { *(self.data.as_ptr().add(byte_offset) as *const f32) }
    }

    pub unsafe fn read(&self, texel_offset: usize, encoding: &ColorEncoding) -> Float {
        match self.format {
            PixelFormat::U8 => encoding.to_linear_scalar(self.read_u8(texel_offset)),
            PixelFormat::F16 => f16_to_f32_software(self.read_f16(texel_offset)),
            PixelFormat::F32 => self.read_f32(texel_offset),
        }
    }

    pub unsafe fn write_u8(&mut self, texel_offset: usize, val: u8) {
        unsafe { *self.data.as_mut_ptr().add(texel_offset) = val };
    }

    pub unsafe fn write_f16(&mut self, texel_offset: usize, val: u16) {
        let byte_offset = texel_offset * 2;
        unsafe { *(self.data.as_mut_ptr().add(byte_offset) as *mut u16) = val };
    }

    pub unsafe fn write_f32(&mut self, texel_offset: usize, val: f32) {
        let byte_offset = texel_offset * 4;
        unsafe { *(self.data.as_mut_ptr().add(byte_offset) as *mut f32) = val };
    }

    pub fn empty(texel_count: usize, format: PixelFormat) -> Self {
        let byte_count = texel_count * format.texel_bytes();
        let mut data = gvec_with_capacity(byte_count);
        data.resize(byte_count, 0u8);
        Self { data, format }
    }

    pub fn from_u8_slice(slice: &[u8]) -> Self {
        let mut data = gvec_with_capacity(slice.len());
        data.extend_from_slice(slice);
        Self {
            data,
            format: PixelFormat::U8,
        }
    }

    pub fn from_f32_slice(slice: &[f32]) -> Self {
        let byte_len = slice.len() * 4;
        let mut data = gvec_with_capacity(byte_len);
        let bytes = unsafe { core::slice::from_raw_parts(slice.as_ptr() as *const u8, byte_len) };
        data.extend_from_slice(bytes);
        Self {
            data,
            format: PixelFormat::F32,
        }
    }

    pub fn from_f16_slice(slice: &[u16]) -> Self {
        let byte_len = slice.len() * 2;
        let mut data = gvec_with_capacity(byte_len);
        let bytes = unsafe { core::slice::from_raw_parts(slice.as_ptr() as *const u8, byte_len) };
        data.extend_from_slice(bytes);
        Self {
            data,
            format: PixelFormat::F16,
        }
    }

    pub fn as_u8(&self) -> &[u8] {
        &self.data
    }

    pub fn as_f16(&mut self) -> &[u16] {
        assert_eq!(self.format, PixelFormat::F16);
        unsafe {
            core::slice::from_raw_parts(self.data.as_ptr() as *const u16, self.data.len() / 2)
        }
    }

    pub fn as_u8_mut(&mut self) -> &mut [u8] {
        &mut self.data
    }

    pub fn as_f16_mut(&mut self) -> &mut [u16] {
        assert_eq!(self.format, PixelFormat::F16);
        unsafe {
            core::slice::from_raw_parts_mut(self.data.as_mut_ptr() as *mut u16, self.data.len() / 2)
        }
    }

    pub fn as_f32_slice(&self) -> &[f32] {
        assert_eq!(self.format, PixelFormat::F32);
        unsafe {
            core::slice::from_raw_parts(self.data.as_ptr() as *const f32, self.data.len() / 4)
        }
    }

    pub fn as_f32_slice_mut(&mut self) -> &mut [f32] {
        assert_eq!(self.format, PixelFormat::F32);
        unsafe {
            core::slice::from_raw_parts_mut(self.data.as_mut_ptr() as *mut f32, self.data.len() / 4)
        }
    }
}

#[derive(Clone, Debug)]
pub struct ImageBase {
    pub format: PixelFormat,
    pub encoding: ColorEncoding,
    pub resolution: Point2i,
    pub n_channels: i32,
}

#[repr(C)]
#[derive(Clone, Debug)]
pub struct Image {
    pub format: PixelFormat,
    pub encoding: ColorEncoding,
    pub resolution: Point2i,
    pub n_channels: i32,
    pub pixels: Pixels,
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct ImageView {
    pub pixels: *const u8,
    pub byte_len: usize,
    pub resolution: Point2i,
    pub n_channels: i32,
    pub format: PixelFormat,
    pub encoding: ColorEncoding,
}

impl Image {
    pub fn new(
        format: PixelFormat,
        resolution: Point2i,
        n_channels: i32,
        encoding: ColorEncoding,
    ) -> Self {
        let texel_count = (resolution.x() * resolution.y()) as usize * n_channels as usize;
        Self {
            format,
            encoding,
            resolution,
            n_channels,
            pixels: Pixels::empty(texel_count, format),
        }
    }

    pub fn from_u8(
        data: &[u8],
        resolution: Point2i,
        n_channels: i32,
        encoding: ColorEncoding,
    ) -> Self {
        let expected = (resolution.x() * resolution.y()) as usize * n_channels as usize;
        assert_eq!(data.len(), expected, "Pixel data size mismatch");
        Self {
            format: PixelFormat::U8,
            encoding,
            resolution,
            n_channels,
            pixels: Pixels::from_u8_slice(data),
        }
    }

    pub fn from_f32(data: &[f32], resolution: Point2i, n_channels: i32) -> Self {
        let expected = (resolution.x() * resolution.y()) as usize * n_channels as usize;
        assert_eq!(data.len(), expected, "Pixel data size mismatch");
        Self {
            format: PixelFormat::F32,
            encoding: LINEAR,
            resolution,
            n_channels,
            pixels: Pixels::from_f32_slice(data),
        }
    }

    pub fn from_f16(data: &[u16], resolution: Point2i, n_channels: i32) -> Self {
        let expected = (resolution.x() * resolution.y()) as usize * n_channels as usize;
        assert_eq!(data.len(), expected, "Pixel data size mismatch");
        Self {
            format: PixelFormat::F16,
            encoding: LINEAR,
            resolution,
            n_channels,
            pixels: Pixels::from_f16_slice(data),
        }
    }

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

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

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

    pub fn is_valid(&self) -> bool {
        self.resolution.x() > 0 && self.resolution.y() > 0
    }

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

    pub fn remap_pixel_coords(&self, p: &mut Point2i, wrap_mode: WrapMode2D) -> bool {
        for i in 0..2 {
            if p[i] >= 0 && p[i] < self.resolution[i] {
                continue;
            }
            match wrap_mode.uv[i] {
                WrapMode::Black => return false,
                WrapMode::Clamp => p[i] = p[i].clamp(0, self.resolution[i] - 1),
                WrapMode::Repeat => p[i] = p[i].rem_euclid(self.resolution[i]),
                WrapMode::OctahedralSphere => {
                    p[i] = p[i].clamp(0, self.resolution[i] - 1);
                }
            }
        }
        true
    }

    pub fn base(&self) -> ImageBase {
        ImageBase {
            format: self.format,
            encoding: self.encoding,
            resolution: self.resolution,
            n_channels: self.n_channels,
        }
    }

    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.remap_pixel_coords(&mut p, wrap_mode) {
            return 0.0;
        }
        let offset = self.pixel_offset(p) + c as usize;
        unsafe { self.pixels.read(offset, &self.encoding) }
    }

    pub fn lookup_nearest_channel(&self, p: Point2f, c: i32) -> Float {
        self.lookup_nearest_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn lookup_nearest_channel_with_wrap(
        &self,
        p: Point2f,
        c: i32,
        wrap_mode: WrapMode2D,
    ) -> Float {
        let pi = Point2i::new(
            p.x() as i32 * self.resolution().x(),
            p.y() as i32 * self.resolution().y(),
        );

        self.get_channel_with_wrap(pi, c, wrap_mode)
    }

    pub fn get_channels<const N: usize>(&self, p: Point2i) -> [Float; N] {
        self.get_channels_with_wrap(p, WrapMode::Clamp.into())
    }

    pub fn get_channels_with_wrap<const N: usize>(
        &self,
        mut p: Point2i,
        wrap_mode: WrapMode2D,
    ) -> [Float; N] {
        debug_assert!(N <= self.n_channels as usize);
        let mut result = [0.0; N];
        if !self.remap_pixel_coords(&mut p, wrap_mode) {
            return result;
        }
        let offset = self.pixel_offset(p);
        for i in 0..N {
            result[i] = unsafe { self.pixels.read(offset + i, &self.encoding) };
        }
        result
    }

    pub fn get_channels_average(&self, p: Point2i) -> Float {
        let offset = self.pixel_offset(p);
        let nc = self.n_channels as usize;
        let mut sum = 0.0;
        for i in 0..nc {
            sum += unsafe { self.pixels.read(offset + i, &self.encoding) };
        }
        sum / nc as Float
    }

    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;
        unsafe {
            match self.format {
                PixelFormat::U8 => {
                    self.pixels
                        .write_u8(offset, self.encoding.from_linear_scalar(value));
                }
                PixelFormat::F16 => {
                    self.pixels
                        .write_f16(offset, half::f16::from_f32(value).to_bits());
                }
                PixelFormat::F32 => {
                    self.pixels.write_f32(offset, value);
                }
            }
        }
    }

    pub fn bilerp_channel(&self, p: Point2f, c: i32) -> Float {
        self.bilerp_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn bilerp_channel_with_wrap(&self, p: Point2f, c: i32, wrap_mode: WrapMode2D) -> Float {
        let x = p.x() * self.resolution.x() as Float - 0.5;
        let y = p.y() * self.resolution.y() as Float - 0.5;
        let xi = x.floor() as i32;
        let yi = y.floor() as i32;
        let dx = x - xi as Float;
        let dy = y - yi as Float;
        let v00 = self.get_channel_with_wrap(Point2i::new(xi, yi), c, wrap_mode);
        let v10 = self.get_channel_with_wrap(Point2i::new(xi + 1, yi), c, wrap_mode);
        let v01 = self.get_channel_with_wrap(Point2i::new(xi, yi + 1), c, wrap_mode);
        let v11 = self.get_channel_with_wrap(Point2i::new(xi + 1, yi + 1), c, wrap_mode);
        lerp(dy, lerp(dx, v00, v10), lerp(dx, v01, v11))
    }

    pub fn has_any_infinite_pixels(&self) -> bool {
        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;
                    }
                }
            }
        }
        false
    }

    pub fn has_any_nan_pixels(&self) -> bool {
        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;
                    }
                }
            }
        }
        false
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FilterFunction {
    Point,
    Bilinear,
    Trilinear,
    Ewa,
}

impl FilterFunction {
    pub fn parse(name: &str) -> Result<FilterFunction> {
        match name {
            "ewa" | "EWA" => Ok(FilterFunction::Ewa),
            "trilinear" => Ok(FilterFunction::Trilinear),
            "bilinear" => Ok(FilterFunction::Bilinear),
            "point" => Ok(FilterFunction::Point),
            _ => bail!("Filter function unknown"),
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct ImagePyramid {
    pub levels: *const Ptr<Image>,
    pub level_count: u32,
    pub wrap_mode: WrapMode,
    pub filter: FilterFunction,
    pub max_aniso: f32,
}