70 lines
2.7 KiB
Rust
70 lines
2.7 KiB
Rust
mod cli;
|
|
mod errors;
|
|
use errors::*;
|
|
use ndarray_image::*;
|
|
const RETINAFACE_MODEL: &[u8] = include_bytes!("../models/retinaface.mnn");
|
|
pub fn main() -> Result<()> {
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter("trace")
|
|
.with_thread_ids(true)
|
|
.with_thread_names(true)
|
|
.with_target(false)
|
|
.init();
|
|
let args = <cli::Cli as clap::Parser>::parse();
|
|
match args.cmd {
|
|
cli::SubCommand::Detect(detect) => {
|
|
use detector::facedet;
|
|
let model = facedet::retinaface::FaceDetection::new_from_bytes(RETINAFACE_MODEL)
|
|
.change_context(errors::Error)
|
|
.attach_printable("Failed to create face detection model")?;
|
|
let image = image::open(detect.image).change_context(Error)?;
|
|
let image = image.into_rgb8();
|
|
let mut array = image
|
|
.into_ndarray()
|
|
.change_context(errors::Error)
|
|
.attach_printable("Failed to convert image to ndarray")?;
|
|
let output = model
|
|
.detect_faces(array.clone())
|
|
.change_context(errors::Error)
|
|
.attach_printable("Failed to detect faces")?;
|
|
// output.print(20);
|
|
let aabbs = output
|
|
.postprocess(Default::default())
|
|
.change_context(errors::Error)
|
|
.attach_printable("Failed to attach context")?;
|
|
for bbox in aabbs {
|
|
println!("Detected face: {:?}", bbox);
|
|
use bounding_box::draw::*;
|
|
let bbox = bbox
|
|
.denormalize(nalgebra::SVector::<f32, 2>::new(
|
|
array.shape()[1] as f32,
|
|
array.shape()[0] as f32,
|
|
))
|
|
.cast()
|
|
.ok_or(errors::Error)
|
|
.attach_printable("Failed to cast f32 to usize")?;
|
|
array.draw(bbox, color::palette::css::GREEN_YELLOW.to_rgba8(), 20);
|
|
break;
|
|
}
|
|
let v = array.view();
|
|
if let Some(output) = detect.output {
|
|
let image: image::RgbImage = v
|
|
.to_image()
|
|
.change_context(errors::Error)
|
|
.attach_printable("Failed to convert ndarray to image")?;
|
|
image
|
|
.save(output)
|
|
.change_context(errors::Error)
|
|
.attach_printable("Failed to save output image")?;
|
|
}
|
|
}
|
|
cli::SubCommand::List(list) => {
|
|
println!("List: {:?}", list);
|
|
}
|
|
cli::SubCommand::Completions { shell } => {
|
|
cli::Cli::completions(shell);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|