use crate::scrfd::SCRFD;
use ort::session::Session;
use std::error::Error;
#[cfg(feature = "async")]
use crate::scrfd_async::SCRFDAsync;
pub struct SCRFDBuilder {
session: Session,
input_size: Option<(i32, i32)>,
conf_thres: Option<f32>,
iou_thres: Option<f32>,
relative_output: bool,
}
impl SCRFDBuilder {
pub fn new(session: Session) -> Self {
SCRFDBuilder {
session,
input_size: Some((640, 640)),
conf_thres: Some(0.25),
iou_thres: Some(0.4),
relative_output: true,
}
}
pub fn set_input_size(mut self, size: (i32, i32)) -> Self {
self.input_size = Some(size);
self
}
pub fn set_conf_thres(mut self, thres: f32) -> Self {
self.conf_thres = Some(thres);
self
}
pub fn set_iou_thres(mut self, thres: f32) -> Self {
self.iou_thres = Some(thres);
self
}
pub fn set_relative_output(mut self, relative: bool) -> Self {
self.relative_output = relative;
self
}
pub fn build(self) -> Result<SCRFD, Box<dyn Error>> {
let input_size = self.input_size.ok_or("Input size not set")?;
let conf_thres = self.conf_thres.ok_or("Confidence threshold not set")?;
let iou_thres = self.iou_thres.ok_or("IoU threshold not set")?;
SCRFD::new(
self.session,
input_size,
conf_thres,
iou_thres,
self.relative_output,
)
}
#[cfg(feature = "async")]
pub fn build_async(self) -> Result<SCRFDAsync, Box<dyn Error>> {
let input_size = self.input_size.ok_or("Input size not set")?;
let conf_thres = self.conf_thres.ok_or("Confidence threshold not set")?;
let iou_thres = self.iou_thres.ok_or("IoU threshold not set")?;
SCRFDAsync::new(
self.session,
input_size,
conf_thres,
iou_thres,
self.relative_output,
)
}
}