Fix a bug and document RandomZoomOut (#289)

Similar to pytorch/vision#5278
zhiqwang · Jan 25, 2022 · 3d943e7 · 3d943e7
1 parent 51f9d41
commit 3d943e7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/yolort/data/transforms.py b/yolort/data/transforms.py
@@ -203,7 +203,7 @@ def forward(
             elif image.ndimension() == 2:
                 image = image.unsqueeze(0)
 
-        if torch.rand(1) < self.p:
+        if torch.rand(1) >= self.p:
             return image, target
 
         orig_w, orig_h = get_image_size(image)
@@ -225,6 +225,7 @@ def forward(
 
         image = F.pad(image, [left, top, right, bottom], fill=fill)
         if isinstance(image, Tensor):
+            # PyTorch's pad supports only integers on fill. So we need to overwrite the colour
             v = torch.tensor(self.fill, device=image.device, dtype=image.dtype).view(-1, 1, 1)
             image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h) :, :] = image[
                 ..., :, (left + orig_w) :