add dlparkimg example

Cargo.toml

@@ -17,7 +17,7 @@ half = { version = "2.3", optional = true }
 pyo3 = { version = "0.19", optional = true }
 
 [workspace]
-members = ["examples/from_numpy", "examples/with_pyo3"]
+members = ["examples/from_numpy", "examples/with_pyo3", "examples/dlparkimg"]
 
 [features]
 default = []

examples/dlparkimg/.github/workflows/CI.yml

@@ -0,0 +1,120 @@
+# This file is autogenerated by maturin v1.1.0
+# To update, run
+#
+# maturin generate-ci github
+#
+name: CI
+
+on:
+ push:
+ branches:
+ - main
+ - master
+ tags:
+ - '*'
+ pull_request:
+ workflow_dispatch:
+
+permissions:
+ contents: read
+
+jobs:
+ linux:
+ runs-on: ubuntu-latest
+ strategy:
+ matrix:
+ target: [x86_64, x86, aarch64, armv7, s390x, ppc64le]
+ steps:
+ - uses: actions/checkout@v3
+ - uses: actions/setup-python@v4
+ with:
+ python-version: '3.10'
+ - name: Build wheels
+ uses: PyO3/maturin-action@v1
+ with:
+ target: ${{ matrix.target }}
+ args: --release --out dist --find-interpreter
+ sccache: 'true'
+ manylinux: auto
+ - name: Upload wheels
+ uses: actions/upload-artifact@v3
+ with:
+ name: wheels
+ path: dist
+
+ windows:
+ runs-on: windows-latest
+ strategy:
+ matrix:
+ target: [x64, x86]
+ steps:
+ - uses: actions/checkout@v3
+ - uses: actions/setup-python@v4
+ with:
+ python-version: '3.10'
+ architecture: ${{ matrix.target }}
+ - name: Build wheels
+ uses: PyO3/maturin-action@v1
+ with:
+ target: ${{ matrix.target }}
+ args: --release --out dist --find-interpreter
+ sccache: 'true'
+ - name: Upload wheels
+ uses: actions/upload-artifact@v3
+ with:
+ name: wheels
+ path: dist
+
+ macos:
+ runs-on: macos-latest
+ strategy:
+ matrix:
+ target: [x86_64, aarch64]
+ steps:
+ - uses: actions/checkout@v3
+ - uses: actions/setup-python@v4
+ with:
+ python-version: '3.10'
+ - name: Build wheels
+ uses: PyO3/maturin-action@v1
+ with:
+ target: ${{ matrix.target }}
+ args: --release --out dist --find-interpreter
+ sccache: 'true'
+ - name: Upload wheels
+ uses: actions/upload-artifact@v3
+ with:
+ name: wheels
+ path: dist
+
+ sdist:
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v3
+ - name: Build sdist
+ uses: PyO3/maturin-action@v1
+ with:
+ command: sdist
+ args: --out dist
+ - name: Upload sdist
+ uses: actions/upload-artifact@v3
+ with:
+ name: wheels
+ path: dist
+
+ release:
+ name: Release
+ runs-on: ubuntu-latest
+ if: "startsWith(github.ref, 'refs/tags/')"
+ needs: [linux, windows, macos, sdist]
+ steps:
+ - uses: actions/download-artifact@v3
+ with:
+ name: wheels
+ - name: Publish to PyPI
+ uses: PyO3/maturin-action@v1
+ env:
+ MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
+ with:
+ command: upload
+ args: --skip-existing *

examples/dlparkimg/.gitignore

@@ -0,0 +1,72 @@
+/target
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+.pytest_cache/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+.venv/
+env/
+bin/
+build/
+develop-eggs/
+dist/
+eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+include/
+man/
+venv/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+pip-selfcheck.json
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.cache
+nosetests.xml
+coverage.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# Rope
+.ropeproject
+
+# Django stuff:
+*.log
+*.pot
+
+.DS_Store
+
+# Sphinx documentation
+docs/_build/
+
+# PyCharm
+.idea/
+
+# VSCode
+.vscode/
+
+# Pyenv
+.python-version

examples/dlparkimg/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "dlparkimg"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[lib]
+name = "dlparkimg"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = "0.19.0"
+dlpark = { path = "../../", features = ["pyo3"] }
+image = "0.24.7"

examples/dlparkimg/README.md

@@ -0,0 +1,16 @@
+# Example for [dlpark](https://github.com/SunDoge/dlpark)
+
+This is an example of how to transfer `image::RgbImage` to Python and how to transfer `torch.Tensor` to Rust.
+
+## Usage
+
+```shell
+pip install maturin
+maturin develop
+pip install torch matplotlib
+python main.py
+```
+
+| Input RGB Image | Output BGR Image |
+| ------------------- | ---------------- |
+| ![candy](candy.jpg) | ![bgr](bgr.jpg) |

examples/dlparkimg/main.py

@@ -0,0 +1,12 @@
+import dlparkimg
+from torch.utils.dlpack import to_dlpack, from_dlpack
+# import matplotlib.pyplot as plt
+
+tensor = from_dlpack(dlparkimg.read_image("candy.jpg"))
+
+print(tensor.shape)
+# plt.imshow(tensor.numpy())
+# plt.show()
+
+bgr_img = tensor[..., [2, 1, 0]]
+dlparkimg.write_image('bgr.jpg', to_dlpack(bgr_img))

