pbrt/shared/src/core/image.rs

use crate::core::color::{ColorEncoding, ColorEncodingTrait, LINEAR};
use crate::core::geometry::{Bounds2f, Point2f, Point2fi, Point2i};
use crate::core::pbrt::Float;
use crate::utils::containers::Array2D;
use crate::utils::math::{f16_to_f32, lerp, square};
use core::hash;
use half::f16;
use smallvec::{SmallVec, smallvec};
use std::ops::{Deref, DerefMut};

#[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 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 enum Pixels {
    U8(*const u8),
    F16(*const u16),
    F32(*const f32),
}

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

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

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

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

    pub fn n_channels(&self) -> i32 {
        self.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 get_channel_with_wrap(&self, p: Point2i, c: i32, wrap_mode: WrapMode2D) -> Float {
        if !self.remap_pixel_coords(&mut p, wrap_mode) {
            return 0.;
        }

        let offset = self.pixel_offset(p) + c;
        unsafe {
            match self.pixels {
                Pixels::U8(ptr) => {
                    let raw_u8 = *ptr.add(offset);
                    self.encoding.to_linear_scalar(raw_u8)
                }
                Pixels::F16(ptr) => {
                    let half_bits = *ptr.add(offset);
                    f16_to_f32(f16::from_bits(half_bits))
                }
                Pixels::F32(ptr) => *ptr.add(offset),
            }
        }
    }

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

    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 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 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 lookup_nearest_channel(&self, p: Point2f, c: i32) -> Float {
        self.lookup_nearest_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }
}