diff --git a/python/src/bindings/keypoint.rs b/python/src/bindings/keypoint.rs index 73a6c369..77c0e509 100644 --- a/python/src/bindings/keypoint.rs +++ b/python/src/bindings/keypoint.rs @@ -147,6 +147,21 @@ impl From for u8 { } } +impl PyKeypoint { + /// Rust-side accessor for the inner Keypoint (used by pose.rs). + /// Not exposed to Python — Python users go through the + /// #[pymethods] getters above. + pub(crate) fn inner(&self) -> &Keypoint { + &self.inner + } + + /// Rust-side constructor from a core Keypoint (used by pose.rs + /// when re-wrapping outputs of PersonPose methods). + pub(crate) fn from_rust(k: Keypoint) -> Self { + Self { inner: k } + } +} + // ─── Keypoint ──────────────────────────────────────────────────────── /// Single skeletal joint with COCO type, 2D-or-3D position, and a diff --git a/python/src/bindings/pose.rs b/python/src/bindings/pose.rs new file mode 100644 index 00000000..70e8a01f --- /dev/null +++ b/python/src/bindings/pose.rs @@ -0,0 +1,376 @@ +//! 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; 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 { + 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 { + 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> { + // 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 { + 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 { + 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 { + 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, + confidence: f32, + latency_ms: f32, + model_version: String, + ) -> PyResult { + let conf = Confidence::new(confidence).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(e.to_string()) + })?; + let rust_persons: Vec = + 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 { + 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 { + 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 { + HashMap::new() +} + +pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/python/src/lib.rs b/python/src/lib.rs index 5ffcaa93..ec3e7313 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -18,6 +18,7 @@ use pyo3::prelude::*; mod bindings { pub mod keypoint; + pub mod pose; } /// Version of the bound Rust core. Surfaced to Python as @@ -36,6 +37,7 @@ fn build_features() -> Vec<&'static str> { let mut feats: Vec<&'static str> = Vec::new(); feats.push("p1-scaffold"); feats.push("p2-keypoint-bindings"); // Keypoint + KeypointType + feats.push("p2-pose-bindings"); // BoundingBox + PersonPose + PoseEstimate feats } @@ -67,5 +69,7 @@ fn wifi_densepose_native(m: &Bound<'_, PyModule>) -> PyResult<()> { // P2 — Keypoint + KeypointType bindings. bindings::keypoint::register(m)?; + // P2 — BoundingBox + PersonPose + PoseEstimate bindings. + bindings::pose::register(m)?; Ok(()) } diff --git a/python/tests/test_pose.py b/python/tests/test_pose.py new file mode 100644 index 00000000..60867aa9 --- /dev/null +++ b/python/tests/test_pose.py @@ -0,0 +1,248 @@ +"""ADR-117 P2 tests — BoundingBox + PersonPose + PoseEstimate bindings. + +Run with: cd python && .venv/Scripts/python -m pytest tests/test_pose.py -v +""" + +from __future__ import annotations + +import pytest + +from wifi_densepose import ( + BoundingBox, + Keypoint, + KeypointType, + PersonPose, + PoseEstimate, +) + + +# ─── BoundingBox ───────────────────────────────────────────────────── + + +def test_bounding_box_construct() -> None: + bb = BoundingBox(0.1, 0.2, 0.5, 0.7) + assert bb.x_min == pytest.approx(0.1) + assert bb.y_min == pytest.approx(0.2) + assert bb.x_max == pytest.approx(0.5) + assert bb.y_max == pytest.approx(0.7) + + +def test_bounding_box_dimensions() -> None: + bb = BoundingBox(0.0, 0.0, 4.0, 3.0) + assert bb.width == pytest.approx(4.0) + assert bb.height == pytest.approx(3.0) + assert bb.area == pytest.approx(12.0) + assert bb.center == pytest.approx((2.0, 1.5)) + + +def test_bounding_box_from_center() -> None: + bb = BoundingBox.from_center(2.0, 3.0, 4.0, 6.0) + assert bb.x_min == pytest.approx(0.0) + assert bb.y_min == pytest.approx(0.0) + assert bb.x_max == pytest.approx(4.0) + assert bb.y_max == pytest.approx(6.0) + + +def test_bounding_box_iou_no_overlap() -> None: + a = BoundingBox(0.0, 0.0, 1.0, 1.0) + b = BoundingBox(2.0, 2.0, 3.0, 3.0) + assert a.iou(b) == pytest.approx(0.0) + + +def test_bounding_box_iou_full_overlap() -> None: + a = BoundingBox(0.0, 0.0, 1.0, 1.0) + b = BoundingBox(0.0, 0.0, 1.0, 1.0) + assert a.iou(b) == pytest.approx(1.0) + + +def test_bounding_box_iou_partial() -> None: + a = BoundingBox(0.0, 0.0, 10.0, 10.0) + b = BoundingBox(5.0, 5.0, 15.0, 15.0) + # intersection 25, union 175 → 1/7 + assert a.iou(b) == pytest.approx(25.0 / 175.0) + + +def test_bounding_box_eq() -> None: + assert BoundingBox(1, 2, 3, 4) == BoundingBox(1, 2, 3, 4) + assert BoundingBox(1, 2, 3, 4) != BoundingBox(1, 2, 3, 5) + + +def test_bounding_box_repr() -> None: + bb = BoundingBox(0.1, 0.2, 0.5, 0.7) + assert "BoundingBox" in repr(bb) + assert "x_min=0.1" in repr(bb) + + +# ─── PersonPose ────────────────────────────────────────────────────── + + +def test_person_pose_empty() -> None: + p = PersonPose() + assert p.id is None + assert p.visible_keypoint_count == 0 + assert p.bounding_box is None + assert p.confidence == 0.0 + + +def test_person_pose_set_get_keypoint() -> None: + p = PersonPose() + kp = Keypoint(KeypointType.Nose, 0.5, 0.3, 0.95) + p.set_keypoint(kp) + got = p.get_keypoint(KeypointType.Nose) + assert got is not None + assert got.x == pytest.approx(0.5) + assert got.confidence == pytest.approx(0.95) + + +def test_person_pose_get_missing_returns_none() -> None: + p = PersonPose() + p.set_keypoint(Keypoint(KeypointType.Nose, 0.5, 0.3, 0.95)) + assert p.get_keypoint(KeypointType.LeftWrist) is None + + +def test_person_pose_visible_count() -> None: + p = PersonPose() + p.set_keypoint(Keypoint(KeypointType.Nose, 0.0, 0.0, 0.9)) # visible + p.set_keypoint(Keypoint(KeypointType.LeftEar, 0.0, 0.0, 0.2)) # invisible + p.set_keypoint(Keypoint(KeypointType.RightEar, 0.0, 0.0, 0.8)) # visible + assert p.visible_keypoint_count == 2 + + +def test_person_pose_visible_keypoints_list() -> None: + p = PersonPose() + p.set_keypoint(Keypoint(KeypointType.Nose, 0.0, 0.0, 0.9)) + p.set_keypoint(Keypoint(KeypointType.LeftEar, 0.0, 0.0, 0.2)) + vis = p.visible_keypoints() + assert len(vis) == 1 + assert vis[0].keypoint_type == KeypointType.Nose + + +def test_person_pose_keypoints_dict_excludes_missing() -> None: + p = PersonPose() + p.set_keypoint(Keypoint(KeypointType.Nose, 0.0, 0.0, 0.9)) + p.set_keypoint(Keypoint(KeypointType.LeftWrist, 0.5, 0.5, 0.6)) + d = p.keypoints() + assert KeypointType.Nose in d + assert KeypointType.LeftWrist in d + assert KeypointType.RightAnkle not in d + assert len(d) == 2 + + +def test_person_pose_set_id() -> None: + p = PersonPose() + p.set_id(7) + assert p.id == 7 + + +def test_person_pose_set_bounding_box() -> None: + p = PersonPose() + bb = BoundingBox(0.1, 0.1, 0.5, 0.9) + p.set_bounding_box(bb) + assert p.bounding_box == bb + + +def test_person_pose_compute_bbox_returns_none_when_empty() -> None: + p = PersonPose() + assert p.compute_bounding_box() is None + + +def test_person_pose_compute_bbox_from_keypoints() -> None: + p = PersonPose() + p.set_keypoint(Keypoint(KeypointType.Nose, 0.0, 0.0, 0.95)) + p.set_keypoint(Keypoint(KeypointType.RightAnkle, 1.0, 2.0, 0.95)) + bb = p.compute_bounding_box() + assert bb is not None + # bbox should span both keypoints + assert bb.x_min <= 0.0 + assert bb.y_min <= 0.0 + assert bb.x_max >= 1.0 + assert bb.y_max >= 2.0 + # also stored + assert p.bounding_box is not None + + +def test_person_pose_set_confidence_validation() -> None: + p = PersonPose() + p.set_confidence(0.85) + assert p.confidence == pytest.approx(0.85) + with pytest.raises(ValueError): + p.set_confidence(1.5) + + +def test_person_pose_repr() -> None: + p = PersonPose() + p.set_id(3) + p.set_keypoint(Keypoint(KeypointType.Nose, 0.0, 0.0, 0.9)) + r = repr(p) + assert "PersonPose" in r + assert "id=Some(3)" in r or "id=3" in r + + +# ─── PoseEstimate ──────────────────────────────────────────────────── + + +def test_pose_estimate_construct_empty() -> None: + e = PoseEstimate([], 0.5, 1.0, "test-v0") + assert e.person_count == 0 + assert not e.has_detections + assert e.confidence == pytest.approx(0.5) + assert e.latency_ms == pytest.approx(1.0) + assert e.model_version == "test-v0" + + +def test_pose_estimate_construct_with_persons() -> None: + p1 = PersonPose() + p1.set_id(1) + p1.set_confidence(0.8) + p2 = PersonPose() + p2.set_id(2) + p2.set_confidence(0.9) + e = PoseEstimate([p1, p2], 0.85, 5.2, "v0.7.0") + assert e.person_count == 2 + assert e.has_detections + assert e.confidence == pytest.approx(0.85) + + +def test_pose_estimate_highest_confidence_person() -> None: + p1 = PersonPose() + p1.set_confidence(0.5) + p2 = PersonPose() + p2.set_confidence(0.95) + p3 = PersonPose() + p3.set_confidence(0.7) + e = PoseEstimate([p1, p2, p3], 0.85, 5.2, "v0.7.0") + best = e.highest_confidence_person() + assert best is not None + assert best.confidence == pytest.approx(0.95) + + +def test_pose_estimate_highest_confidence_returns_none_when_empty() -> None: + e = PoseEstimate([], 0.5, 1.0, "test") + assert e.highest_confidence_person() is None + + +def test_pose_estimate_metadata_strings_nonempty() -> None: + e = PoseEstimate([], 0.5, 1.0, "test") + assert isinstance(e.id, str) + assert isinstance(e.timestamp, str) + assert e.id # non-empty + assert e.timestamp # non-empty + + +def test_pose_estimate_confidence_validation() -> None: + with pytest.raises(ValueError): + PoseEstimate([], 1.5, 0.0, "test") + + +def test_pose_estimate_repr_contains_counts() -> None: + e = PoseEstimate([], 0.5, 2.3, "v0.7.0") + r = repr(e) + assert "PoseEstimate" in r + assert "v0.7.0" in r + + +def test_build_features_marks_p2_complete() -> None: + import wifi_densepose + + assert "p2-keypoint-bindings" in wifi_densepose.__build_features__ + assert "p2-pose-bindings" in wifi_densepose.__build_features__ diff --git a/python/wifi_densepose/__init__.py b/python/wifi_densepose/__init__.py index d71162c4..d24e26f7 100644 --- a/python/wifi_densepose/__init__.py +++ b/python/wifi_densepose/__init__.py @@ -41,6 +41,9 @@ from wifi_densepose import _native # import Keypoint, KeypointType` — never reach into `_native`. Keypoint = _native.Keypoint KeypointType = _native.KeypointType +BoundingBox = _native.BoundingBox +PersonPose = _native.PersonPose +PoseEstimate = _native.PoseEstimate __rust_version__: str = _native.__rust_version__ @@ -74,4 +77,7 @@ __all__ = [ # P2 — core types "Keypoint", "KeypointType", + "BoundingBox", + "PersonPose", + "PoseEstimate", ]