kornia_image/
image.rs

1use crate::{allocator::ImageAllocator, error::ImageError};
2use kornia_tensor::{Tensor, Tensor2, Tensor3};
3use rayon::prelude::*;
4
5/// Image size in pixels
6///
7/// A struct to represent the size of an image in pixels.
8///
9/// # Examples
10///
11/// ```
12/// use kornia_image::ImageSize;
13///
14/// let image_size = ImageSize {
15///   width: 10,
16///   height: 20,
17/// };
18///
19/// assert_eq!(image_size.width, 10);
20/// assert_eq!(image_size.height, 20);
21/// ```
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23pub struct ImageSize {
24    /// Width of the image in pixels
25    pub width: usize,
26    /// Height of the image in pixels
27    pub height: usize,
28}
29
30/// Pixel data type stored in an image buffer.
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub enum PixelFormat {
33    /// Unsigned 8-bit integer pixels.
34    U8,
35    /// Unsigned 16-bit integer pixels.
36    U16,
37    /// 32-bit floating point pixels.
38    F32,
39}
40
41/// Interpolation mode for the image operations
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum InterpolationMode {
44    /// Bilinear interpolation
45    Bilinear,
46    /// Nearest neighbor interpolation
47    Nearest,
48    /// Lanczos interpolation
49    Lanczos,
50    /// Bicubic interpolation
51    Bicubic,
52}
53
54/// Image layout metadata (size, channels, pixel format).
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub struct ImageLayout {
57    /// Image size in pixels.
58    pub image_size: ImageSize,
59    /// Number of channels per pixel.
60    pub channels: u8,
61    /// Pixel data format.
62    pub pixel_format: PixelFormat,
63}
64
65impl ImageLayout {
66    /// Create a new image layout descriptor.
67    pub fn new(image_size: ImageSize, channels: u8, pixel_format: PixelFormat) -> Self {
68        Self {
69            image_size,
70            channels,
71            pixel_format,
72        }
73    }
74}
75
76impl std::fmt::Display for ImageSize {
77    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
78        write!(
79            f,
80            "ImageSize {{ width: {}, height: {} }}",
81            self.width, self.height
82        )
83    }
84}
85
86impl From<[usize; 2]> for ImageSize {
87    fn from(size: [usize; 2]) -> Self {
88        ImageSize {
89            width: size[0],
90            height: size[1],
91        }
92    }
93}
94
95impl From<ImageSize> for [u32; 2] {
96    fn from(size: ImageSize) -> Self {
97        [size.width as u32, size.height as u32]
98    }
99}
100
101impl ImageSize {
102    /// Converts (y, x) coordinates to a flat index (row-major order).
103    #[inline]
104    pub fn index(&self, y: usize, x: usize) -> usize {
105        y * self.width + x
106    }
107
108    /// Converts a flat index to (y, x) coordinates (row-major order).
109    #[inline]
110    pub fn coords(&self, idx: usize) -> (usize, usize) {
111        let y = idx / self.width;
112        let x = idx % self.width;
113        (y, x)
114    }
115}
116
117#[derive(Clone)]
118/// Represents an image with pixel data.
119///
120/// The image is represented as a 3D Tensor with shape (H, W, C), where H is the height of the image,
121pub struct Image<T, const C: usize, A: ImageAllocator>(pub Tensor3<T, A>);
122
123/// helper to deference the inner tensor
124impl<T, const C: usize, A: ImageAllocator> std::ops::Deref for Image<T, C, A> {
125    type Target = Tensor3<T, A>;
126
127    // Define the deref method to return a reference to the inner Tensor3<T>.
128    fn deref(&self) -> &Self::Target {
129        &self.0
130    }
131}
132
133/// helper to deference the inner tensor
134impl<T, const C: usize, A: ImageAllocator> std::ops::DerefMut for Image<T, C, A> {
135    // Define the deref_mut method to return a mutable reference to the inner Tensor3<T>.
136    fn deref_mut(&mut self) -> &mut Self::Target {
137        &mut self.0
138    }
139}
140
141impl<T, const C: usize, A: ImageAllocator> Image<T, C, A> {
142    /// Create a new image from pixel data.
143    ///
144    /// # Arguments
145    ///
146    /// * `size` - The size of the image in pixels.
147    /// * `data` - The pixel data of the image.
148    /// * `alloc` - The allocator of the image
149    ///
150    /// # Returns
151    ///
152    /// A new image with the given pixel data.
153    ///
154    /// # Errors
155    ///
156    /// If the length of the pixel data does not match the image size, an error is returned.
157    ///
158    /// # Examples
159    ///
160    /// ```
161    /// use kornia_image::{Image, ImageSize};
162    /// use kornia_image::allocator::CpuAllocator;
163    ///
164    /// let image = Image::<u8, 3, _>::new(
165    ///    ImageSize {
166    ///       width: 10,
167    ///      height: 20,
168    ///  },
169    /// vec![0u8; 10 * 20 * 3],
170    /// CpuAllocator
171    /// ).unwrap();
172    ///
173    /// assert_eq!(image.size().width, 10);
174    /// assert_eq!(image.size().height, 20);
175    /// assert_eq!(image.num_channels(), 3);
176    /// ```
177    pub fn new(size: ImageSize, data: Vec<T>, alloc: A) -> Result<Self, ImageError> {
178        // check if the data length matches the image size
179        if data.len() != size.width * size.height * C {
180            return Err(ImageError::InvalidChannelShape(
181                data.len(),
182                size.width * size.height * C,
183            ));
184        }
185
186        // allocate the image data
187        Ok(Self(Tensor3::from_shape_vec(
188            [size.height, size.width, C],
189            data,
190            alloc,
191        )?))
192    }
193
194    /// Create a new image with the given size and default pixel data.
195    ///
196    /// # Arguments
197    ///
198    /// * `size` - The size of the image in pixels.
199    /// * `val` - The default value of the pixel data.
200    /// * `alloc` - The allocator of the image
201    ///
202    /// # Returns
203    ///
204    /// A new image with the given size and default pixel data.
205    ///
206    /// # Errors
207    ///
208    /// If the length of the pixel data does not match the image size, an error is returned.
209    ///
210    /// # Examples
211    ///
212    /// ```
213    /// use kornia_image::{Image, ImageSize};
214    /// use kornia_image::allocator::CpuAllocator;
215    ///
216    /// let image = Image::<u8, 3, _>::from_size_val(
217    ///   ImageSize {
218    ///     width: 10,
219    ///    height: 20,
220    /// }, 0u8, CpuAllocator).unwrap();
221    ///
222    /// assert_eq!(image.size().width, 10);
223    /// assert_eq!(image.size().height, 20);
224    /// assert_eq!(image.num_channels(), 3);
225    /// ```
226    pub fn from_size_val(size: ImageSize, val: T, alloc: A) -> Result<Self, ImageError>
227    where
228        T: Clone,
229    {
230        let data = vec![val; size.width * size.height * C];
231        let image = Image::new(size, data, alloc)?;
232
233        Ok(image)
234    }
235
236    /// Create a new image from raw parts.
237    ///
238    /// # Arguments
239    ///
240    /// * `size` - The size of the image in pixels.
241    /// * `data` - A pointer to the pixel data.
242    /// * `len` - The length of the pixel data.
243    ///
244    /// # Returns
245    ///
246    /// A new image created from the given size and pixel data.
247    ///
248    /// # Safety
249    ///
250    /// The pointer must be non-null and the length must be valid.
251    pub unsafe fn from_raw_parts(
252        size: ImageSize,
253        data: *const T,
254        len: usize,
255        alloc: A,
256    ) -> Result<Self, ImageError>
257    where
258        T: Clone,
259    {
260        Tensor::from_raw_parts([size.height, size.width, C], data, len, alloc)?.try_into()
261    }
262
263    /// Create a new image from a slice of pixel data.
264    ///
265    /// # Arguments
266    ///
267    /// * `size` - The size of the image in pixels.
268    /// * `data` - A slice containing the pixel data.
269    ///
270    /// # Returns
271    ///
272    /// A new image created from the given size and pixel data.
273    ///
274    /// # Errors
275    ///
276    /// Returns an error if the length of the data slice doesn't match the image dimensions,
277    /// or if there's an issue creating the tensor or image.
278    pub fn from_size_slice(size: ImageSize, data: &[T], alloc: A) -> Result<Self, ImageError>
279    where
280        T: Clone,
281    {
282        let tensor: Tensor3<T, A> =
283            Tensor::from_shape_slice([size.height, size.width, C], data, alloc)?;
284        Image::try_from(tensor)
285    }
286
287    /// Map the pixel data of the image to a different type.
288    ///
289    /// # Arguments
290    ///
291    /// * `f` - A function that takes a pixel value and returns a new pixel value.
292    ///
293    /// # Returns
294    ///
295    /// A new image with the pixel data mapped to the new type.
296    pub fn map<U>(&self, f: impl Fn(&T) -> U) -> Result<Image<U, C, A>, ImageError> {
297        let data = self.as_slice().iter().map(f).collect::<Vec<U>>();
298        let alloc = self.storage.alloc();
299        Image::<U, C, A>::new(self.size(), data, alloc.clone())
300    }
301
302    /// Cast the pixel data of the image to a different type.
303    ///
304    /// Each pixel value is converted using [`num_traits::NumCast`], which supports
305    /// lossy casts (e.g. `f32` → `u8` via truncation). If any value cannot be
306    /// represented in the target type, the entire operation returns
307    /// [`ImageError::CastError`].
308    ///
309    /// The output image has the same shape (height × width × channels) as the
310    /// source image and uses the same allocator.
311    ///
312    /// # Returns
313    ///
314    /// A new image with the pixel data cast to the target type `U`.
315    ///
316    /// # Errors
317    ///
318    /// Returns [`ImageError::CastError`] if any pixel value cannot be represented
319    /// as type `U` (e.g. `f32::MAX` cast to `u8`).
320    ///
321    /// # Example
322    ///
323    /// ```
324    /// use kornia_image::{Image, ImageSize};
325    /// use kornia_image::allocator::CpuAllocator;
326    ///
327    /// let image_u8 = Image::<u8, 3, CpuAllocator>::new(
328    ///     ImageSize { width: 2, height: 1 },
329    ///     vec![0u8, 128, 255, 10, 20, 30],
330    ///     CpuAllocator,
331    /// ).unwrap();
332    ///
333    /// let image_f32 = image_u8.cast::<f32>().unwrap();
334    /// assert_eq!(image_f32.as_slice()[2], 255.0f32);
335    /// ```
336    pub fn cast<U>(&self) -> Result<Image<U, C, A>, ImageError>
337    where
338        U: num_traits::NumCast + Copy,
339        T: num_traits::NumCast + Copy,
340    {
341        let data = self
342            .as_slice()
343            .iter()
344            .map(|&x| U::from(x).ok_or(ImageError::CastError))
345            .collect::<Result<Vec<U>, ImageError>>()?;
346
347        let alloc = self.storage.alloc().clone();
348        let tensor = Tensor3::from_shape_vec(self.0.shape, data, alloc)?;
349        Ok(Image(tensor))
350    }
351
352    /// Get a channel of the image.
353    /// # Arguments
354    ///
355    /// * `channel` - The channel to get.
356    ///
357    /// # Returns
358    ///
359    /// A new image with the given channel.
360    ///
361    /// # Errors
362    ///
363    /// If the channel index is out of bounds, an error is returned.
364    pub fn channel(&self, channel: usize) -> Result<Image<T, 1, A>, ImageError>
365    where
366        T: Clone,
367    {
368        if channel >= C {
369            return Err(ImageError::ChannelIndexOutOfBounds(channel, C));
370        }
371
372        let channel_data = self
373            .as_slice()
374            .iter()
375            .skip(channel)
376            .step_by(C)
377            .cloned()
378            .collect();
379
380        let alloc = self.storage.alloc();
381
382        Image::new(self.size(), channel_data, alloc.clone())
383    }
384
385    /// Split the image into its channels.
386    ///
387    /// # Returns
388    ///
389    /// A vector of images, each containing one channel of the original image.
390    ///
391    /// # Examples
392    ///
393    /// ```
394    /// use kornia_image::{Image, ImageSize};
395    /// use kornia_image::allocator::CpuAllocator;
396    ///
397    /// let image = Image::<f32, 2, _>::from_size_val(
398    ///   ImageSize {
399    ///    width: 10,
400    ///   height: 20,
401    /// },
402    /// 0.0f32,
403    /// CpuAllocator).unwrap();
404    ///
405    /// let channels = image.split_channels().unwrap();
406    /// assert_eq!(channels.len(), 2);
407    /// ```
408    pub fn split_channels(&self) -> Result<Vec<Image<T, 1, A>>, ImageError>
409    where
410        T: Copy,
411    {
412        let mut channels = Vec::with_capacity(C);
413
414        for i in 0..C {
415            channels.push(self.channel(i)?);
416        }
417
418        Ok(channels)
419    }
420
421    /// Get the size of the image in pixels.
422    pub fn size(&self) -> ImageSize {
423        ImageSize {
424            width: self.shape[1],
425            height: self.shape[0],
426        }
427    }
428
429    /// Get the number of columns of the image.
430    pub fn cols(&self) -> usize {
431        self.shape[1]
432    }
433
434    /// Get the number of rows of the image.
435    pub fn rows(&self) -> usize {
436        self.shape[0]
437    }
438
439    /// Get the width of the image in pixels.
440    pub fn width(&self) -> usize {
441        self.cols()
442    }
443
444    /// Get the height of the image in pixels.
445    pub fn height(&self) -> usize {
446        self.rows()
447    }
448
449    /// Get the number of channels in the image.
450    pub fn num_channels(&self) -> usize {
451        C
452    }
453
454    /// Cast the pixel data to a different type and scale it.
455    ///
456    /// # Arguments
457    ///
458    /// * `scale` - The scale to multiply the pixel data with.
459    ///
460    /// # Returns
461    ///
462    /// A new image with the pixel data cast to the new type and scaled.
463    ///
464    /// # Errors
465    ///
466    /// If the pixel data cannot be cast to the new type, an error is returned.
467    ///
468    /// # Examples
469    ///
470    /// ```
471    /// use kornia_image::{Image, ImageSize};
472    /// use kornia_image::allocator::CpuAllocator;
473    ///
474    /// let data = vec![0u8, 0, 255, 0, 0, 255];
475    ///
476    /// let image_u8 = Image::<u8, 3, _>::new(
477    /// ImageSize {
478    ///   height: 2,
479    ///   width: 1,
480    /// },
481    /// data,
482    /// CpuAllocator
483    /// ).unwrap();
484    ///
485    /// let image_f32 = image_u8.cast_and_scale::<f32>(1. / 255.0).unwrap();
486    ///
487    /// assert_eq!(image_f32.get([1, 0, 2]), Some(&1.0f32));
488    /// ```
489    #[allow(clippy::uninit_vec)]
490    pub fn cast_and_scale<U>(self, scale: U) -> Result<Image<U, C, A>, ImageError>
491    where
492        U: num_traits::NumCast + std::ops::Mul<Output = U> + Clone + Copy + Send + Sync,
493        T: num_traits::NumCast + Clone + Copy + Send + Sync,
494    {
495        let slice = self.as_slice();
496        let mut casted_data = Vec::with_capacity(slice.len());
497        // SAFETY: Each element is written to with no reads beforehand.
498        unsafe {
499            casted_data.set_len(slice.len());
500        }
501
502        slice
503            .par_iter()
504            .zip(casted_data.par_iter_mut())
505            .try_for_each(|(&x, out)| {
506                let xu = U::from(x).ok_or(ImageError::CastError)?;
507                *out = xu * scale;
508                Ok::<(), ImageError>(())
509            })?;
510
511        let alloc = self.storage.alloc();
512        Image::new(self.size(), casted_data, alloc.clone())
513    }
514
515    /// Cast the pixel data to a different type and scale it.
516    ///
517    /// # Arguments
518    ///
519    /// * `scale` - The scale to multiply the pixel data with.
520    ///
521    /// # Returns
522    ///
523    /// A new image with the pixel data cast to the new type and scaled.
524    #[allow(clippy::uninit_vec)]
525    pub fn scale_and_cast<U>(&self, scale: T) -> Result<Image<U, C, A>, ImageError>
526    where
527        U: num_traits::NumCast + Clone + Copy + Send + Sync,
528        T: num_traits::NumCast + std::ops::Mul<Output = T> + Clone + Copy + Send + Sync,
529    {
530        let slice = self.as_slice();
531        let mut casted_data = Vec::with_capacity(slice.len());
532        // SAFETY: Each element is written to with no reads beforehand.
533        unsafe {
534            casted_data.set_len(slice.len());
535        }
536
537        slice
538            .par_iter()
539            .zip(casted_data.par_iter_mut())
540            .try_for_each(|(&x, out)| {
541                *out = U::from(x * scale).ok_or(ImageError::CastError)?;
542                Ok::<(), ImageError>(())
543            })?;
544
545        let alloc = self.storage.alloc();
546        Image::new(self.size(), casted_data, alloc.clone())
547    }
548
549    /// Get the pixel data of the image.
550    ///
551    /// NOTE: this is method is for convenience and not optimized for performance.
552    /// We recommend using iterators over the data slice.
553    ///
554    /// # Arguments
555    ///
556    /// * `x` - The x-coordinate of the pixel.
557    /// * `y` - The y-coordinate of the pixel.
558    /// * `ch` - The channel index of the pixel.
559    ///
560    /// # Returns
561    ///
562    /// The pixel value at the given coordinates.
563    pub fn get_pixel(&self, x: usize, y: usize, ch: usize) -> Result<&T, ImageError> {
564        if x >= self.width() || y >= self.height() {
565            return Err(ImageError::PixelIndexOutOfBounds(
566                x,
567                y,
568                self.width(),
569                self.height(),
570            ));
571        }
572
573        if ch >= C {
574            return Err(ImageError::ChannelIndexOutOfBounds(ch, C));
575        }
576
577        let val = match self.get([y, x, ch]) {
578            Some(v) => v,
579            None => return Err(ImageError::ImageDataNotContiguous),
580        };
581
582        Ok(val)
583    }
584
585    /// Set the pixel value at the given coordinates.
586    ///
587    /// NOTE: this is method is for convenience and not optimized for performance.
588    /// We recommend creating a mutable slice and operating on it directly.
589    ///
590    /// # Arguments
591    ///
592    /// * `x` - The x-coordinate of the pixel.
593    /// * `y` - The y-coordinate of the pixel.
594    /// * `ch` - The channel index of the pixel.
595    /// * `val` - The value to set the pixel to.
596    ///
597    /// # Returns
598    ///
599    /// The pixel value at the given coordinates.
600    pub fn set_pixel(&mut self, x: usize, y: usize, ch: usize, val: T) -> Result<(), ImageError> {
601        if x >= self.width() || y >= self.height() {
602            return Err(ImageError::PixelIndexOutOfBounds(
603                x,
604                y,
605                self.width(),
606                self.height(),
607            ));
608        }
609
610        if ch >= C {
611            return Err(ImageError::ChannelIndexOutOfBounds(ch, C));
612        }
613
614        let idx = y * self.width() * C + x * C + ch;
615        self.as_slice_mut()[idx] = val;
616
617        Ok(())
618    }
619
620    /// Convert the image to a vector.
621    pub fn into_vec(self) -> Vec<T> {
622        self.0.into_vec()
623    }
624
625    /// Get a copy of the image data as a vector.
626    pub fn to_vec(&self) -> Vec<T>
627    where
628        T: Clone,
629    {
630        self.as_slice().to_vec()
631    }
632}
633
634/// helper to convert an single channel tensor to a kornia image with try into
635impl<T, A: ImageAllocator> TryFrom<Tensor2<T, A>> for Image<T, 1, A>
636where
637    T: Clone,
638{
639    type Error = ImageError;
640
641    fn try_from(value: Tensor2<T, A>) -> Result<Self, Self::Error> {
642        let alloc = value.storage.alloc();
643
644        Self::from_size_slice(
645            ImageSize {
646                width: value.shape[1],
647                height: value.shape[0],
648            },
649            value.as_slice(),
650            alloc.clone(),
651        )
652    }
653}
654
655/// helper to convert an multi channel tensor to a kornia image with try into
656impl<T, const C: usize, A: ImageAllocator> TryFrom<Tensor3<T, A>> for Image<T, C, A> {
657    type Error = ImageError;
658
659    fn try_from(value: Tensor3<T, A>) -> Result<Self, Self::Error> {
660        if value.shape[2] != C {
661            return Err(ImageError::InvalidChannelShape(value.shape[2], C));
662        }
663        Ok(Self(value))
664    }
665}
666
667impl<T, const C: usize, A: ImageAllocator> TryInto<Tensor3<T, A>> for Image<T, C, A> {
668    type Error = ImageError;
669
670    fn try_into(self) -> Result<Tensor3<T, A>, Self::Error> {
671        Ok(self.0)
672    }
673}
674
675#[cfg(test)]
676mod tests {
677    use crate::image::{Image, ImageError, ImageSize};
678    use kornia_tensor::{CpuAllocator, Tensor};
679
680    #[test]
681    fn test_image_size() {
682        let image_size = ImageSize {
683            width: 10,
684            height: 20,
685        };
686        assert_eq!(image_size.width, 10);
687        assert_eq!(image_size.height, 20);
688    }
689
690    #[test]
691    fn test_image_size_index_coords() {
692        let size = ImageSize {
693            width: 8,
694            height: 6,
695        };
696        // Test (y, x) -> idx
697        assert_eq!(size.index(0, 0), 0);
698        assert_eq!(size.index(1, 0), 8);
699        assert_eq!(size.index(2, 3), 2 * 8 + 3);
700        assert_eq!(size.index(5, 7), 5 * 8 + 7);
701
702        // Test idx -> (y, x)
703        assert_eq!(size.coords(0), (0, 0));
704        assert_eq!(size.coords(8), (1, 0));
705        assert_eq!(size.coords(19), (2, 3));
706        assert_eq!(size.coords(47), (5, 7));
707    }
708
709    #[test]
710    fn test_image_smoke() -> Result<(), ImageError> {
711        let image = Image::<u8, 3, CpuAllocator>::new(
712            ImageSize {
713                width: 10,
714                height: 20,
715            },
716            vec![0u8; 10 * 20 * 3],
717            CpuAllocator,
718        )?;
719        assert_eq!(image.size().width, 10);
720        assert_eq!(image.size().height, 20);
721        assert_eq!(image.num_channels(), 3);
722
723        Ok(())
724    }
725
726    #[test]
727    fn test_image_from_vec() -> Result<(), ImageError> {
728        let image: Image<f32, 3, CpuAllocator> = Image::new(
729            ImageSize {
730                height: 3,
731                width: 2,
732            },
733            vec![0.0; 3 * 2 * 3],
734            CpuAllocator,
735        )?;
736        assert_eq!(image.size().width, 2);
737        assert_eq!(image.size().height, 3);
738        assert_eq!(image.num_channels(), 3);
739
740        Ok(())
741    }
742
743    #[test]
744    fn test_image_from_empty_vec() -> Result<(), ImageError> {
745        let image: Result<Image<f32, 1, CpuAllocator>, ImageError> = Image::new(
746            ImageSize {
747                height: 0,
748                width: 0,
749            },
750            vec![0.0; 0],
751            CpuAllocator,
752        );
753        assert!(
754            image.is_ok(),
755            "Image::new should create an empty image and drop it without segfault"
756        );
757
758        Ok(())
759    }
760
761    #[test]
762    fn test_image_cast() -> Result<(), ImageError> {
763        let data = vec![0, 1, 2, 3, 4, 5];
764        let image_u8 = Image::<_, 3, CpuAllocator>::new(
765            ImageSize {
766                height: 2,
767                width: 1,
768            },
769            data,
770            CpuAllocator,
771        )?;
772        assert_eq!(image_u8.get([1, 0, 2]), Some(&5u8));
773
774        let image_i32: Image<i32, 3, CpuAllocator> = image_u8.cast()?;
775        assert_eq!(image_i32.get([1, 0, 2]), Some(&5i32));
776
777        Ok(())
778    }
779
780    #[test]
781    fn test_image_rgbd() -> Result<(), ImageError> {
782        let image = Image::<f32, 4, CpuAllocator>::new(
783            ImageSize {
784                height: 2,
785                width: 3,
786            },
787            vec![0f32; 2 * 3 * 4],
788            CpuAllocator,
789        )?;
790        assert_eq!(image.size().width, 3);
791        assert_eq!(image.size().height, 2);
792        assert_eq!(image.num_channels(), 4);
793
794        Ok(())
795    }
796
797    #[test]
798    fn test_image_channel() -> Result<(), ImageError> {
799        let image = Image::<f32, 3, CpuAllocator>::new(
800            ImageSize {
801                height: 2,
802                width: 1,
803            },
804            vec![0., 1., 2., 3., 4., 5.],
805            CpuAllocator,
806        )?;
807
808        let channel = image.channel(2)?;
809        assert_eq!(channel.get([1, 0, 0]), Some(&5.0f32));
810
811        Ok(())
812    }
813
814    #[test]
815    fn test_image_split_channels() -> Result<(), ImageError> {
816        let image = Image::<f32, 3, CpuAllocator>::new(
817            ImageSize {
818                height: 2,
819                width: 1,
820            },
821            vec![0., 1., 2., 3., 4., 5.],
822            CpuAllocator,
823        )
824        .unwrap();
825        let channels = image.split_channels()?;
826        assert_eq!(channels.len(), 3);
827        assert_eq!(channels[0].get([1, 0, 0]), Some(&3.0f32));
828        assert_eq!(channels[1].get([1, 0, 0]), Some(&4.0f32));
829        assert_eq!(channels[2].get([1, 0, 0]), Some(&5.0f32));
830
831        Ok(())
832    }
833
834    #[test]
835    fn test_scale_and_cast() -> Result<(), ImageError> {
836        let data = vec![0u8, 0, 255, 0, 0, 255];
837        let image_u8 = Image::<u8, 3, CpuAllocator>::new(
838            ImageSize {
839                height: 2,
840                width: 1,
841            },
842            data,
843            CpuAllocator,
844        )?;
845        let image_f32 = image_u8.cast_and_scale::<f32>(1. / 255.0)?;
846        assert_eq!(image_f32.get([1, 0, 2]), Some(&1.0f32));
847
848        Ok(())
849    }
850
851    #[test]
852    fn test_cast_and_scale() -> Result<(), ImageError> {
853        let data = vec![0u8, 0, 255, 0, 0, 255];
854        let image_u8 = Image::<u8, 3, CpuAllocator>::new(
855            ImageSize {
856                height: 2,
857                width: 1,
858            },
859            data,
860            CpuAllocator,
861        )?;
862        let image_f32 = image_u8.cast_and_scale::<f32>(1. / 255.0)?;
863        assert_eq!(image_f32.get([1, 0, 2]), Some(&1.0f32));
864
865        Ok(())
866    }
867
868    #[test]
869    fn test_image_from_tensor() -> Result<(), ImageError> {
870        let data = vec![0u8, 1, 2, 3, 4, 5];
871        let tensor = Tensor::<u8, 2, _>::from_shape_vec([2, 3], data, CpuAllocator)?;
872
873        let image = Image::<u8, 1, CpuAllocator>::try_from(tensor.clone())?;
874        assert_eq!(image.size().width, 3);
875        assert_eq!(image.size().height, 2);
876        assert_eq!(image.num_channels(), 1);
877
878        let image_2: Image<u8, 1, CpuAllocator> = tensor.try_into()?;
879        assert_eq!(image_2.size().width, 3);
880        assert_eq!(image_2.size().height, 2);
881        assert_eq!(image_2.num_channels(), 1);
882
883        Ok(())
884    }
885
886    #[test]
887    fn test_image_from_tensor_3d() -> Result<(), ImageError> {
888        let tensor = Tensor::<u8, 3, CpuAllocator>::from_shape_vec(
889            [2, 3, 4],
890            vec![0u8; 2 * 3 * 4],
891            CpuAllocator,
892        )?;
893
894        let image = Image::<u8, 4, CpuAllocator>::try_from(tensor.clone())?;
895        assert_eq!(image.size().width, 3);
896        assert_eq!(image.size().height, 2);
897        assert_eq!(image.num_channels(), 4);
898
899        let image_2: Image<u8, 4, CpuAllocator> = tensor.try_into()?;
900        assert_eq!(image_2.size().width, 3);
901        assert_eq!(image_2.size().height, 2);
902        assert_eq!(image_2.num_channels(), 4);
903
904        Ok(())
905    }
906
907    #[test]
908    fn test_image_from_raw_parts() -> Result<(), ImageError> {
909        let data = vec![0u8, 1, 2, 3, 4, 5];
910        let image = unsafe {
911            Image::<_, 1, CpuAllocator>::from_raw_parts(
912                [2, 3].into(),
913                data.as_ptr(),
914                data.len(),
915                CpuAllocator,
916            )?
917        };
918        std::mem::forget(data);
919        assert_eq!(image.size().width, 2);
920        assert_eq!(image.size().height, 3);
921        assert_eq!(image.num_channels(), 1);
922        Ok(())
923    }
924
925    #[test]
926    fn test_get_pixel() -> Result<(), ImageError> {
927        let image = Image::<u8, 3, CpuAllocator>::new(
928            ImageSize {
929                height: 2,
930                width: 1,
931            },
932            vec![1, 2, 5, 19, 255, 128],
933            CpuAllocator,
934        )?;
935        assert_eq!(image.get_pixel(0, 0, 0)?, &1);
936        assert_eq!(image.get_pixel(0, 0, 1)?, &2);
937        assert_eq!(image.get_pixel(0, 0, 2)?, &5);
938        assert_eq!(image.get_pixel(0, 1, 0)?, &19);
939        assert_eq!(image.get_pixel(0, 1, 1)?, &255);
940        assert_eq!(image.get_pixel(0, 1, 2)?, &128);
941        Ok(())
942    }
943
944    #[test]
945    fn test_set_pixel() -> Result<(), ImageError> {
946        let mut image = Image::<u8, 3, CpuAllocator>::new(
947            ImageSize {
948                height: 2,
949                width: 1,
950            },
951            vec![1, 2, 5, 19, 255, 128],
952            CpuAllocator,
953        )?;
954
955        image.set_pixel(0, 0, 0, 128)?;
956        image.set_pixel(0, 1, 1, 25)?;
957
958        assert_eq!(image.get_pixel(0, 0, 0)?, &128);
959        assert_eq!(image.get_pixel(0, 1, 1)?, &25);
960
961        Ok(())
962    }
963
964    #[test]
965    fn test_image_map() -> Result<(), ImageError> {
966        let image_u8 = Image::<u8, 1, CpuAllocator>::new(
967            ImageSize {
968                height: 2,
969                width: 1,
970            },
971            vec![0, 128],
972            CpuAllocator,
973        )?;
974
975        let image_f32 = image_u8.map(|x| (x + 2) as f32)?;
976
977        assert_eq!(image_f32.size().width, 1);
978        assert_eq!(image_f32.size().height, 2);
979        assert_eq!(image_f32.num_channels(), 1);
980        assert_eq!(image_f32.get([0, 0, 0]), Some(&2.0f32));
981        assert_eq!(image_f32.get([1, 0, 0]), Some(&130.0f32));
982
983        Ok(())
984    }
985
986    #[test]
987    fn test_cast_round_trip() -> Result<(), ImageError> {
988        // f32 → u8 → f32: values representable as u8 should survive the round trip.
989        let data_f32 = vec![0.0f32, 128.0, 255.0, 10.0, 20.0, 30.0];
990        let image_f32 = Image::<f32, 3, CpuAllocator>::new(
991            ImageSize {
992                height: 2,
993                width: 1,
994            },
995            data_f32.clone(),
996            CpuAllocator,
997        )?;
998
999        let image_u8 = image_f32.cast::<u8>()?;
1000        let image_f32_rt = image_u8.cast::<f32>()?;
1001
1002        assert_eq!(image_f32_rt.size(), image_f32.size());
1003        assert_eq!(image_f32_rt.num_channels(), 3);
1004        for (original, round_tripped) in data_f32.iter().zip(image_f32_rt.as_slice()) {
1005            assert!((*original - round_tripped).abs() < 1.0);
1006        }
1007
1008        Ok(())
1009    }
1010
1011    #[test]
1012    fn test_cast_out_of_range() -> Result<(), ImageError> {
1013        // f32::MAX cannot be represented as u8, so cast must fail.
1014        let image_f32 = Image::<f32, 1, CpuAllocator>::new(
1015            ImageSize {
1016                height: 1,
1017                width: 1,
1018            },
1019            vec![f32::MAX],
1020            CpuAllocator,
1021        )?;
1022
1023        let result = image_f32.cast::<u8>();
1024        assert!(
1025            matches!(result, Err(ImageError::CastError)),
1026            "expected CastError for out-of-range f32 value"
1027        );
1028
1029        Ok(())
1030    }
1031
1032    #[test]
1033    fn test_cast_shape_preservation() -> Result<(), ImageError> {
1034        // Cast a 2×1 image with C=3; result must have identical size and correct values.
1035        let data = vec![0u8, 64, 128, 192, 200, 255];
1036        let image_u8 = Image::<u8, 3, CpuAllocator>::new(
1037            ImageSize {
1038                height: 2,
1039                width: 1,
1040            },
1041            data.clone(),
1042            CpuAllocator,
1043        )?;
1044
1045        let image_f32 = image_u8.cast::<f32>()?;
1046
1047        assert_eq!(image_f32.size(), image_u8.size());
1048        assert_eq!(image_f32.num_channels(), 3);
1049        for (original, casted) in data.iter().zip(image_f32.as_slice()) {
1050            assert_eq!(*casted, *original as f32);
1051        }
1052
1053        Ok(())
1054    }
1055}
kornia_image/image.rs

kornia_image/
image.rs
