418 lines
14 KiB
Rust
418 lines
14 KiB
Rust
// use rayon::prelude::*;
|
|
use crate::core::pbrt::Float;
|
|
use crate::geometry::{Bounds2i, Point2i};
|
|
use crate::image::pixel::PixelStorage;
|
|
use crate::image::{Image, PixelData, PixelFormat, WrapMode, WrapMode2D};
|
|
use crate::utils::math::windowed_sinc;
|
|
use rayon::prelude::*;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct ResampleWeight {
|
|
pub first_pixel: i32,
|
|
pub weight: [Float; 4],
|
|
}
|
|
|
|
impl Image {
|
|
pub fn flip_y(&mut self) {
|
|
let res = self.resolution;
|
|
let nc = self.n_channels();
|
|
|
|
match &mut self.pixels {
|
|
PixelData::U8(d) => flip_y_kernel(d, res, nc),
|
|
PixelData::F16(d) => flip_y_kernel(d, res, nc),
|
|
PixelData::F32(d) => flip_y_kernel(d, res, nc),
|
|
}
|
|
}
|
|
|
|
pub fn crop(&self, bounds: Bounds2i) -> Image {
|
|
let new_res = Point2i::new(
|
|
bounds.p_max.x() - bounds.p_min.x(),
|
|
bounds.p_max.y() - bounds.p_min.y(),
|
|
);
|
|
|
|
let mut new_image = Image::new(
|
|
self.format,
|
|
new_res,
|
|
self.channel_names.clone(),
|
|
self.encoding,
|
|
);
|
|
|
|
match (&self.pixels, &mut new_image.pixels) {
|
|
(PixelData::U8(src), PixelData::U8(dst)) => {
|
|
crop_kernel(src, dst, self.resolution, bounds, self.n_channels())
|
|
}
|
|
(PixelData::F16(src), PixelData::F16(dst)) => {
|
|
crop_kernel(src, dst, self.resolution, bounds, self.n_channels())
|
|
}
|
|
(PixelData::F32(src), PixelData::F32(dst)) => {
|
|
crop_kernel(src, dst, self.resolution, bounds, self.n_channels())
|
|
}
|
|
_ => panic!("Format mismatch in crop"),
|
|
}
|
|
|
|
new_image
|
|
}
|
|
|
|
pub fn copy_rect_out(&self, extent: Bounds2i, buf: &mut [Float], wrap: WrapMode2D) {
|
|
match &self.pixels {
|
|
PixelData::U8(d) => copy_rect_out_kernel(d, self, extent, buf, wrap),
|
|
PixelData::F16(d) => copy_rect_out_kernel(d, self, extent, buf, wrap),
|
|
PixelData::F32(d) => copy_rect_out_kernel(d, self, extent, buf, wrap),
|
|
}
|
|
}
|
|
|
|
pub fn copy_rect_in(&mut self, extent: Bounds2i, buf: &[Float]) {
|
|
let resolution = self.resolution;
|
|
let n_channels = self.n_channels();
|
|
let encoding = self.encoding;
|
|
|
|
match &mut self.pixels {
|
|
PixelData::U8(d) => {
|
|
copy_rect_in_kernel(d, resolution, n_channels, encoding, extent, buf)
|
|
}
|
|
PixelData::F16(d) => {
|
|
copy_rect_in_kernel(d, resolution, n_channels, encoding, extent, buf)
|
|
}
|
|
PixelData::F32(d) => {
|
|
copy_rect_in_kernel(d, resolution, n_channels, encoding, extent, buf)
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn float_resize_up(&self, new_res: Point2i, wrap_mode: WrapMode2D) -> Image {
|
|
assert!(new_res.x() >= self.resolution.x() && new_res.y() >= self.resolution.y());
|
|
assert!(
|
|
matches!(self.format, PixelFormat::F32),
|
|
"ResizeUp requires Float format"
|
|
);
|
|
|
|
// Create destination image wrapped in Mutex for tile-based write access
|
|
let resampled_image = Arc::new(Mutex::new(Image::new(
|
|
PixelFormat::F32, // Force float output
|
|
new_res,
|
|
self.channel_names.clone(),
|
|
self.encoding,
|
|
)));
|
|
|
|
// Precompute weights
|
|
let x_weights = resample_weights(self.resolution.x() as usize, new_res.x() as usize);
|
|
let y_weights = resample_weights(self.resolution.y() as usize, new_res.y() as usize);
|
|
let n_channels = self.n_channels();
|
|
|
|
// Generate Tiles
|
|
let tile_size = 16;
|
|
let tiles = generate_tiles(new_res, tile_size);
|
|
|
|
// Parallel Execution
|
|
tiles.par_iter().for_each(|out_extent| {
|
|
let x_start = x_weights[out_extent.p_min.x() as usize].first_pixel;
|
|
let x_end = x_weights[(out_extent.p_max.x() - 1) as usize].first_pixel + 4;
|
|
let y_start = y_weights[out_extent.p_min.y() as usize].first_pixel;
|
|
let y_end = y_weights[(out_extent.p_max.y() - 1) as usize].first_pixel + 4;
|
|
|
|
let in_extent =
|
|
Bounds2i::from_points(Point2i::new(x_start, y_start), Point2i::new(x_end, y_end));
|
|
|
|
let mut in_buf = vec![0.0; in_extent.area() as usize * n_channels];
|
|
self.copy_rect_out(in_extent, &mut in_buf, wrap_mode);
|
|
|
|
let out_buf = compute_resize_tile(
|
|
&in_buf,
|
|
in_extent,
|
|
*out_extent,
|
|
n_channels,
|
|
&x_weights,
|
|
&y_weights,
|
|
);
|
|
|
|
let mut guard = resampled_image.lock().unwrap();
|
|
guard.copy_rect_in(*out_extent, &out_buf);
|
|
});
|
|
|
|
Arc::try_unwrap(resampled_image)
|
|
.unwrap()
|
|
.into_inner()
|
|
.unwrap()
|
|
}
|
|
|
|
pub fn generate_pyramid(base: Image, _wrap: WrapMode) -> Vec<Image> {
|
|
let mut levels = vec![base];
|
|
let internal_wrap = WrapMode2D {
|
|
uv: [WrapMode::Clamp; 2],
|
|
};
|
|
|
|
loop {
|
|
let prev = levels.last().unwrap();
|
|
let old = prev.resolution;
|
|
if old.x() == 1 && old.y() == 1 {
|
|
break;
|
|
}
|
|
|
|
let new_res = Point2i::new((old.x() / 2).max(1), (old.y() / 2).max(1));
|
|
let mut next = Image::new(
|
|
prev.format,
|
|
new_res,
|
|
prev.channel_names.clone(),
|
|
prev.encoding,
|
|
);
|
|
|
|
match &mut next.pixels {
|
|
PixelData::U8(d) => downsample_kernel(d, new_res, prev, internal_wrap),
|
|
PixelData::F16(d) => downsample_kernel(d, new_res, prev, internal_wrap),
|
|
PixelData::F32(d) => downsample_kernel(d, new_res, prev, internal_wrap),
|
|
}
|
|
levels.push(next);
|
|
}
|
|
levels
|
|
}
|
|
}
|
|
|
|
fn flip_y_kernel<T: PixelStorage>(pixels: &mut [T], res: Point2i, channels: usize) {
|
|
let w = res.x() as usize;
|
|
let h = res.y() as usize;
|
|
let stride = w * channels;
|
|
|
|
for y in 0..(h / 2) {
|
|
let bot = h - 1 - y;
|
|
for i in 0..stride {
|
|
pixels.swap(y * stride + i, bot * stride + i);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn crop_kernel<T: PixelStorage>(
|
|
src: &[T],
|
|
dst: &mut [T],
|
|
src_res: Point2i,
|
|
bounds: Bounds2i,
|
|
channels: usize,
|
|
) {
|
|
let dst_w = (bounds.p_max.x() - bounds.p_min.x()) as usize;
|
|
// let dst_h = (bounds.p_max.y() - bounds.p_min.y()) as usize;
|
|
|
|
dst.par_chunks_mut(dst_w * channels)
|
|
.enumerate()
|
|
.for_each(|(dy, dst_row)| {
|
|
let sy = bounds.p_min.y() as usize + dy;
|
|
let sx_start = bounds.p_min.x() as usize;
|
|
|
|
let src_offset = (sy * src_res.x() as usize + sx_start) * channels;
|
|
let count = dst_w * channels;
|
|
|
|
dst_row.copy_from_slice(&src[src_offset..src_offset + count]);
|
|
});
|
|
}
|
|
|
|
fn copy_rect_out_kernel<T: PixelStorage>(
|
|
src: &[T],
|
|
image: &Image,
|
|
extent: Bounds2i,
|
|
buf: &mut [Float],
|
|
wrap: WrapMode2D,
|
|
) {
|
|
let w = (extent.p_max.x() - extent.p_min.x()) as usize;
|
|
let channels = image.n_channels();
|
|
let enc = image.encoding;
|
|
let res = image.resolution;
|
|
|
|
buf.par_chunks_mut(w * channels)
|
|
.enumerate()
|
|
.for_each(|(y_rel, row_buf)| {
|
|
let y = extent.p_min.y() + y_rel as i32;
|
|
for x_rel in 0..w {
|
|
let x = extent.p_min.x() + x_rel as i32;
|
|
|
|
// This allows us to use 'src' directly (Fast Path).
|
|
if x >= 0 && x < res.x() && y >= 0 && y < res.y() {
|
|
let offset = (y as usize * res.x() as usize + x as usize) * channels;
|
|
|
|
for c in 0..channels {
|
|
row_buf[x_rel * channels + c] = T::to_linear(src[offset + c], enc);
|
|
}
|
|
} else {
|
|
// Slow path: Out of bounds, requires Wrap Mode logic.
|
|
// We fall back to get_channel which handles the wrapping math.
|
|
let p = Point2i::new(x, y);
|
|
for c in 0..channels {
|
|
row_buf[x_rel * channels + c] = image.get_channel(p, c, wrap);
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
fn copy_rect_in_kernel<T: PixelStorage>(
|
|
dst: &mut [T],
|
|
res: Point2i,
|
|
channels: usize,
|
|
enc: crate::utils::color::ColorEncoding,
|
|
extent: Bounds2i,
|
|
buf: &[Float],
|
|
) {
|
|
let w = (extent.p_max.x() - extent.p_min.x()) as usize;
|
|
let res_x = res.x() as usize;
|
|
|
|
let rows = buf.chunks(w * channels);
|
|
for (y_rel, row) in rows.enumerate() {
|
|
let y = extent.p_min.y() + y_rel as i32;
|
|
if y < 0 || y >= res.y() {
|
|
continue;
|
|
}
|
|
|
|
let dst_row_start = (y as usize * res_x) * channels;
|
|
|
|
for (x_rel, &val) in row.iter().enumerate() {
|
|
let c = x_rel % channels;
|
|
let x_pixel = x_rel / channels;
|
|
let x = extent.p_min.x() + x_pixel as i32;
|
|
|
|
if x >= 0 && x < res.x() {
|
|
let idx = dst_row_start + (x as usize * channels) + c;
|
|
dst[idx] = T::from_linear(val, enc);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn downsample_kernel<T: PixelStorage>(
|
|
dst: &mut [T],
|
|
dst_res: Point2i,
|
|
prev: &Image,
|
|
wrap: WrapMode2D,
|
|
) {
|
|
let w = dst_res.x() as usize;
|
|
let channels = prev.n_channels();
|
|
let enc = prev.encoding;
|
|
let old_res = prev.resolution;
|
|
|
|
dst.par_chunks_mut(w * channels)
|
|
.enumerate()
|
|
.for_each(|(y, row)| {
|
|
let src_y = y * 2;
|
|
for x in 0..w {
|
|
let src_x = x * 2;
|
|
for c in 0..channels {
|
|
let mut sum = 0.0;
|
|
let mut count = 0.0;
|
|
|
|
for dy in 0..2 {
|
|
for dx in 0..2 {
|
|
let sx = src_x as i32 + dx;
|
|
let sy = src_y as i32 + dy;
|
|
if sx < old_res.x() && sy < old_res.y() {
|
|
sum += prev.get_channel(Point2i::new(sx, sy), c, wrap);
|
|
count += 1.0;
|
|
}
|
|
}
|
|
}
|
|
|
|
let avg = if count > 0.0 { sum / count } else { 0.0 };
|
|
row[x * channels + c] = T::from_linear(avg, enc);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
fn resample_weights(old_res: usize, new_res: usize) -> Vec<ResampleWeight> {
|
|
let filter_radius = 2.0;
|
|
let tau = 2.0;
|
|
(0..new_res)
|
|
.map(|i| {
|
|
let center = (i as Float + 0.5) * old_res as Float / new_res as Float;
|
|
let first_pixel = ((center - filter_radius) + 0.5).floor() as i32;
|
|
let mut weights = [0.0; 4];
|
|
let mut sum = 0.0;
|
|
for j in 0..4 {
|
|
let pos = (first_pixel + j) as Float + 0.5;
|
|
weights[j as usize] = windowed_sinc(pos - center, filter_radius, tau);
|
|
sum += weights[j as usize];
|
|
}
|
|
let inv_sum = 1.0 / sum;
|
|
for w in &mut weights {
|
|
*w *= inv_sum;
|
|
}
|
|
ResampleWeight {
|
|
first_pixel,
|
|
weight: weights,
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn generate_tiles(res: Point2i, tile_size: i32) -> Vec<Bounds2i> {
|
|
let nx = (res.x() + tile_size - 1) / tile_size;
|
|
let ny = (res.y() + tile_size - 1) / tile_size;
|
|
let mut tiles = Vec::new();
|
|
for y in 0..ny {
|
|
for x in 0..nx {
|
|
let p_min = Point2i::new(x * tile_size, y * tile_size);
|
|
let p_max = Point2i::new(
|
|
(x * tile_size + tile_size).min(res.x()),
|
|
(y * tile_size + tile_size).min(res.y()),
|
|
);
|
|
tiles.push(Bounds2i::from_points(p_min, p_max));
|
|
}
|
|
}
|
|
tiles
|
|
}
|
|
|
|
fn compute_resize_tile(
|
|
in_buf: &[Float],
|
|
in_extent: Bounds2i,
|
|
out_extent: Bounds2i,
|
|
n_channels: usize,
|
|
x_w: &[ResampleWeight],
|
|
y_w: &[ResampleWeight],
|
|
) -> Vec<Float> {
|
|
let nx_out = (out_extent.p_max.x() - out_extent.p_min.x()) as usize;
|
|
let ny_out = (out_extent.p_max.y() - out_extent.p_min.y()) as usize;
|
|
let nx_in = (in_extent.p_max.x() - in_extent.p_min.x()) as usize;
|
|
let ny_in = (in_extent.p_max.y() - in_extent.p_min.y()) as usize;
|
|
|
|
let mut x_buf = vec![0.0; n_channels * ny_in * nx_out];
|
|
|
|
// Resize X
|
|
for y in 0..ny_in {
|
|
for x in 0..nx_out {
|
|
let x_global = out_extent.p_min.x() + x as i32;
|
|
let w = &x_w[x_global as usize];
|
|
let x_in_start = (w.first_pixel - in_extent.p_min.x()) as usize;
|
|
|
|
let in_idx_base = (y * nx_in + x_in_start) * n_channels;
|
|
let out_idx_base = (y * nx_out + x) * n_channels;
|
|
|
|
for c in 0..n_channels {
|
|
let mut val = 0.0;
|
|
for k in 0..4 {
|
|
val += w.weight[k] * in_buf[in_idx_base + k * n_channels + c];
|
|
}
|
|
x_buf[out_idx_base + c] = val;
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut out_buf = vec![0.0; n_channels * nx_out * ny_out];
|
|
|
|
// Resize Y
|
|
for x in 0..nx_out {
|
|
for y in 0..ny_out {
|
|
let y_global = out_extent.p_min.y() + y as i32;
|
|
let w = &y_w[y_global as usize];
|
|
let y_in_start = (w.first_pixel - in_extent.p_min.y()) as usize;
|
|
|
|
let in_idx_base = (y_in_start * nx_out + x) * n_channels;
|
|
let out_idx_base = (y * nx_out + x) * n_channels;
|
|
let stride = nx_out * n_channels;
|
|
|
|
for c in 0..n_channels {
|
|
let mut val = 0.0;
|
|
for k in 0..4 {
|
|
val += w.weight[k] * x_buf[in_idx_base + k * stride + c];
|
|
}
|
|
out_buf[out_idx_base + c] = val.max(0.0);
|
|
}
|
|
}
|
|
}
|
|
out_buf
|
|
}
|