636 lines
19 KiB
Rust
636 lines
19 KiB
Rust
use shared::core::geometry::Point2i;
|
|
use shared::core::image::{Image, PixelFormat};
|
|
use shared::utils::math::f16_to_f32;
|
|
use std::ops::Deref;
|
|
|
|
pub mod io;
|
|
pub mod metadata;
|
|
pub mod ops;
|
|
pub mod pixel;
|
|
|
|
pub use metadata::*;
|
|
|
|
impl std::fmt::Display for PixelFormat {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
PixelFormat::U8 => write!(f, "U256"),
|
|
PixelFormat::F16 => write!(f, "Half"),
|
|
PixelFormat::F32 => write!(f, "Float"),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[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 enum PixelStorage {
|
|
U8(Vec<u8>),
|
|
F16(Vec<f16>),
|
|
F32(Vec<f32>),
|
|
}
|
|
|
|
pub struct ImageBuffer {
|
|
pub view: Image,
|
|
pub channel_names: Vec<String>,
|
|
_storage: PixelStorage,
|
|
}
|
|
|
|
impl Deref for ImageBuffer {
|
|
type Target = Image;
|
|
fn deref(&self) -> &Self::Target {
|
|
&self.view
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct ImageAndMetadata {
|
|
pub image: ImageBuffer,
|
|
pub metadata: ImageMetadata,
|
|
}
|
|
|
|
impl ImageBuffer {
|
|
fn resolution(&self) {
|
|
self.view.resolution()
|
|
}
|
|
|
|
pub fn new_empty() -> Self {
|
|
Self {
|
|
channel_names: Vec::new(),
|
|
_storage: PixelStorage::U8(Vec::new()),
|
|
view: Image {
|
|
format: PixelFormat::U256,
|
|
resolution: Point2i::new(0, 0),
|
|
n_channels: 0,
|
|
pixels: Pixels::U8(std::ptr::null()),
|
|
encoding: ColorEncoding::default(),
|
|
},
|
|
}
|
|
}
|
|
|
|
fn from_vector(
|
|
format: PixelFormat,
|
|
resolution: Point2i,
|
|
channel_names: Vec<String>,
|
|
encoding: ColorEncoding,
|
|
) -> Self {
|
|
let n_channels = channel_names.len() as i32;
|
|
let (format, pixels_view) = match &storage {
|
|
PixelStorage::U8(vec) => (PixelFormat::U8, ImagePixels::U8(vec.as_ptr())),
|
|
PixelStorage::F16(vec) => (PixelFormat::F16, ImagePixels::F16(vec.as_ptr())),
|
|
PixelStorage::F32(vec) => (PixelFormat::F32, ImagePixels::F32(vec.as_ptr())),
|
|
};
|
|
|
|
let view = Image {
|
|
format,
|
|
resolution,
|
|
n_channels,
|
|
encoding,
|
|
pixels: pixels_view,
|
|
};
|
|
|
|
Self {
|
|
view,
|
|
_storage: storage,
|
|
channel_names,
|
|
}
|
|
}
|
|
|
|
pub fn from_u8(
|
|
data: Vec<u8>,
|
|
resolution: Point2i,
|
|
channel_names: Vec<String>,
|
|
encoding: ColorEncoding,
|
|
) -> Self {
|
|
let n_channels = channel_names.len() as i32;
|
|
let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;
|
|
if data.len() != expected_len {
|
|
panic!(
|
|
"ImageBuffer::from_u8: Data length {} does not match resolution {:?} * channels {}",
|
|
data.len(),
|
|
resolution,
|
|
n_channels
|
|
);
|
|
}
|
|
|
|
let storage = PixelStorage::U8(data);
|
|
|
|
let ptr = match &storage {
|
|
PixelStorage::U8(v) => v.as_ptr(),
|
|
_ => unreachable!(),
|
|
};
|
|
|
|
let view = Image {
|
|
format: PixelFormat::U256,
|
|
resolution,
|
|
n_channels,
|
|
pixels: Pixels::U8(ptr),
|
|
encoding: encoding.clone(),
|
|
};
|
|
|
|
Self {
|
|
view,
|
|
channel_names,
|
|
_storage: storage,
|
|
}
|
|
}
|
|
|
|
pub fn from_u8(
|
|
data: Vec<u8>,
|
|
resolution: Point2i,
|
|
channel_names: Vec<String>,
|
|
encoding: ColorEncoding,
|
|
) -> Self {
|
|
let n_channels = channel_names.len() as i32;
|
|
let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;
|
|
|
|
if data.len() != expected_len {
|
|
panic!(
|
|
"ImageBuffer::from_u8: Data length {} does not match resolution {:?} * channels {}",
|
|
data.len(),
|
|
resolution,
|
|
n_channels
|
|
);
|
|
}
|
|
|
|
let storage = PixelStorage::U8(data);
|
|
|
|
let ptr = match &storage {
|
|
PixelStorage::U8(v) => v.as_ptr(),
|
|
_ => unreachable!(),
|
|
};
|
|
|
|
let view = Image {
|
|
format: PixelFormat::U256,
|
|
resolution,
|
|
n_channels,
|
|
pixels: Pixels::U8(ptr),
|
|
encoding: encoding.clone(),
|
|
};
|
|
|
|
Self {
|
|
view,
|
|
channel_names,
|
|
_storage: storage,
|
|
}
|
|
}
|
|
|
|
pub fn from_f16(data: Vec<half::f16>, resolution: Point2i, channel_names: Vec<String>) -> Self {
|
|
let n_channels = channel_names.len() as i32;
|
|
let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;
|
|
|
|
if data.len() != expected_len {
|
|
panic!("ImageBuffer::from_f16: Data length mismatch");
|
|
}
|
|
|
|
let storage = PixelStorage::F16(data);
|
|
|
|
let ptr = match &storage {
|
|
PixelStorage::F16(v) => v.as_ptr() as *const u16,
|
|
_ => unreachable!(),
|
|
};
|
|
|
|
let view = Image {
|
|
format: PixelFormat::Half,
|
|
resolution,
|
|
n_channels,
|
|
pixels: Pixels::F16(ptr),
|
|
encoding: ColorEncoding::default(),
|
|
};
|
|
|
|
Self {
|
|
view,
|
|
channel_names,
|
|
_storage: storage,
|
|
}
|
|
}
|
|
|
|
pub fn from_f32(data: Vec<f32>, resolution: Point2i, channel_names: Vec<String>) -> Self {
|
|
let n_channels = channel_names.len() as i32;
|
|
let expected_len = (resolution.x * resolution.y) as usize * n_channels as usize;
|
|
|
|
if data.len() != expected_len {
|
|
panic!("ImageBuffer::from_f32: Data length mismatch");
|
|
}
|
|
|
|
let storage = PixelStorage::F32(data);
|
|
|
|
let ptr = match &storage {
|
|
PixelStorage::F32(v) => v.as_ptr(),
|
|
_ => unreachable!(),
|
|
};
|
|
|
|
let view = Image {
|
|
format: PixelFormat::Float,
|
|
resolution,
|
|
n_channels,
|
|
pixels: Pixels::F32(ptr),
|
|
encoding: ColorEncoding::default(),
|
|
};
|
|
|
|
Self {
|
|
view,
|
|
channel_names,
|
|
_storage: storage,
|
|
}
|
|
}
|
|
|
|
pub fn new(
|
|
format: PixelFormat,
|
|
resolution: Point2i,
|
|
channel_names: &[&str],
|
|
encoding: *const ColorEncoding,
|
|
) -> Self {
|
|
let n_channels = channel_names.len();
|
|
let pixel_count = (resolution.x * resolution.y) as usize * n_channels;
|
|
let owned_names: Vec<String> = channel_names.iter().map(|s| s.to_string()).collect();
|
|
|
|
let storage = match format {
|
|
PixelFormat::U8 => PixelStorage::U8(vec![0; pixel_count]),
|
|
PixelFormat::F16 => PixelStorage::F16(vec![0; pixel_count]),
|
|
PixelFormat::F32 => PixelStorage::F32(vec![0.0; pixel_count]),
|
|
};
|
|
|
|
Self::from_vector(storage, resolution, owned_names, encoding)
|
|
}
|
|
|
|
pub fn channel_names(&self) -> Vec<&str> {
|
|
self.channel_names.iter().map(|s| s.as_str()).collect()
|
|
}
|
|
|
|
pub fn encoding(&self) -> ColorEncoding {
|
|
self.view.encoding
|
|
}
|
|
|
|
pub fn set_channel(&self, p: Point2i, c: i32, mut value: Float) {
|
|
if value.is_nan() {
|
|
value = 0.0;
|
|
}
|
|
|
|
if !self.view.resolution.inside_exclusive(p) {
|
|
return;
|
|
}
|
|
|
|
let offset = self.view.pixel_offset(p) + c as usize;
|
|
match &mut self._storage {
|
|
PixelStorage::U8(data) => {
|
|
if !self.view.encoding.is_null() {
|
|
data[offset] = self.view.encoding.from_linear_scalar(value);
|
|
} else {
|
|
let val = (value * 255.0 + 0.5).clamp(0.0, 255.0);
|
|
data[offset] = val as u8;
|
|
}
|
|
}
|
|
PixelStorage::F16(data) => {
|
|
data[offset] = f16_to_f32(value);
|
|
}
|
|
PixelStorage::F32(data) => {
|
|
data[offset] = value;
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn get_channels_with_wrap(&self, p: Point2i, wrap_mode: WrapMode2D) -> ImageChannelValues {
|
|
let mut pp = p;
|
|
if !self.view.remap_pixel_coords(&mut pp, wrap_mode) {
|
|
return ImageChannelValues(smallvec![0.0; desc.offset.len()]);
|
|
}
|
|
|
|
let pixel_offset = self.view.pixel_offset(pp);
|
|
let mut values: SmallVec<[Float; 4]> = SmallVec::with_capacity(desc.offset.len());
|
|
match &self.pixels {
|
|
PixelData::U8(data) => {
|
|
for i in 0..self.view.n_channels() {
|
|
let raw_val = data[pixel_offset + i];
|
|
let linear = self.view.encoding.to_linear_scalar(raw_val);
|
|
values.push(linear);
|
|
}
|
|
}
|
|
PixelData::F16(data) => {
|
|
for i in 0..self.view.n_channels() {
|
|
let raw_val = data[pixel_offset];
|
|
values.push(f16_to_f32(raw_val));
|
|
}
|
|
}
|
|
PixelData::F32(data) => {
|
|
for i in 0..self.view.n_channels() {
|
|
let val = data[pixel_offset + i];
|
|
values.push(val);
|
|
}
|
|
}
|
|
}
|
|
|
|
ImageChannelValues(values)
|
|
}
|
|
|
|
pub fn get_channels(&self, p: Point2i) -> ImageChannelValues {
|
|
self.get_channels_with_wrap(p, WrapMode::Clamp.into())
|
|
}
|
|
|
|
pub fn get_channels_with_desc(
|
|
&self,
|
|
p: Point2i,
|
|
desc: &ImageChannelDesc,
|
|
wrap: WrapMode2D,
|
|
) -> ImageChannelValues {
|
|
let mut pp = p;
|
|
if !self.view.remap_pixel_coords(&mut pp, wrap) {
|
|
return ImageChannelValues(smallvec![0.0; desc.offset.len()]);
|
|
}
|
|
|
|
let pixel_offset = self.view.pixel_offset(pp);
|
|
|
|
let mut values: SmallVec<[Float; 4]> = SmallVec::with_capacity(desc.offset.len());
|
|
|
|
match &self.pixels {
|
|
PixelData::U8(data) => {
|
|
for &channel_idx in &desc.offset {
|
|
let raw_val = data[pixel_offset + channel_idx as usize];
|
|
let linear = self.view.encoding.to_linear_scalar(raw_val);
|
|
values.push(linear);
|
|
}
|
|
}
|
|
PixelData::F16(data) => {
|
|
for &channel_idx in &desc.offset {
|
|
let raw_val = data[pixel_offset + channel_idx as usize];
|
|
values.push(f16_to_f32(raw_val));
|
|
}
|
|
}
|
|
PixelData::F32(data) => {
|
|
for &channel_idx in &desc.offset {
|
|
let val = data[pixel_offset + channel_idx as usize];
|
|
values.push(val);
|
|
}
|
|
}
|
|
}
|
|
|
|
ImageChannelValues(values)
|
|
}
|
|
|
|
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_channel_desc(
|
|
&self,
|
|
requested_channels: &[&str],
|
|
) -> Result<ImageChannelDesc, String> {
|
|
let mut offset = Vec::with_capacity(requested_channels.len());
|
|
|
|
for &req in requested_channels.iter() {
|
|
match self.channel_names.iter().position(|n| n == req) {
|
|
Some(idx) => {
|
|
offset.push(idx);
|
|
}
|
|
None => {
|
|
return Err(format!(
|
|
"Image is missing requested channel '{}'. Available channels: {:?}",
|
|
req, self.channel_names
|
|
));
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(ImageChannelDesc { offset })
|
|
}
|
|
|
|
pub fn all_channels_desc(&self) -> ImageChannelDesc {
|
|
ImageChannelDesc {
|
|
offset: (0..self.n_channels()).collect(),
|
|
}
|
|
}
|
|
|
|
pub fn select_channels(&self, desc: &ImageChannelDesc) -> Self {
|
|
let desc_channel_names: Vec<String> = desc
|
|
.offset
|
|
.iter()
|
|
.map(|&i| self.channel_names[i as usize])
|
|
.collect();
|
|
|
|
let new_storage = match &self._storage {
|
|
PixelStorage::U8(src_data) => {
|
|
// Allocate destination buffer
|
|
let mut dst_data = vec![0u8; pixel_count * dst_n_channels];
|
|
|
|
// Iterate over every pixel (Flat loop is faster than nested x,y)
|
|
for i in 0..pixel_count {
|
|
let src_pixel_start = i * src_n_channels;
|
|
let dst_pixel_start = i * dst_n_channels;
|
|
|
|
// Copy specific channels based on desc
|
|
for (out_idx, &in_channel_offset) in desc.offset.iter().enumerate() {
|
|
let val = src_data[src_pixel_start + in_channel_offset as usize];
|
|
dst_data[dst_pixel_start + out_idx] = val;
|
|
}
|
|
}
|
|
PixelStorage::U8(dst_data)
|
|
}
|
|
PixelStorage::F16(src_data) => {
|
|
let mut dst_data = vec![half::f16::ZERO; pixel_count * dst_n_channels];
|
|
|
|
for i in 0..pixel_count {
|
|
let src_pixel_start = i * src_n_channels;
|
|
let dst_pixel_start = i * dst_n_channels;
|
|
|
|
for (out_idx, &in_channel_offset) in desc.offset.iter().enumerate() {
|
|
let val = src_data[src_pixel_start + in_channel_offset as usize];
|
|
dst_data[dst_pixel_start + out_idx] = val;
|
|
}
|
|
}
|
|
PixelStorage::F16(dst_data)
|
|
}
|
|
PixelStorage::F32(src_data) => {
|
|
let mut dst_data = vec![0.0; pixel_count * dst_n_channels];
|
|
|
|
for i in 0..pixel_count {
|
|
let src_pixel_start = i * src_n_channels;
|
|
let dst_pixel_start = i * dst_n_channels;
|
|
|
|
for (out_idx, &in_channel_offset) in desc.offset.iter().enumerate() {
|
|
let val = src_data[src_pixel_start + in_channel_offset as usize];
|
|
dst_data[dst_pixel_start + out_idx] = val;
|
|
}
|
|
}
|
|
PixelStorage::F32(dst_data)
|
|
}
|
|
};
|
|
|
|
let image = Self::new(
|
|
self.format,
|
|
self.resolution,
|
|
desc_channel_names,
|
|
self.encoding(),
|
|
);
|
|
}
|
|
|
|
pub fn set_channels(
|
|
&mut self,
|
|
p: Point2i,
|
|
desc: &ImageChannelDesc,
|
|
mut values: &ImageChannelValues,
|
|
) {
|
|
assert_eq!(desc.size(), values.len());
|
|
for i in 0..desc.size() {
|
|
self.view.set_channel(p, desc.offset[i], values[i]);
|
|
}
|
|
}
|
|
|
|
pub fn set_channels_all(&mut self, p: Point2i, values: &ImageChannelValues) {
|
|
self.set_channels(p, &self.all_channels_desc(), values)
|
|
}
|
|
|
|
pub fn get_sampling_distribution<F>(&self, dxd_a: F, domain: Bounds2f) -> Array2D<Float>
|
|
where
|
|
F: Fn(Point2f) -> Float + Sync + Send,
|
|
{
|
|
let width = self.resolution().x();
|
|
let height = self.resolution().y();
|
|
|
|
let mut dist = Array2D::new_with_dims(width as usize, height as usize);
|
|
|
|
dist.values
|
|
.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_default(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 default_domain = Bounds2f::from_points(Point2f::new(0.0, 0.0), Point2f::new(1.0, 1.0));
|
|
|
|
self.get_sampling_distribution(|_| 1.0, default_domain)
|
|
}
|
|
|
|
pub fn mse(
|
|
&self,
|
|
desc: ImageChannelDesc,
|
|
ref_img: &Image,
|
|
generate_mse_image: bool,
|
|
) -> (ImageChannelValues, Option<Image>) {
|
|
let mut sum_se: Vec<f64> = vec![0.; desc.size()];
|
|
let names_ref = self.channel_names_from_desc(&desc);
|
|
let ref_desc = ref_img
|
|
.get_channel_desc(&self.channel_names_from_desc(&desc))
|
|
.expect("Channels not found in image");
|
|
assert_eq!(self.resolution(), ref_img.resolution());
|
|
|
|
let width = self.resolution.x() as usize;
|
|
let height = self.resolution.y() as usize;
|
|
let n_channels = desc.offset.len();
|
|
let mut mse_pixels = if generate_mse_image {
|
|
vec![0.0f32; width * height * n_channels]
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
for y in 0..self.resolution().y() {
|
|
for x in 0..self.resolution().x() {
|
|
let v =
|
|
self.get_channel_with_desc(Point2i::new(x, y), &desc, WrapMode::Clamp.into());
|
|
let v_ref = self.get_channel_with_desc(
|
|
Point2i::new(x, y),
|
|
&ref_desc,
|
|
WrapMode::Clamp.into(),
|
|
);
|
|
for c in 0..desc.size() {
|
|
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 = (self.resolution().x() * self.resolution.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 {
|
|
Some(Image::new(
|
|
PixelFormat::F32,
|
|
self.resolution,
|
|
&names_ref,
|
|
LINEAR,
|
|
))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
(ImageChannelValues(mse_values), mse_image)
|
|
}
|
|
|
|
pub fn update_view_pointers(&mut self) {
|
|
self.view.pixels = match &self._storage {
|
|
PixelStorage::U8(vec) => Pixels::U8(vec.as_ptr()),
|
|
PixelStorage::F16(vec) => Pixels::F16(vec.as_ptr() as *const u16),
|
|
PixelStorage::F32(vec) => Pixels::F32(vec.as_ptr()),
|
|
};
|
|
}
|
|
}
|