Fix bbox scaling estimation for Large Scale Jitter (#5446)

* Fix bbox scaling estimation for Large Scale Jitter * Fix linter
pytorch · Feb 20, 2022 · 034c222 · 034c222
1 parent 7bb5e41
commit 034c222
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/references/detection/transforms.py b/references/detection/transforms.py
@@ -317,14 +317,17 @@ def forward(
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
+        orig_width, orig_height = F.get_image_size(image)
+
         r = self.scale_range[0] + torch.rand(1) * (self.scale_range[1] - self.scale_range[0])
         new_width = int(self.target_size[1] * r)
         new_height = int(self.target_size[0] * r)
 
         image = F.resize(image, [new_height, new_width], interpolation=self.interpolation)
 
         if target is not None:
-            target["boxes"] *= r
+            target["boxes"][:, 0::2] *= new_width / orig_width
+            target["boxes"][:, 1::2] *= new_height / orig_height
             if "masks" in target:
                 target["masks"] = F.resize(
                     target["masks"], [new_height, new_width], interpolation=InterpolationMode.NEAREST