scale detection bboxes to match original dimensions

mixaal · Sep 12, 2024 · 1433a59 · 1433a59
1 parent 4d3df0a
commit 1433a59
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 23 deletions.
diff --git a/src/image.rs b/src/image.rs
@@ -56,15 +56,12 @@ impl Image {
 
     pub fn draw_rectangle(&mut self, bboxes: &Vec<BBox>) {
         let image = &mut self.image;
-        let (_, initial_h, initial_w) = image.size3().expect("can't get image size");
-        let w_ratio = initial_w as f64 / self.width as f64;
-        let h_ratio = initial_h as f64 / self.height as f64;
 
         for bbox in bboxes.iter() {
-            let xmin = ((bbox.xmin * w_ratio) as i64).clamp(0, initial_w - 1);
-            let ymin = ((bbox.ymin * h_ratio) as i64).clamp(0, initial_h - 1);
-            let xmax = ((bbox.xmax * w_ratio) as i64).clamp(0, initial_w - 1);
-            let ymax = ((bbox.ymax * h_ratio) as i64).clamp(0, initial_h - 1);
+            let xmin = bbox.xmin as i64;
+            let ymin = bbox.ymin as i64;
+            let xmax = bbox.xmax as i64;
+            let ymax = bbox.ymax as i64;
             Self::draw_line(image, xmin, xmax, ymin, ymax.min(ymin + 2));
             Self::draw_line(image, xmin, xmax, ymin.max(ymax - 2), ymax);
             Self::draw_line(image, xmin, xmax.min(xmin + 2), ymin, ymax);

diff --git a/src/lib.rs b/src/lib.rs
@@ -276,10 +276,10 @@ mod test {
         assert_eq!(3, detection.len());
         bbox_eq(
             BBox {
-                xmin: 548.6856384277344,
-                ymin: 311.5385515507371,
-                xmax: 578.6725158691406,
-                ymax: 383.35467598219884,
+                xmin: 1028.787,
+                ymin: 387.95,
+                xmax: 1085.012,
+                ymax: 477.393,
                 conf: 0.5158080458641052,
                 cls: 0,
                 name: "person",
@@ -289,10 +289,10 @@ mod test {
 
         bbox_eq(
             BBox {
-                xmin: 475.7906494140625,
-                ymin: 282.6662423051286,
-                xmax: 520.7713012695313,
-                ymax: 367.16750926325284,
+                xmin: 892.104,
+                ymin: 352.016,
+                xmax: 976.45,
+                ymax: 457.242,
                 conf: 0.4701675474643707,
                 cls: 0,
                 name: "person",
@@ -301,10 +301,10 @@ mod test {
         );
         bbox_eq(
             BBox {
-                xmin: 13.92529296875,
-                ymin: 111.91644110972788,
-                xmax: 607.813720703125,
-                ymax: 531.2043928770977,
+                xmin: 26.115,
+                ymin: 139.375,
+                xmax: 1139.667,
+                ymax: 661.516,
                 conf: 0.9144105911254883,
                 cls: 5,
                 name: "bus",

diff --git a/src/utils.rs b/src/utils.rs
@@ -177,6 +177,11 @@ impl DetectionTools {
         let prediction = prediction.transpose(1, 0);
         let (anchors, classes_no) = prediction.size2().unwrap();
 
+        let initial_w = image_dim.2 as f64;
+        let initial_h = image_dim.1 as f64;
+        let w_ratio = initial_w / scaled_image_dim.2 as f64;
+        let h_ratio = initial_h / scaled_image_dim.1 as f64;
+
         let nclasses = (classes_no - 4) as usize;
         // println!("classes_no={classes_no}, anchors={anchors}");
 
@@ -218,11 +223,16 @@ impl DetectionTools {
                     let x = cx + dx;
                     let y = cy + dy;
 
+                    let xmin = ((x - w / 2.) * w_ratio).clamp(0.0, initial_w - 1.0);
+                    let ymin = ((y - h / 2.) * h_ratio).clamp(0.0, initial_h - 1.0);
+                    let xmax = ((x + w / 2.) * w_ratio).clamp(0.0, initial_w - 1.0);
+                    let ymax = ((y + h / 2.) * h_ratio).clamp(0.0, initial_h - 1.0);
+
                     let bbox = BBox {
-                        xmin: x - w / 2.,
-                        ymin: y - h / 2.,
-                        xmax: x + w / 2.,
-                        ymax: y + h / 2.,
+                        xmin,
+                        ymin,
+                        xmax,
+                        ymax,
                         conf: confidence,
                         cls: class_index,
                         name: crate::classes::DETECT_CLASSES[class_index],