Fix memory release

pyproject.toml

@@ -27,6 +27,7 @@ classifiers = [
 [project.optional-dependencies]
 dev = [
   "pytest",
+  "pytest-asyncio",
   "kornia==0.6.3",
   "numpy==1.23.0",
   "jax[cpu]",

src/dlpack_py.rs

@@ -92,14 +92,13 @@ pub fn cvtensor_to_dlpack(x: &cv::Tensor, py: Python) -> PyResult<PyObject> {
     let dlm_tensor_bx = Box::new(dlm_tensor);
 
     // create python capsule
-    let capsule: PyObject = unsafe {
+    let capsule = unsafe {
         let ptr = pyo3::ffi::PyCapsule_New(
-            &*dlm_tensor_bx as *const dlpack::DLManagedTensor as *mut c_void,
+            Box::into_raw(dlm_tensor_bx) as *mut c_void,
             DLPACK_CAPSULE_NAME.as_ptr() as *const c_char,
             Some(dlpack_capsule_destructor as pyo3::ffi::PyCapsule_Destructor),
         );
         PyObject::from_owned_ptr(py, ptr)
     };
-    Box::leak(dlm_tensor_bx); // to hold reference until program exits
     Ok(capsule)
 }

src/tensor.rs

@@ -40,14 +40,13 @@ pub mod cv {
         }
 
         #[pyo3(name = "__dlpack__")]
-        pub fn to_dlpack_py(&self, py: Python) -> PyResult<PyObject> {
+        pub fn dlpack_py(&self, py: Python) -> PyResult<PyObject> {
             cvtensor_to_dlpack(self, py)
         }
 
         #[pyo3(name = "__dlpack_device__")]
-        pub fn to_dlpack_device_py(&self) -> (i32, i32) {
-            let tensor_bx = Box::new(self);
-            let dl_tensor = cvtensor_to_dltensor(&tensor_bx);
+        pub fn dlpack_device_py(&self) -> (i32, i32) {
+            let dl_tensor = cvtensor_to_dltensor(self);
             (
                 dl_tensor.device.device_type as i32,
                 dl_tensor.device.device_id,

test/test_io.py

@@ -1,4 +1,7 @@
 from pathlib import Path
+import pytest
+import random
+import asyncio
 
 import kornia_rs as K
 from kornia_rs import Tensor as cvTensor
@@ -84,3 +87,22 @@ def test_write_read_jpeg():
     read_image = np.from_dlpack(read_tensor)
 
     np.testing.assert_allclose(decoded_image, read_image)
+
+async def encode_frame(i: int) -> bytes:
+    img = (np.random.rand(480, 640, 3) * 255).astype(np.uint8)
+    frame = K.ImageEncoder().encode(img.tobytes(), img.shape)
+    await asyncio.sleep(random.random())
+    img_decoded = K.ImageDecoder().decode(bytes(frame))
+    await asyncio.sleep(random.random())
+    img = np.from_dlpack(img_decoded)
+    await asyncio.sleep(random.random())
+    print(f"End: {i}")
+    return img.mean()
+
+
+@pytest.mark.asyncio
+async def test_receive_stream_task():
+    tasks = [asyncio.create_task(encode_frame(i)) for i in range(3)]
+    results = await asyncio.gather(*tasks)
+    mean = sum(results) / len(results)
+    print(mean)