diff --git a/crates/kornia-core/src/storage.rs b/crates/kornia-core/src/storage.rs index c9b6ec32..319efc69 100644 --- a/crates/kornia-core/src/storage.rs +++ b/crates/kornia-core/src/storage.rs @@ -160,6 +160,54 @@ where } } + /// Creates a new `TensorStorage` from a slice of data. + /// + /// # Arguments + /// + /// * `data` - A slice containing the data to be stored. + /// * `alloc` - The allocator to use for creating the storage. + /// + /// # Returns + /// + /// A new `TensorStorage` instance containing a copy of the input data. + /// + /// # Errors + /// + /// Returns a `TensorAllocatorError` if the allocation fails. + pub fn from_slice(data: &[T], alloc: A) -> Self { + let buffer = Buffer::from_slice_ref(data); + Self { + data: buffer.into(), + alloc: Arc::new(alloc), + } + } + + /// Creates a new tensor storage from an existing raw pointer with the given allocator. + /// + /// # Arguments + /// + /// * `ptr` - The existing raw pointer to the tensor data. + /// * `len` - The number of elements in the tensor storage. + /// * `alloc` - A reference to the allocator used to allocate the tensor storage. + /// + /// # Safety + /// + /// The pointer must be properly aligned and have the correct length. + pub unsafe fn from_ptr(ptr: *mut T, len: usize, alloc: &A) -> Self { + // create the buffer + let buffer = Buffer::from_custom_allocation( + NonNull::new_unchecked(ptr as *mut u8), + len * std::mem::size_of::(), + Arc::new(()), + ); + + // create tensor storage + Self { + data: buffer.into(), + alloc: Arc::new(alloc.clone()), + } + } + /// Returns the allocator used to allocate the tensor storage. #[inline] pub fn alloc(&self) -> &A { @@ -213,52 +261,6 @@ where pub fn get_unchecked(&self, index: usize) -> &T { unsafe { self.data.get_unchecked(index) } } - - /// Creates a new `TensorStorage` from a slice of data. - /// - /// # Arguments - /// - /// * `data` - A slice containing the data to be stored. - /// * `alloc` - The allocator to use for creating the storage. - /// - /// # Returns - /// - /// A new `TensorStorage` instance containing a copy of the input data. - /// - /// # Errors - /// - /// Returns a `TensorAllocatorError` if the allocation fails. - pub fn from_slice(data: &[T], alloc: A) -> Result { - let mut storage = Self::new(data.len(), alloc)?; - storage.as_mut_slice().copy_from_slice(data); - Ok(storage) - } - - /// Creates a new tensor storage from an existing raw pointer with the given allocator. - /// - /// # Arguments - /// - /// * `ptr` - The existing raw pointer to the tensor data. - /// * `len` - The number of elements in the tensor storage. - /// * `alloc` - A reference to the allocator used to allocate the tensor storage. - /// - /// # Safety - /// - /// The pointer must be properly aligned and have the correct length. - pub unsafe fn from_ptr(ptr: *mut T, len: usize, alloc: &A) -> Self { - // create the buffer - let buffer = Buffer::from_custom_allocation( - NonNull::new_unchecked(ptr as *mut u8), - len * std::mem::size_of::(), - Arc::new(()), - ); - - // create tensor storage - Self { - data: buffer.into(), - alloc: Arc::new(alloc.clone()), - } - } } /// A new `TensorStorage` instance with cloned data if successful, otherwise an error. diff --git a/crates/kornia-core/src/tensor.rs b/crates/kornia-core/src/tensor.rs index 983bdb4a..9001d67d 100644 --- a/crates/kornia-core/src/tensor.rs +++ b/crates/kornia-core/src/tensor.rs @@ -237,7 +237,7 @@ where if numel != data.len() { return Err(TensorError::InvalidShape(numel)); } - let storage = TensorStorage::from_slice(data, alloc)?; + let storage = TensorStorage::from_slice(data, alloc); let strides = get_strides_from_shape(shape); Ok(Self { storage, diff --git a/crates/kornia-image/src/image.rs b/crates/kornia-image/src/image.rs index 08556ce7..211f27c9 100644 --- a/crates/kornia-image/src/image.rs +++ b/crates/kornia-image/src/image.rs @@ -1,4 +1,4 @@ -use kornia_core::{CpuAllocator, SafeTensorType, Tensor3}; +use kornia_core::{CpuAllocator, SafeTensorType, Tensor, Tensor2, Tensor3}; use crate::error::ImageError; @@ -177,6 +177,27 @@ where Ok(image) } + /// Create a new image from a slice of pixel data. + /// + /// # Arguments + /// + /// * `size` - The size of the image in pixels. + /// * `data` - A slice containing the pixel data. + /// + /// # Returns + /// + /// A new image created from the given size and pixel data. + /// + /// # Errors + /// + /// Returns an error if the length of the data slice doesn't match the image dimensions, + /// or if there's an issue creating the tensor or image. + pub fn from_size_slice(size: ImageSize, data: &[T]) -> Result { + let tensor: Tensor3 = + Tensor::from_shape_slice([size.height, size.width, C], data, CpuAllocator)?; + Image::try_from(tensor) + } + /// Cast the pixel data of the image to a different type. /// /// # Returns @@ -414,9 +435,54 @@ where } } +/// helper to convert an single channel tensor to a kornia image with try into +impl TryFrom> for Image +where + T: SafeTensorType, +{ + type Error = ImageError; + + fn try_from(value: Tensor2) -> Result { + Self::from_size_slice( + ImageSize { + width: value.shape[1], + height: value.shape[0], + }, + value.as_slice(), + ) + } +} + +/// helper to convert an multi channel tensor to a kornia image with try into +impl TryFrom> for Image +where + T: SafeTensorType, +{ + type Error = ImageError; + + fn try_from(value: Tensor3) -> Result { + if value.shape[2] != C { + return Err(ImageError::InvalidChannelShape(value.shape[2], C)); + } + Ok(Self(value)) + } +} + +impl TryInto> for Image +where + T: SafeTensorType, +{ + type Error = ImageError; + + fn try_into(self) -> Result, Self::Error> { + Ok(self.0) + } +} + #[cfg(test)] mod tests { use crate::image::{Image, ImageError, ImageSize}; + use kornia_core::{CpuAllocator, Tensor}; #[test] fn image_size() { @@ -560,4 +626,43 @@ mod tests { Ok(()) } + + #[test] + fn image_from_tensor() -> Result<(), ImageError> { + let tensor = + Tensor::::from_shape_vec([2, 3], vec![0u8; 2 * 3], CpuAllocator)?; + + let image = Image::::try_from(tensor.clone())?; + assert_eq!(image.size().width, 3); + assert_eq!(image.size().height, 2); + assert_eq!(image.num_channels(), 1); + + let image_2: Image = tensor.try_into()?; + assert_eq!(image_2.size().width, 3); + assert_eq!(image_2.size().height, 2); + assert_eq!(image_2.num_channels(), 1); + + Ok(()) + } + + #[test] + fn image_from_tensor_3d() -> Result<(), ImageError> { + let tensor = Tensor::::from_shape_vec( + [2, 3, 4], + vec![0u8; 2 * 3 * 4], + CpuAllocator, + )?; + + let image = Image::::try_from(tensor.clone())?; + assert_eq!(image.size().width, 3); + assert_eq!(image.size().height, 2); + assert_eq!(image.num_channels(), 4); + + let image_2: Image = tensor.try_into()?; + assert_eq!(image_2.size().width, 3); + assert_eq!(image_2.size().height, 2); + assert_eq!(image_2.num_channels(), 4); + + Ok(()) + } }