huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Automatic mask generation (#779)

Browse Source 
* A few more contiguous fixes for cuda.

* Mask generation.

* Generic bbox.

* Generate all the masks.
This commit is contained in:
Laurent Mazare
2023-09-08 19:11:34 +01:00
committed by GitHub
parent 158ff3c609
commit 0906acab91
7 changed files with 125 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-examples/src/object_detection.rs
View File

@ -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());