377 lines
12 KiB
Rust
377 lines
12 KiB
Rust
//! ADR-117 P2 — PyO3 bindings for `BoundingBox`, `PersonPose`,
|
|
//! `PoseEstimate`.
|
|
//!
|
|
//! Design notes:
|
|
//!
|
|
//! 1. **`PersonPose` exposes the 17-keypoint array as a Python dict
|
|
//! keyed by `KeypointType`**, not as a fixed-length list with
|
|
//! `None` slots. Pythonistas don't want to know that the underlying
|
|
//! storage is `[Option<Keypoint>; 17]`.
|
|
//!
|
|
//! 2. **`PoseEstimate` metadata `id` and `timestamp` are exposed as
|
|
//! strings** (UUID + RFC 3339) rather than as bound types. Users
|
|
//! in notebooks rarely need to compare UUIDs structurally; strings
|
|
//! are good enough and don't require binding `FrameId` /
|
|
//! `Timestamp` as separate classes.
|
|
//!
|
|
//! 3. **`PersonPose` is mutable** via `set_keypoint` / `set_bbox` /
|
|
//! `set_id` — it's a builder-style type users construct
|
|
//! incrementally. Hence NOT `#[pyclass(frozen)]`.
|
|
//!
|
|
//! 4. **`PoseEstimate` is frozen** — once constructed, the list of
|
|
//! persons + the metadata don't change.
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use pyo3::prelude::*;
|
|
use pyo3::types::PyDict;
|
|
|
|
use wifi_densepose_core::{
|
|
BoundingBox, Confidence, KeypointType, PersonPose, PoseEstimate,
|
|
};
|
|
|
|
use super::keypoint::{PyKeypoint, PyKeypointType};
|
|
|
|
// ─── BoundingBox ─────────────────────────────────────────────────────
|
|
|
|
/// Axis-aligned bounding box around a detected person.
|
|
///
|
|
/// Python:
|
|
/// ```python
|
|
/// from wifi_densepose import BoundingBox
|
|
///
|
|
/// bb = BoundingBox(0.1, 0.2, 0.5, 0.7)
|
|
/// print(bb.width, bb.height, bb.area, bb.center)
|
|
/// bb2 = BoundingBox.from_center(0.3, 0.45, 0.4, 0.5)
|
|
/// print(bb.iou(bb2))
|
|
/// ```
|
|
#[pyclass(frozen, name = "BoundingBox")]
|
|
#[derive(Clone)]
|
|
pub struct PyBoundingBox {
|
|
inner: BoundingBox,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl PyBoundingBox {
|
|
#[new]
|
|
fn new(x_min: f32, y_min: f32, x_max: f32, y_max: f32) -> Self {
|
|
Self { inner: BoundingBox::new(x_min, y_min, x_max, y_max) }
|
|
}
|
|
|
|
/// Construct from center point + width + height.
|
|
#[staticmethod]
|
|
fn from_center(cx: f32, cy: f32, width: f32, height: f32) -> Self {
|
|
Self { inner: BoundingBox::from_center(cx, cy, width, height) }
|
|
}
|
|
|
|
#[getter]
|
|
fn x_min(&self) -> f32 { self.inner.x_min }
|
|
#[getter]
|
|
fn y_min(&self) -> f32 { self.inner.y_min }
|
|
#[getter]
|
|
fn x_max(&self) -> f32 { self.inner.x_max }
|
|
#[getter]
|
|
fn y_max(&self) -> f32 { self.inner.y_max }
|
|
#[getter]
|
|
fn width(&self) -> f32 { self.inner.width() }
|
|
#[getter]
|
|
fn height(&self) -> f32 { self.inner.height() }
|
|
#[getter]
|
|
fn area(&self) -> f32 { self.inner.area() }
|
|
#[getter]
|
|
fn center(&self) -> (f32, f32) { self.inner.center() }
|
|
|
|
/// Intersection over Union (IoU) with another box. Range [0.0, 1.0].
|
|
fn iou(&self, other: &PyBoundingBox) -> f32 {
|
|
self.inner.iou(&other.inner)
|
|
}
|
|
|
|
fn __repr__(&self) -> String {
|
|
format!(
|
|
"BoundingBox(x_min={}, y_min={}, x_max={}, y_max={})",
|
|
self.inner.x_min, self.inner.y_min, self.inner.x_max, self.inner.y_max,
|
|
)
|
|
}
|
|
|
|
fn __eq__(&self, other: &PyBoundingBox) -> bool {
|
|
self.inner == other.inner
|
|
}
|
|
}
|
|
|
|
impl PyBoundingBox {
|
|
pub(crate) fn from_rust(bb: BoundingBox) -> Self {
|
|
Self { inner: bb }
|
|
}
|
|
}
|
|
|
|
// ─── PersonPose ──────────────────────────────────────────────────────
|
|
|
|
/// A single detected person with optional ID, up to 17 keypoints, and
|
|
/// an optional bounding box.
|
|
///
|
|
/// Python:
|
|
/// ```python
|
|
/// from wifi_densepose import PersonPose, Keypoint, KeypointType, BoundingBox
|
|
///
|
|
/// pose = PersonPose()
|
|
/// pose.set_keypoint(Keypoint(KeypointType.Nose, 0.5, 0.3, 0.95))
|
|
/// pose.set_keypoint(Keypoint(KeypointType.LeftShoulder, 0.4, 0.5, 0.92))
|
|
/// pose.set_id(7)
|
|
/// print(pose.visible_keypoint_count) # 2
|
|
/// print(pose.get_keypoint(KeypointType.Nose).confidence) # 0.95
|
|
/// print(pose.compute_bounding_box()) # auto-derived from visible kp
|
|
/// ```
|
|
#[pyclass(name = "PersonPose")]
|
|
#[derive(Clone)]
|
|
pub struct PyPersonPose {
|
|
inner: PersonPose,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl PyPersonPose {
|
|
/// Construct an empty person pose. Set keypoints + bbox + id with
|
|
/// the dedicated methods.
|
|
#[new]
|
|
fn new() -> Self {
|
|
Self { inner: PersonPose::new() }
|
|
}
|
|
|
|
/// Per-person track ID. None until set.
|
|
#[getter]
|
|
fn id(&self) -> Option<u32> {
|
|
self.inner.id
|
|
}
|
|
|
|
fn set_id(&mut self, id: u32) {
|
|
self.inner.id = Some(id);
|
|
}
|
|
|
|
/// Set or replace a keypoint. The keypoint's type determines its
|
|
/// slot in the internal 17-element array.
|
|
fn set_keypoint(&mut self, keypoint: PyKeypoint) {
|
|
self.inner.set_keypoint(*keypoint.inner());
|
|
}
|
|
|
|
/// Get a keypoint by type, or None if not set.
|
|
fn get_keypoint(&self, keypoint_type: PyKeypointType) -> Option<PyKeypoint> {
|
|
let kp = self.inner.get_keypoint(keypoint_type.as_rust())?;
|
|
// Re-wrap the inner Rust Keypoint for Python.
|
|
Some(PyKeypoint::from_rust(*kp))
|
|
}
|
|
|
|
/// All keypoints as a dict keyed by KeypointType. Missing
|
|
/// keypoints are omitted (NOT included with None values).
|
|
fn keypoints<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
|
// PyO3 0.22 — PyDict::new_bound returns a Bound, the legacy
|
|
// PyDict::new (returning &PyDict) was removed in 0.21.
|
|
let dict = PyDict::new_bound(py);
|
|
for (i, kp_opt) in self.inner.keypoints.iter().enumerate() {
|
|
if let Some(kp) = kp_opt {
|
|
let kpt = match KeypointType::all().get(i) {
|
|
Some(t) => *t,
|
|
None => continue,
|
|
};
|
|
// Convert through IntoPy to satisfy ToPyObject bound
|
|
// for dict.set_item — #[pyclass] types impl IntoPy but
|
|
// not ToPyObject directly in PyO3 0.22.
|
|
use pyo3::IntoPy;
|
|
let k_obj: PyObject = PyKeypointType::from_rust(kpt).into_py(py);
|
|
let v_obj: PyObject = PyKeypoint::from_rust(*kp).into_py(py);
|
|
dict.set_item(k_obj, v_obj)?;
|
|
}
|
|
}
|
|
Ok(dict)
|
|
}
|
|
|
|
/// Number of visible keypoints (confidence >= 0.5).
|
|
#[getter]
|
|
fn visible_keypoint_count(&self) -> usize {
|
|
self.inner.visible_keypoint_count()
|
|
}
|
|
|
|
/// List of visible keypoints (subset of the dict from
|
|
/// `keypoints()`).
|
|
fn visible_keypoints(&self) -> Vec<PyKeypoint> {
|
|
self.inner
|
|
.visible_keypoints()
|
|
.into_iter()
|
|
.map(|k| PyKeypoint::from_rust(*k))
|
|
.collect()
|
|
}
|
|
|
|
/// Bounding box, if previously set or computed.
|
|
#[getter]
|
|
fn bounding_box(&self) -> Option<PyBoundingBox> {
|
|
self.inner.bounding_box.map(PyBoundingBox::from_rust)
|
|
}
|
|
|
|
fn set_bounding_box(&mut self, bb: PyBoundingBox) {
|
|
self.inner.bounding_box = Some(bb.inner);
|
|
}
|
|
|
|
/// Auto-compute bounding box from visible keypoints, set it
|
|
/// internally, and return it. Returns None if no keypoints visible.
|
|
fn compute_bounding_box(&mut self) -> Option<PyBoundingBox> {
|
|
let bb = self.inner.compute_bounding_box()?;
|
|
self.inner.bounding_box = Some(bb);
|
|
Some(PyBoundingBox::from_rust(bb))
|
|
}
|
|
|
|
/// Overall confidence in [0.0, 1.0].
|
|
#[getter]
|
|
fn confidence(&self) -> f32 {
|
|
self.inner.confidence.value()
|
|
}
|
|
|
|
fn set_confidence(&mut self, c: f32) -> PyResult<()> {
|
|
self.inner.confidence = Confidence::new(c).map_err(|e| {
|
|
pyo3::exceptions::PyValueError::new_err(e.to_string())
|
|
})?;
|
|
Ok(())
|
|
}
|
|
|
|
fn __repr__(&self) -> String {
|
|
format!(
|
|
"PersonPose(id={:?}, visible_keypoints={}, confidence={:.4})",
|
|
self.inner.id,
|
|
self.inner.visible_keypoint_count(),
|
|
self.inner.confidence.value(),
|
|
)
|
|
}
|
|
}
|
|
|
|
impl PyPersonPose {
|
|
pub(crate) fn from_rust(pose: PersonPose) -> Self {
|
|
Self { inner: pose }
|
|
}
|
|
}
|
|
|
|
// ─── PoseEstimate ────────────────────────────────────────────────────
|
|
|
|
/// Top-level result of a pose-estimation pass — a list of detected
|
|
/// persons plus metadata about the inference run.
|
|
///
|
|
/// Python:
|
|
/// ```python
|
|
/// from wifi_densepose import PoseEstimate, PersonPose
|
|
///
|
|
/// est = PoseEstimate([pose1, pose2], confidence=0.87, latency_ms=8.4,
|
|
/// model_version="v0.1.0")
|
|
/// print(est.person_count, est.has_detections)
|
|
/// best = est.highest_confidence_person()
|
|
/// ```
|
|
#[pyclass(frozen, name = "PoseEstimate")]
|
|
pub struct PyPoseEstimate {
|
|
inner: PoseEstimate,
|
|
}
|
|
|
|
#[pymethods]
|
|
impl PyPoseEstimate {
|
|
/// Construct a pose estimate from a list of detected persons,
|
|
/// an overall confidence, inference latency, and model version
|
|
/// string.
|
|
#[new]
|
|
fn new(
|
|
persons: Vec<PyPersonPose>,
|
|
confidence: f32,
|
|
latency_ms: f32,
|
|
model_version: String,
|
|
) -> PyResult<Self> {
|
|
let conf = Confidence::new(confidence).map_err(|e| {
|
|
pyo3::exceptions::PyValueError::new_err(e.to_string())
|
|
})?;
|
|
let rust_persons: Vec<PersonPose> =
|
|
persons.into_iter().map(|p| p.inner).collect();
|
|
Ok(Self {
|
|
inner: PoseEstimate::new(
|
|
Vec::new(),
|
|
rust_persons,
|
|
conf,
|
|
latency_ms,
|
|
model_version,
|
|
),
|
|
})
|
|
}
|
|
|
|
/// Unique frame identifier as a UUID string.
|
|
#[getter]
|
|
fn id(&self) -> String {
|
|
format!("{:?}", self.inner.id)
|
|
.trim_start_matches("FrameId(")
|
|
.trim_end_matches(')')
|
|
.to_string()
|
|
}
|
|
|
|
/// Frame timestamp as an RFC 3339 / ISO 8601 string in UTC.
|
|
#[getter]
|
|
fn timestamp(&self) -> String {
|
|
// Timestamp's Debug impl is usable; for a fully spec-compliant
|
|
// ISO format, a future refactor binds chrono. P2 string-form
|
|
// is "good enough" for diagnostics.
|
|
format!("{:?}", self.inner.timestamp)
|
|
}
|
|
|
|
#[getter]
|
|
fn persons(&self) -> Vec<PyPersonPose> {
|
|
self.inner.persons.iter().cloned().map(PyPersonPose::from_rust).collect()
|
|
}
|
|
|
|
#[getter]
|
|
fn confidence(&self) -> f32 {
|
|
self.inner.confidence.value()
|
|
}
|
|
|
|
#[getter]
|
|
fn latency_ms(&self) -> f32 {
|
|
self.inner.latency_ms
|
|
}
|
|
|
|
#[getter]
|
|
fn model_version(&self) -> &str {
|
|
&self.inner.model_version
|
|
}
|
|
|
|
#[getter]
|
|
fn person_count(&self) -> usize {
|
|
self.inner.person_count()
|
|
}
|
|
|
|
#[getter]
|
|
fn has_detections(&self) -> bool {
|
|
self.inner.has_detections()
|
|
}
|
|
|
|
/// Get the person with the highest individual confidence, or None
|
|
/// if no persons detected.
|
|
fn highest_confidence_person(&self) -> Option<PyPersonPose> {
|
|
self.inner
|
|
.highest_confidence_person()
|
|
.cloned()
|
|
.map(PyPersonPose::from_rust)
|
|
}
|
|
|
|
fn __repr__(&self) -> String {
|
|
format!(
|
|
"PoseEstimate(persons={}, confidence={:.4}, latency_ms={:.2}, model_version={:?})",
|
|
self.inner.person_count(),
|
|
self.inner.confidence.value(),
|
|
self.inner.latency_ms,
|
|
self.inner.model_version,
|
|
)
|
|
}
|
|
}
|
|
|
|
/// Suppress unused-import warnings for HashMap (held for future
|
|
/// keypoint-map helpers in P3).
|
|
#[allow(dead_code)]
|
|
fn _hashmap_kept_for_future_use() -> HashMap<u8, u8> {
|
|
HashMap::new()
|
|
}
|
|
|
|
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|
m.add_class::<PyBoundingBox>()?;
|
|
m.add_class::<PyPersonPose>()?;
|
|
m.add_class::<PyPoseEstimate>()?;
|
|
Ok(())
|
|
}
|