432 lines
13 KiB
Rust
432 lines
13 KiB
Rust
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::Float;
|
|
use crate::GVec;
|
|
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, Debug)]
|
|
pub struct Pixels {
|
|
pixels: 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 {
|
|
*self.data.as_ptr().add(texel_offset)
|
|
}
|
|
|
|
pub unsafe fn read_f16(&self, texel_offset: usize) -> u16 {
|
|
let byte_offset = texel_offset * 2;
|
|
*(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;
|
|
*(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) {
|
|
*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;
|
|
*(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;
|
|
*(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_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 Image {
|
|
pub format: PixelFormat,
|
|
pub encoding: ColorEncoding,
|
|
pub resolution: Point2i,
|
|
pub n_channels: i32,
|
|
pub pixels: Pixels,
|
|
}
|
|
|
|
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 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_array<const N: usize>(&self, p: Point2i) -> [Float; N] {
|
|
self.get_channels_array_with_wrap(p, WrapMode::Clamp.into())
|
|
}
|
|
|
|
pub fn get_channels_array_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
|
|
}
|
|
}
|