pbrt/shared/src/core/image.rs

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

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

#[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, Copy, Debug)]
pub struct Pixels {
    ptr: Ptr<u8>,
    format: PixelFormat,
}

impl Pixels {
    pub fn new_u8(ptr: Ptr<u8>) -> Self {
        Self {
            ptr,
            format: PixelFormat::U8,
        }
    }

    pub fn new_f16(ptr: Ptr<u16>) -> Self {
        Self {
            ptr: Ptr::from_raw(ptr.as_raw() as *const u8),
            format: PixelFormat::F16,
        }
    }

    pub fn new_f32(ptr: Ptr<f32>) -> Self {
        Self {
            ptr: Ptr::from_raw(ptr.as_raw() as *const u8),
            format: PixelFormat::F32,
        }
    }

    unsafe fn read_u8(&self, byte_offset: usize) -> u8 {
        unsafe { *self.ptr.as_raw().add(byte_offset) }
    }

    unsafe fn read_f16(&self, elem_offset: usize) -> u16 {
        let byte_offset = elem_offset * 2;
        unsafe { *(self.ptr.as_raw().add(byte_offset) as *const u16) }
    }

    unsafe fn read_f32(&self, elem_offset: usize) -> f32 {
        let byte_offset = elem_offset * 4;
        unsafe { *(self.ptr.as_raw().add(byte_offset) as *const f32) }
    }

    // Unsure if ill need this
    // pub unsafe fn read(&self, offset: usize) -> Float {
    //     match self.format {
    //         PixelFormat::U8 => unsafe { self.read_u8(offset) as Float / 255.0 },
    //         PixelFormat::F16 => unsafe { f16_to_f32_software(self.read_f16(offset)) },
    //         PixelFormat::F32 => unsafe { self.read_f32(offset) },
    //     }
    // }
}

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

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

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct DeviceImage {
    pub base: ImageBase,
    pub pixels: Pixels,
}

impl DeviceImage {
    pub fn base(&self) -> ImageBase {
        self.base
    }

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

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

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

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

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

    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 trait ImageAccess {
    fn get_channel_with_wrap(&self, p: Point2i, c: i32, wrap_mode: WrapMode2D) -> Float;
    fn get_channel(&self, p: Point2i, c: i32) -> Float;
    fn lookup_nearest_channel_with_wrap(&self, p: Point2f, c: i32, wrap_mode: WrapMode2D) -> Float;
    fn lookup_nearest_channel(&self, p: Point2f, c: i32) -> Float;
}

impl ImageAccess for DeviceImage {
    fn get_channel_with_wrap(&self, mut p: Point2i, c: i32, wrap_mode: WrapMode2D) -> Float {
        if !self.base.remap_pixel_coords(&mut p, wrap_mode) {
            return 0.;
        }

        let offset = (self.pixel_offset(p) + c as u32) as usize;
        unsafe {
            match self.pixels.format {
                PixelFormat::U8 => {
                    let raw_val = self.pixels.read_u8(offset);
                    self.base.encoding.to_linear_scalar(raw_val)
                }
                PixelFormat::F16 => {
                    let raw_val = self.pixels.read_f16(offset);
                    f16_to_f32_software(raw_val)
                }
                PixelFormat::F32 => self.pixels.read_f32(offset),
            }
        }
    }

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

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

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