use crate::{allocator::ImageAllocator, Image, ImageError};
pub fn cast_and_scale<T, U, const C: usize, A1: ImageAllocator, A2: ImageAllocator>(
src: &Image<T, C, A1>,
dst: &mut Image<U, C, A2>,
scale: U,
) -> Result<(), ImageError>
where
T: Copy + num_traits::NumCast,
U: Copy + num_traits::NumCast + std::ops::Mul<U, Output = U>,
{
if src.size() != dst.size() {
return Err(ImageError::InvalidImageSize(
src.width(),
src.height(),
dst.width(),
dst.height(),
));
}
dst.as_slice_mut()
.iter_mut()
.zip(src.as_slice().iter())
.try_for_each(|(out, &inp)| {
let x = U::from(inp).ok_or(ImageError::CastError)?;
*out = x * scale;
Ok::<(), ImageError>(())
})?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::image::ImageSize;
use kornia_tensor::CpuAllocator;
#[test]
fn test_cast_and_scale() -> Result<(), ImageError> {
let image = Image::<u8, 3, CpuAllocator>::new(
ImageSize {
height: 2,
width: 1,
},
vec![0u8, 0, 255, 0, 0, 255],
CpuAllocator,
)?;
let mut image_f64: Image<f64, 3, CpuAllocator> =
Image::from_size_val(image.size(), 0.0, CpuAllocator)?;
super::cast_and_scale(&image, &mut image_f64, 1. / 255.0)?;
let expected = vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0];
assert_eq!(image_f64.as_slice(), expected);
Ok(())
}
}