Implement Tensor from_raw_parts

crates/kornia-image/src/image.rs

@@ -193,6 +193,25 @@ impl<T, const C: usize> Image<T, C> {
         Image::try_from(tensor)
     }
 
+    /// Create a new image from raw parts.
+    ///
+    /// # Arguments
+    ///
+    /// * `size` - The size of the image in pixels.
+    /// * `data` - A pointer to the pixel data.
+    /// * `len` - The length of the pixel data.
+    ///
+    /// # Returns
+    ///
+    /// A new image created from the given size and pixel data.
+    pub fn from_raw_parts(size: ImageSize, data: *const T, len: usize) -> Result<Self, ImageError>
+    where
+        T: Clone,
+    {
+        let tensor = Tensor::from_raw_parts([size.height, size.width, C], data, len, CpuAllocator)?;
+        Image::try_from(tensor)
+    }
+
     /// Cast the pixel data of the image to a different type.
     ///
     /// # Returns

crates/kornia-io/src/stream/capture.rs

@@ -2,7 +2,7 @@ use std::sync::{Arc, Mutex};
 
 use crate::stream::error::StreamCaptureError;
 use gst::prelude::*;
-use kornia_image::{Image, ImageSize};
+use kornia_image::Image;
 
 /// Represents a stream capture pipeline using GStreamer.
 pub struct StreamCapture {
@@ -160,8 +160,10 @@ impl StreamCapture {
             .ok_or_else(|| StreamCaptureError::GetBufferError)?
             .map_readable()?;
 
-        Image::<u8, 3>::new(ImageSize { width, height }, buffer.as_slice().to_vec())
-            .map_err(|_| StreamCaptureError::CreateImageFrameError)
+        let image = Image::from_raw_parts([width, height].into(), buffer.as_ptr(), buffer.len())
+            .map_err(|_| StreamCaptureError::CreateImageFrameError)?;
+
+        Ok(image)
     }
 }
 

crates/kornia-tensor/src/tensor.rs

@@ -199,7 +199,9 @@ where
         if numel != data.len() {
             return Err(TensorError::InvalidShape(numel));
         }
+
         let storage = TensorStorage::from_vec(data.to_vec(), alloc);
+
         let strides = get_strides_from_shape(shape);
         Ok(Self {
             storage,
@@ -208,6 +210,27 @@ where
         })
     }
 
+    /// Creates a new `Tensor` with the given shape and raw parts.
+    ///
+    /// # Arguments
+    ///
+    /// * `shape` - An array containing the shape of the tensor.
+    /// * `data` - A pointer to the data of the tensor.
+    /// * `len` - The length of the data.
+    /// * `alloc` - The allocator to use.
+    pub fn from_raw_parts(
+        shape: [usize; N],
+        data: *const T,
+        len: usize,
+        alloc: A,
+    ) -> Result<Self, TensorError>
+    where
+        T: Clone,
+    {
+        let data_vec = unsafe { Vec::from_raw_parts(data as *mut T, len, len) };
+        Self::from_shape_vec(shape, data_vec, alloc)
+    }
+
     /// Creates a new `Tensor` with the given shape and a default value.
     ///
     /// # Arguments
@@ -1374,6 +1397,16 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn from_raw_parts() -> Result<(), TensorError> {
+        let data: Vec<u8> = vec![1, 2, 3, 4];
+        let t = Tensor::from_raw_parts([2, 2], data.as_ptr(), data.len(), CpuAllocator)?;
+        std::mem::forget(data);
+        assert_eq!(t.shape, [2, 2]);
+        assert_eq!(t.as_slice(), &[1, 2, 3, 4]);
+        Ok(())
+    }
+
     #[test]
     fn display_2d() -> Result<(), TensorError> {
         let data: [u8; 4] = [1, 2, 3, 4];