diff options
Diffstat (limited to 'candle-examples/src/object_detection.rs')
| -rw-r--r-- | candle-examples/src/object_detection.rs | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/candle-examples/src/object_detection.rs b/candle-examples/src/object_detection.rs index c7c60136..ce579316 100644 --- a/candle-examples/src/object_detection.rs +++ b/candle-examples/src/object_detection.rs @@ -1,12 +1,12 @@ /// A bounding box around an object. #[derive(Debug, Clone)] -pub struct Bbox { +pub struct Bbox<D> { pub xmin: f32, pub ymin: f32, pub xmax: f32, pub ymax: f32, pub confidence: f32, - pub keypoints: Vec<KeyPoint>, + pub data: D, } #[derive(Debug, Clone, Copy, PartialEq)] @@ -17,7 +17,7 @@ pub struct KeyPoint { } /// Intersection over union of two bounding boxes. -pub fn iou(b1: &Bbox, b2: &Bbox) -> f32 { +pub fn iou<D>(b1: &Bbox<D>, b2: &Bbox<D>) -> f32 { let b1_area = (b1.xmax - b1.xmin + 1.) * (b1.ymax - b1.ymin + 1.); let b2_area = (b2.xmax - b2.xmin + 1.) * (b2.ymax - b2.ymin + 1.); let i_xmin = b1.xmin.max(b2.xmin); @@ -28,7 +28,7 @@ pub fn iou(b1: &Bbox, b2: &Bbox) -> f32 { i_area / (b1_area + b2_area - i_area) } -pub fn non_maximum_suppression(bboxes: &mut [Vec<Bbox>], threshold: f32) { +pub fn non_maximum_suppression<D>(bboxes: &mut [Vec<Bbox<D>>], threshold: f32) { // Perform non-maximum suppression. for bboxes_for_class in bboxes.iter_mut() { bboxes_for_class.sort_by(|b1, b2| b2.confidence.partial_cmp(&b1.confidence).unwrap()); |