diff --git a/pyproject.toml b/pyproject.toml index 56bdfcaf..5a9cd607 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ classifiers = [ [project.optional-dependencies] dev = [ "pytest", + "pytest-asyncio", "kornia==0.6.3", "numpy==1.23.0", "jax[cpu]", diff --git a/src/dlpack_py.rs b/src/dlpack_py.rs index 6f451e2e..d62044f6 100644 --- a/src/dlpack_py.rs +++ b/src/dlpack_py.rs @@ -92,14 +92,13 @@ pub fn cvtensor_to_dlpack(x: &cv::Tensor, py: Python) -> PyResult { 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) } diff --git a/src/tensor.rs b/src/tensor.rs index c7dec03b..079f6241 100644 --- a/src/tensor.rs +++ b/src/tensor.rs @@ -40,14 +40,13 @@ pub mod cv { } #[pyo3(name = "__dlpack__")] - pub fn to_dlpack_py(&self, py: Python) -> PyResult { + pub fn dlpack_py(&self, py: Python) -> PyResult { 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, diff --git a/test/test_io.py b/test/test_io.py index c6d5aa20..83ba2d38 100644 --- a/test/test_io.py +++ b/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)