use nv_core::{BBox, DetectionId, TypedMetadata};
#[derive(Clone, Debug)]
pub struct Detection {
pub id: DetectionId,
pub class_id: u32,
pub confidence: f32,
pub bbox: BBox,
pub embedding: Option<Vec<f32>>,
pub metadata: TypedMetadata,
}
impl Detection {
#[must_use]
pub fn builder(
id: DetectionId,
class_id: u32,
confidence: f32,
bbox: BBox,
) -> DetectionBuilder {
DetectionBuilder {
id,
class_id,
confidence,
bbox,
embedding: None,
metadata: TypedMetadata::new(),
}
}
}
pub struct DetectionBuilder {
id: DetectionId,
class_id: u32,
confidence: f32,
bbox: BBox,
embedding: Option<Vec<f32>>,
metadata: TypedMetadata,
}
impl DetectionBuilder {
#[must_use]
pub fn embedding(mut self, embedding: Vec<f32>) -> Self {
self.embedding = Some(embedding);
self
}
#[must_use]
pub fn meta<T: Clone + Send + Sync + 'static>(mut self, val: T) -> Self {
self.metadata.insert(val);
self
}
#[must_use]
pub fn build(self) -> Detection {
Detection {
id: self.id,
class_id: self.class_id,
confidence: clamp_unit(self.confidence),
bbox: self.bbox,
embedding: self.embedding,
metadata: self.metadata,
}
}
}
fn clamp_unit(v: f32) -> f32 {
if v.is_finite() {
v.clamp(0.0, 1.0)
} else {
0.0
}
}
#[derive(Clone, Debug, Default)]
pub struct DetectionSet {
pub detections: Vec<Detection>,
}
impl DetectionSet {
#[must_use]
pub fn empty() -> Self {
Self {
detections: Vec::new(),
}
}
#[must_use]
pub fn len(&self) -> usize {
self.detections.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.detections.is_empty()
}
}
impl From<Vec<Detection>> for DetectionSet {
fn from(detections: Vec<Detection>) -> Self {
Self { detections }
}
}
#[cfg(test)]
mod tests {
use super::*;
use nv_core::BBox;
#[test]
fn builder_required_fields_only() {
let det =
Detection::builder(DetectionId::new(1), 0, 0.95, BBox::new(0.1, 0.2, 0.3, 0.4)).build();
assert_eq!(det.id, DetectionId::new(1));
assert_eq!(det.class_id, 0);
assert!((det.confidence - 0.95).abs() < f32::EPSILON);
assert!(det.embedding.is_none());
assert!(det.metadata.is_empty());
}
#[test]
fn builder_with_embedding_and_meta() {
#[derive(Clone, Debug, PartialEq)]
struct Extra(u32);
let det = Detection::builder(DetectionId::new(2), 5, 0.8, BBox::new(0.0, 0.0, 1.0, 1.0))
.embedding(vec![0.1, 0.2, 0.3])
.meta(Extra(42))
.build();
assert_eq!(det.embedding.as_ref().unwrap().len(), 3);
assert_eq!(det.metadata.get::<Extra>(), Some(&Extra(42)));
}
#[test]
fn detection_set_from_vec() {
let dets = vec![
Detection::builder(DetectionId::new(1), 0, 0.9, BBox::new(0.0, 0.0, 0.5, 0.5)).build(),
Detection::builder(DetectionId::new(2), 1, 0.7, BBox::new(0.5, 0.5, 1.0, 1.0)).build(),
];
let set: DetectionSet = dets.into();
assert_eq!(set.len(), 2);
assert!(!set.is_empty());
}
}