cutout bugfix

agrichron · Nov 12, 2020 · 29dbf5e · 29dbf5e
1 parent 7a94d5b
commit 29dbf5e
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/example/utility/cutout.py b/example/utility/cutout.py
@@ -8,12 +8,13 @@ def __init__(self, size=16, p=0.5):
         self.p = p
 
     def __call__(self, image):
-        if torch.rand([1]).item() > self.p: return image
+        if torch.rand([1]).item() > self.p:
+            return image
 
-        left = torch.randint(-self.half_size, image.shape[0] - self.half_size, [1]).item()
-        top = torch.randint(-self.half_size, image.shape[1] - self.half_size, [1]).item()
-        right = min(image.shape[0], left + self.size)
-        bottom = min(image.shape[1], top + self.size)
+        left = torch.randint(-self.half_size, image.size(1) - self.half_size, [1]).item()
+        top = torch.randint(-self.half_size, image.size(2) - self.half_size, [1]).item()
+        right = min(image.size(1), left + self.size)
+        bottom = min(image.size(2), top + self.size)
 
-        image[max(0,left):right, max(0,top):bottom, :] = 0
+        image[:, max(0, left): right, max(0, top): bottom] = 0
         return image