kornia/kornia-rs

🦀 Low-level 3D Computer Vision library in Rust

What it solves

Kornia-rs 提供了一個高性能、低層級的電腦視覺函式庫，彌補了原始圖像數據與機器學習模型之間的差距。它消除了 Python 中主機到設備（host-to-device）複製與 GIL 瓶頸，從而實現能夠在 CPU 或 NVIDIA GPU 上無縫運行的實時圖像處理流水線。

How it works

該函式庫使用 Rust 編寫，具備記憶體與執行緒安全，並使用單一 API 自動根據數據所在位置分派操作。它利用 SIMD (NEON/AVX2) 進行 CPU 加速，並使用自定義 CUDA kernels 進行 GPU 加速。對於 ML 集成，它使用 DLPack 和 CUDA Array Interface 來將張量（tensors）移交給 PyTorch 和 TensorRT，而無需進行數據複製。一個專門的 Preprocessor 會將多個步驟——縮放、歸一化與格式轉換（例如，NV12 到 CHW）——融合進一個單一的 CUDA kernel，以最大化推理流水線的效率。

Who it’s for

開發建構實時電腦視覺應用程式的開發者、需要快速相機到模型（camera-to-model）流水線的機器人工程師，以及希望在擁有 Python 綁定（bindings）的同時獲得 Rust 級別性能的 ML 研究人員。

Highlights

  • Zero-copy interop: 透過 DLPack 直接移交給 PyTorch 和 TensorRT。
  • Unified API: 相同的算子（operators）可在 CPU 和 GPU 上運行，無需區分不同的類型。
  • Real-time optimized: 用於相機幀預處理的融合 CUDA kernels。
  • Broad format support: 原生支持廣泛的圖像格式，包括 AVIF、OpenEXR 和 PNG-16。
  • Python bindings: 透過 PyO3 提供對 free-threaded Python (no GIL) 的完整支持。