examples/dlparkimg/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["maturin>=1.1,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "dlparkimg"
+requires-python = ">=3.7"
+classifiers = [
+ "Programming Language :: Rust",
+ "Programming Language :: Python :: Implementation :: CPython",
+ "Programming Language :: Python :: Implementation :: PyPy",
+]
+
+
+[tool.maturin]
+features = ["pyo3/extension-module"]

examples/dlparkimg/src/lib.rs

@@ -0,0 +1,60 @@
+use std::ffi::c_void;
+
+use dlpark::prelude::*;
+use image::Rgb;
+use pyo3::prelude::*;
+
+struct PyRgbImage(image::RgbImage);
+
+impl ToTensor for PyRgbImage {
+ fn data_ptr(&self) -> *mut std::ffi::c_void {
+ self.0.as_ptr() as *const c_void as *mut c_void
+ }
+
+ fn byte_offset(&self) -> u64 {
+ 0
+ }
+
+ fn device(&self) -> Device {
+ Device::CPU
+ }
+
+ fn dtype(&self) -> DataType {
+ DataType::U8
+ }
+
+ fn shape_and_strides(&self) -> ShapeAndStrides {
+ ShapeAndStrides::new_contiguous_with_strides(
+ &[self.0.height(), self.0.width(), 3].map(|x| x as i64),
+ )
+ }
+}
+
+#[pyfunction]
+fn read_image(filename: &str) -> ManagerCtx<PyRgbImage> {
+ let img = image::open(filename).unwrap();
+ let rgb_img = img.to_rgb8();
+ ManagerCtx::new(PyRgbImage(rgb_img))
+}
+
+#[pyfunction]
+fn write_image(filename: &str, tensor: ManagedTensor) {
+ let buf = tensor.as_slice::<u8>();
+
+ let rgb_img = image::ImageBuffer::<Rgb<u8>, _>::from_raw(
+ tensor.shape()[1] as u32,
+ tensor.shape()[0] as u32,
+ buf,
+ )
+ .unwrap();
+
+ rgb_img.save(filename).unwrap();
+}
+
+/// A Python module implemented in Rust.
+#[pymodule]
+fn dlparkimg(_py: Python, m: &PyModule) -> PyResult<()> {
+ m.add_function(wrap_pyfunction!(read_image, m)?)?;
+ m.add_function(wrap_pyfunction!(write_image, m)?)?;
+ Ok(())
+}

src/shape_and_strides.rs

@@ -31,7 +31,11 @@ impl ShapeAndStrides {
  {
  let shape: Vec<&i64> = shape.into_iter().collect();
  let strides: Vec<&i64> = strides.into_iter().collect();
- assert_eq!(shape.len(), strides.len());
+ assert_eq!(
+ shape.len(),
+ strides.len(),
+ "shape and strides should have same length"
+ );
  let mut buf: Vec<i64> = Vec::with_capacity(shape.len() + strides.len());
  buf.extend(shape);
  buf.extend(strides);
@@ -56,7 +60,11 @@ impl ShapeAndStrides {
 
  pub fn new_borrowed(shape: &[i64], strides: Option<&[i64]>) -> Self {
  if let Some(ref strides) = strides {
- assert_eq!(shape.len(), strides.len());
+ assert_eq!(
+ shape.len(),
+ strides.len(),
+ "shape and strides should have same length"
+ );
  }
  let len = shape.len();
  let shape = NonNull::from(&shape[0]);

src/tensor.rs

@@ -32,6 +32,11 @@ impl ManagedTensor {
 
  /// Access inner data as 1d array.
  pub fn as_slice<A>(&self) -> &[A] {
+ assert_eq!(
+ std::mem::size_of::<A>(),
+ self.dtype().size(),
+ "dtype and A size mismatch"
+ );
  unsafe {
  let ptr = self.data_ptr().add(self.byte_offset() as usize);
  std::slice::from_raw_parts(ptr.cast(), self.num_elements())
@@ -140,4 +145,11 @@ mod tests {
  let strides = make_contiguous_strides(&shape);
  assert_eq!(&strides, &[6, 3, 1]);
  }
+
+ #[test]
+ fn test_as_slice() {
+ let v: Vec<f32> = (0..10).map(|x| x as f32).collect();
+ let tensor = ManagedTensor::from_dlpack(v.clone().into_dlpack());
+ assert_eq!(tensor.as_slice::<f32>(), &v[..]);
+ }
 }

src/utils.rs

@@ -8,7 +8,11 @@ pub fn make_contiguous_strides(shape: &[i64]) -> Vec<i64> {
 }
 
 pub fn is_contiguous(shape: &[i64], strides: &[i64]) -> bool {
- assert_eq!(shape.len(), strides.len());
+ assert_eq!(
+ shape.len(),
+ strides.len(),
+ "shape and strides should have same length"
+ );
  let mut expected_stride = 1;
  for (&dim, &stride) in shape.iter().rev().zip(strides.iter().rev()) {
  if stride != expected_stride {