Fix transforms bug.

PaddlePaddle · May 11, 2021 · 245f122 · 245f122
1 parent ecfac2a
commit 245f122
Showing 1 changed file with 13 additions and 11 deletions.
diff --git a/paddleseg/transforms/transforms.py b/paddleseg/transforms/transforms.py
@@ -718,13 +718,14 @@ def __call__(self, im, label=None):
                 flags=cv2.INTER_LINEAR,
                 borderMode=cv2.BORDER_CONSTANT,
                 borderValue=self.im_padding_value)
-            label = cv2.warpAffine(
-                label,
-                r,
-                dsize=dsize,
-                flags=cv2.INTER_NEAREST,
-                borderMode=cv2.BORDER_CONSTANT,
-                borderValue=self.label_padding_value)
+            if label is not None:
+                label = cv2.warpAffine(
+                    label,
+                    r,
+                    dsize=dsize,
+                    flags=cv2.INTER_NEAREST,
+                    borderMode=cv2.BORDER_CONSTANT,
+                    borderValue=self.label_padding_value)
 
         if label is None:
             return (im, )
@@ -778,13 +779,14 @@ def __call__(self, im, label=None):
                     w1 = np.random.randint(0, img_width - dw)
 
                     im = im[h1:(h1 + dh), w1:(w1 + dw), :]
-                    label = label[h1:(h1 + dh), w1:(w1 + dw)]
                     im = cv2.resize(
                         im, (img_width, img_height),
                         interpolation=cv2.INTER_LINEAR)
-                    label = cv2.resize(
-                        label, (img_width, img_height),
-                        interpolation=cv2.INTER_NEAREST)
+                    if label is not None:
+                        label = label[h1:(h1 + dh), w1:(w1 + dw)]
+                        label = cv2.resize(
+                            label, (img_width, img_height),
+                            interpolation=cv2.INTER_NEAREST)
                     break
         if label is None:
             return (im, )