From 9e7abe09452e3ab596be485e7571307fb426d5d4 Mon Sep 17 00:00:00 2001 From: Jules Youngberg Date: Mon, 13 Jun 2022 20:57:57 -0700 Subject: [PATCH] break submodules out into separate files --- src/face_mesh.rs | 37 +++++++ src/hands.rs | 65 ++++++++++++ src/lib.rs | 238 ++++---------------------------------------- src/pose.rs | 74 ++++++++++++++ src/segmentation.rs | 28 ++++++ 5 files changed, 223 insertions(+), 219 deletions(-) create mode 100644 src/face_mesh.rs create mode 100644 src/hands.rs create mode 100644 src/pose.rs create mode 100644 src/segmentation.rs diff --git a/src/face_mesh.rs b/src/face_mesh.rs new file mode 100644 index 0000000..f6046df --- /dev/null +++ b/src/face_mesh.rs @@ -0,0 +1,37 @@ +//! Face detection utilities. +use super::*; + +pub struct FaceMeshDetector { + graph: Detector, +} + +impl FaceMeshDetector { + pub fn new() -> Self { + let graph = Detector::new( + FACE_GRAPH_TYPE, + include_str!("graphs/face_mesh_desktop_live.pbtxt"), + "multi_face_landmarks", + ); + + Self { graph } + } + + /// Processes the input frame, returns a face mesh if detected. + pub fn process(&mut self, input: &Mat) -> Option { + let landmarks = self.graph.process(input); + + if landmarks.is_empty() { + return None; + } + + let mut face_mesh = FaceMesh::default(); + face_mesh.data.copy_from_slice(landmarks); + Some(face_mesh) + } +} + +impl Default for FaceMeshDetector { + fn default() -> Self { + Self::new() + } +} diff --git a/src/hands.rs b/src/hands.rs new file mode 100644 index 0000000..7906e43 --- /dev/null +++ b/src/hands.rs @@ -0,0 +1,65 @@ +//! Hand detection utilities. +use super::*; + +/// Hand landmark indices. +pub enum HandLandmark { + WRIST = 0, + THUMB_CMC = 1, + THUMB_MCP = 2, + THUMB_IP = 3, + THUMB_TIP = 4, + INDEX_FINGER_MCP = 5, + INDEX_FINGER_PIP = 6, + INDEX_FINGER_DIP = 7, + INDEX_FINGER_TIP = 8, + MIDDLE_FINGER_MCP = 9, + MIDDLE_FINGER_PIP = 10, + MIDDLE_FINGER_DIP = 11, + MIDDLE_FINGER_TIP = 12, + RING_FINGER_MCP = 13, + RING_FINGER_PIP = 14, + RING_FINGER_DIP = 15, + RING_FINGER_TIP = 16, + PINKY_MCP = 17, + PINKY_PIP = 18, + PINKY_DIP = 19, + PINKY_TIP = 20, +} + +pub struct HandDetector { + graph: Detector, +} + +impl HandDetector { + pub fn new() -> Self { + let graph = Detector::new( + HANDS_GRAPH_TYPE, + include_str!("graphs/hand_tracking_desktop_live.pbtxt"), + "hand_landmarks", + ); + + Self { graph } + } + + /// Processes the input frame, returns a tuple of hands if detected. + pub fn process(&mut self, input: &Mat) -> Option<[Hand; 2]> { + let landmarks = self.graph.process(input); + + if landmarks.is_empty() { + return None; + } + + let mut lh = Hand::default(); + let mut rh = Hand::default(); + lh.data.copy_from_slice(&landmarks[0..21]); + rh.data.copy_from_slice(&landmarks[21..42]); + + Some([lh, rh]) + } +} + +impl Default for HandDetector { + fn default() -> Self { + Self::new() + } +} diff --git a/src/lib.rs b/src/lib.rs index f5d742a..e256ff9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,6 +15,10 @@ use opencv::prelude::*; use std::ffi::CString; mod bindings; +pub mod face_mesh; +pub mod hands; +pub mod pose; +pub mod segmentation; use bindings::*; @@ -151,15 +155,27 @@ impl Effect { /// Processes the input frame, returns a slice of landmarks if any are detected. pub fn process(&mut self, input: &Mat) -> Mat { let mut data = input.clone(); + let cols = data.cols(); + let rows = data.rows(); + let typ = data.typ(); let out_data = unsafe { mediagraph_Effect_Process( self.graph as *mut std::ffi::c_void, data.data_mut(), - data.cols(), - data.rows(), + cols, + rows, ) }; - unsafe { Mat::from_raw(out_data as *mut std::ffi::c_void) } + unsafe { + Mat::new_rows_cols_with_data( + rows, + cols, + typ, + out_data as *mut std::ffi::c_void, + opencv::core::Mat_AUTO_STEP, + ) + .unwrap() + } } } @@ -170,219 +186,3 @@ impl Drop for Effect { } } } - -pub mod pose { - //! Pose detection utilities. - use super::*; - - /// Pose landmark indices. - pub enum PoseLandmark { - NOSE = 0, - LEFT_EYE_INNER = 1, - LEFT_EYE = 2, - LEFT_EYE_OUTER = 3, - RIGHT_EYE_INNER = 4, - RIGHT_EYE = 5, - RIGHT_EYE_OUTER = 6, - LEFT_EAR = 7, - RIGHT_EAR = 8, - MOUTH_LEFT = 9, - MOUTH_RIGHT = 10, - LEFT_SHOULDER = 11, - RIGHT_SHOULDER = 12, - LEFT_ELBOW = 13, - RIGHT_ELBOW = 14, - LEFT_WRIST = 15, - RIGHT_WRIST = 16, - LEFT_PINKY = 17, - RIGHT_PINKY = 18, - LEFT_INDEX = 19, - RIGHT_INDEX = 20, - LEFT_THUMB = 21, - RIGHT_THUMB = 22, - LEFT_HIP = 23, - RIGHT_HIP = 24, - LEFT_KNEE = 25, - RIGHT_KNEE = 26, - LEFT_ANKLE = 27, - RIGHT_ANKLE = 28, - LEFT_HEEL = 29, - RIGHT_HEEL = 30, - LEFT_FOOT_INDEX = 31, - RIGHT_FOOT_INDEX = 32, - } - - pub struct PoseDetector { - graph: Detector, - } - - impl PoseDetector { - pub fn new() -> Self { - let graph = Detector::new( - POSE_GRAPH_TYPE, - include_str!("graphs/pose_tracking_cpu.pbtxt"), - "pose_landmarks", - ); - - Self { graph } - } - - /// Processes the input frame, returns a pose if detected. - pub fn process(&mut self, input: &Mat) -> Option { - let landmarks = self.graph.process(input); - - if landmarks.is_empty() { - return None; - } - - let mut pose = Pose::default(); - pose.data.copy_from_slice(landmarks); - Some(pose) - } - } - - impl Default for PoseDetector { - fn default() -> Self { - Self::new() - } - } -} - -pub mod face_mesh { - //! Face detection utilities. - use super::*; - - pub struct FaceMeshDetector { - graph: Detector, - } - - impl FaceMeshDetector { - pub fn new() -> Self { - let graph = Detector::new( - FACE_GRAPH_TYPE, - include_str!("graphs/face_mesh_desktop_live.pbtxt"), - "multi_face_landmarks", - ); - - Self { graph } - } - - /// Processes the input frame, returns a face mesh if detected. - pub fn process(&mut self, input: &Mat) -> Option { - let landmarks = self.graph.process(input); - - if landmarks.is_empty() { - return None; - } - - let mut face_mesh = FaceMesh::default(); - face_mesh.data.copy_from_slice(landmarks); - Some(face_mesh) - } - } - - impl Default for FaceMeshDetector { - fn default() -> Self { - Self::new() - } - } -} - -pub mod hands { - //! Hand detection utilities. - use super::*; - - /// Hand landmark indices. - pub enum HandLandmark { - WRIST = 0, - THUMB_CMC = 1, - THUMB_MCP = 2, - THUMB_IP = 3, - THUMB_TIP = 4, - INDEX_FINGER_MCP = 5, - INDEX_FINGER_PIP = 6, - INDEX_FINGER_DIP = 7, - INDEX_FINGER_TIP = 8, - MIDDLE_FINGER_MCP = 9, - MIDDLE_FINGER_PIP = 10, - MIDDLE_FINGER_DIP = 11, - MIDDLE_FINGER_TIP = 12, - RING_FINGER_MCP = 13, - RING_FINGER_PIP = 14, - RING_FINGER_DIP = 15, - RING_FINGER_TIP = 16, - PINKY_MCP = 17, - PINKY_PIP = 18, - PINKY_DIP = 19, - PINKY_TIP = 20, - } - - pub struct HandDetector { - graph: Detector, - } - - impl HandDetector { - pub fn new() -> Self { - let graph = Detector::new( - HANDS_GRAPH_TYPE, - include_str!("graphs/hand_tracking_desktop_live.pbtxt"), - "hand_landmarks", - ); - - Self { graph } - } - - /// Processes the input frame, returns a tuple of hands if detected. - pub fn process(&mut self, input: &Mat) -> Option<[Hand; 2]> { - let landmarks = self.graph.process(input); - - if landmarks.is_empty() { - return None; - } - - let mut lh = Hand::default(); - let mut rh = Hand::default(); - lh.data.copy_from_slice(&landmarks[0..21]); - rh.data.copy_from_slice(&landmarks[21..42]); - - Some([lh, rh]) - } - } - - impl Default for HandDetector { - fn default() -> Self { - Self::new() - } - } -} - -pub mod segmentation { - //! Selfie segmentation utilities. - use super::*; - - pub struct Segmentor { - graph: Effect, - } - - impl Segmentor { - pub fn new() -> Self { - let graph = Effect::new( - include_str!("graphs/selfie_segmentation_cpu.pbtxt"), - "output_video", - ); - - Self { graph } - } - - /// Processes the input frame, returns the output frame. - pub fn process(&mut self, input: &Mat) -> Mat { - self.graph.process(input) - } - } - - impl Default for Segmentor { - fn default() -> Self { - Self::new() - } - } -} diff --git a/src/pose.rs b/src/pose.rs new file mode 100644 index 0000000..ce48342 --- /dev/null +++ b/src/pose.rs @@ -0,0 +1,74 @@ +//! Pose detection utilities. +use super::*; + +/// Pose landmark indices. +pub enum PoseLandmark { + NOSE = 0, + LEFT_EYE_INNER = 1, + LEFT_EYE = 2, + LEFT_EYE_OUTER = 3, + RIGHT_EYE_INNER = 4, + RIGHT_EYE = 5, + RIGHT_EYE_OUTER = 6, + LEFT_EAR = 7, + RIGHT_EAR = 8, + MOUTH_LEFT = 9, + MOUTH_RIGHT = 10, + LEFT_SHOULDER = 11, + RIGHT_SHOULDER = 12, + LEFT_ELBOW = 13, + RIGHT_ELBOW = 14, + LEFT_WRIST = 15, + RIGHT_WRIST = 16, + LEFT_PINKY = 17, + RIGHT_PINKY = 18, + LEFT_INDEX = 19, + RIGHT_INDEX = 20, + LEFT_THUMB = 21, + RIGHT_THUMB = 22, + LEFT_HIP = 23, + RIGHT_HIP = 24, + LEFT_KNEE = 25, + RIGHT_KNEE = 26, + LEFT_ANKLE = 27, + RIGHT_ANKLE = 28, + LEFT_HEEL = 29, + RIGHT_HEEL = 30, + LEFT_FOOT_INDEX = 31, + RIGHT_FOOT_INDEX = 32, +} + +pub struct PoseDetector { + graph: Detector, +} + +impl PoseDetector { + pub fn new() -> Self { + let graph = Detector::new( + POSE_GRAPH_TYPE, + include_str!("graphs/pose_tracking_cpu.pbtxt"), + "pose_landmarks", + ); + + Self { graph } + } + + /// Processes the input frame, returns a pose if detected. + pub fn process(&mut self, input: &Mat) -> Option { + let landmarks = self.graph.process(input); + + if landmarks.is_empty() { + return None; + } + + let mut pose = Pose::default(); + pose.data.copy_from_slice(landmarks); + Some(pose) + } +} + +impl Default for PoseDetector { + fn default() -> Self { + Self::new() + } +} diff --git a/src/segmentation.rs b/src/segmentation.rs new file mode 100644 index 0000000..3b33374 --- /dev/null +++ b/src/segmentation.rs @@ -0,0 +1,28 @@ +//! Selfie segmentation utilities. +use super::*; + +pub struct Segmentor { + graph: Effect, +} + +impl Segmentor { + pub fn new() -> Self { + let graph = Effect::new( + include_str!("graphs/selfie_segmentation_cpu.pbtxt"), + "output_video", + ); + + Self { graph } + } + + /// Processes the input frame, returns the output frame. + pub fn process(&mut self, input: &Mat) -> Mat { + self.graph.process(input) + } +} + +impl Default for Segmentor { + fn default() -> Self { + Self::new() + } +}