feat: added some post processing for retinaface
This commit is contained in:
@@ -1,13 +1,14 @@
|
||||
use crate::errors::*;
|
||||
use bounding_box::Aabb2;
|
||||
use error_stack::ResultExt;
|
||||
use mnn_bridge::ndarray::*;
|
||||
use nalgebra::Point2;
|
||||
use nalgebra::{Point2, Vector2};
|
||||
use ndarray_resize::NdFir;
|
||||
use std::path::Path;
|
||||
|
||||
pub struct FaceDetectionConfig {
|
||||
min_sizes: Vec<Point2<u32>>,
|
||||
steps: Vec<u32>,
|
||||
min_sizes: Vec<Vector2<usize>>,
|
||||
steps: Vec<usize>,
|
||||
variance: Vec<f32>,
|
||||
}
|
||||
pub struct FaceDetection {
|
||||
@@ -20,6 +21,37 @@ pub struct FaceDetectionModelOutput {
|
||||
pub landmark: ndarray::Array3<f32>,
|
||||
}
|
||||
|
||||
impl FaceDetectionModelOutput {
|
||||
pub fn postprocess(self, config: FaceDetectionConfig) -> Result<Vec<Aabb2<f32>>> {
|
||||
// for k, step in enumerate(cfg['steps']):
|
||||
// feature_size = 640 // step
|
||||
// for i in range(feature_size):
|
||||
// for j in range(feature_size):
|
||||
// for min_size in cfg['min_sizes'][k]:
|
||||
// cx = (j + 0.5) * step / 640
|
||||
// cy = (i + 0.5) * step / 640
|
||||
// s_kx = s_ky = min_size / 640
|
||||
// anchors.append([cx, cy, s_kx, s_ky])
|
||||
let mut anchors = Vec::new();
|
||||
config.steps.iter().enumerate().for_each(|(k, step)| {
|
||||
let feature_size = 640 / step;
|
||||
for i in 0..feature_size {
|
||||
for j in 0..feature_size {
|
||||
for min_size in &config.min_sizes[k] {
|
||||
let cx = (j as f32 + 0.5) * *step as f32 / 640.0;
|
||||
let cy = (i as f32 + 0.5) * *step as f32 / 640.0;
|
||||
let s_kx = *min_size as f32 / 640.0;
|
||||
let s_ky = *min_size as f32 / 640.0;
|
||||
anchors.push([cx, cy, s_kx, s_ky]);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Vec::new())
|
||||
}
|
||||
}
|
||||
|
||||
impl FaceDetectionModelOutput {
|
||||
pub fn print(&self, limit: usize) {
|
||||
tracing::info!("Detected {} faces", self.bbox.shape()[1]);
|
||||
|
||||
Reference in New Issue
Block a user