wifi_densepose_nn/
lib.rs

1//! # WiFi-DensePose Neural Network Crate
2//!
3//! This crate provides neural network inference capabilities for the WiFi-DensePose
4//! pose estimation system. It supports multiple backends including ONNX Runtime,
5//! tch-rs (PyTorch), and Candle for flexible deployment.
6//!
7//! ## Features
8//!
9//! - **DensePose Head**: Body part segmentation and UV coordinate regression
10//! - **Modality Translator**: CSI to visual feature space translation
11//! - **Multi-Backend Support**: ONNX, PyTorch (tch), and Candle backends
12//! - **Inference Optimization**: Batching, GPU acceleration, and model caching
13//!
14//! ## Example
15//!
16//! ```rust,ignore
17//! use wifi_densepose_nn::{InferenceEngine, DensePoseConfig, OnnxBackend};
18//!
19//! // Create inference engine with ONNX backend
20//! let config = DensePoseConfig::default();
21//! let backend = OnnxBackend::from_file("model.onnx")?;
22//! let engine = InferenceEngine::new(backend, config)?;
23//!
24//! // Run inference
25//! let input = ndarray::Array4::zeros((1, 256, 64, 64));
26//! let output = engine.infer(&input)?;
27//! ```
28
29#![warn(missing_docs)]
30#![warn(rustdoc::missing_doc_code_examples)]
31#![deny(unsafe_code)]
32
33pub mod densepose;
34pub mod error;
35pub mod inference;
36#[cfg(feature = "onnx")]
37pub mod onnx;
38pub mod tensor;
39pub mod translator;
40
41// Re-exports for convenience
42pub use densepose::{DensePoseConfig, DensePoseHead, DensePoseOutput};
43pub use error::{NnError, NnResult};
44pub use inference::{Backend, InferenceEngine, InferenceOptions};
45#[cfg(feature = "onnx")]
46pub use onnx::{OnnxBackend, OnnxSession};
47pub use tensor::{Tensor, TensorShape};
48pub use translator::{ModalityTranslator, TranslatorConfig, TranslatorOutput};
49
50/// Prelude module for convenient imports
51pub mod prelude {
52    pub use crate::densepose::{DensePoseConfig, DensePoseHead, DensePoseOutput};
53    pub use crate::error::{NnError, NnResult};
54    pub use crate::inference::{Backend, InferenceEngine, InferenceOptions};
55    #[cfg(feature = "onnx")]
56    pub use crate::onnx::{OnnxBackend, OnnxSession};
57    pub use crate::tensor::{Tensor, TensorShape};
58    pub use crate::translator::{ModalityTranslator, TranslatorConfig, TranslatorOutput};
59}
60
61/// Version information
62pub const VERSION: &str = env!("CARGO_PKG_VERSION");
63
64/// Number of body parts in DensePose model (standard configuration)
65pub const NUM_BODY_PARTS: usize = 24;
66
67/// Number of UV coordinates (U and V)
68pub const NUM_UV_COORDINATES: usize = 2;
69
70/// Default hidden channel sizes for networks
71pub const DEFAULT_HIDDEN_CHANNELS: &[usize] = &[256, 128, 64];
wifi_densepose_nn/lib.rs

wifi_densepose_nn/
lib.rs
