fix padding error

jfzhang95 · Dec 2, 2018 · acfff55 · acfff55
1 parent b7cbf5a
commit acfff55
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/dataloaders/custom_transforms.py b/dataloaders/custom_transforms.py
@@ -86,9 +86,10 @@ def __call__(self, sample):
 
 
 class RandomScaleCrop(object):
-    def __init__(self, base_size, crop_size):
+    def __init__(self, base_size, crop_size, fill=0):
         self.base_size = base_size
         self.crop_size = crop_size
+        self.fill = fill
 
     def __call__(self, sample):
         img = sample['image']
@@ -109,7 +110,7 @@ def __call__(self, sample):
             padh = self.crop_size - oh if oh < self.crop_size else 0
             padw = self.crop_size - ow if ow < self.crop_size else 0
             img = ImageOps.expand(img, border=(0, 0, padw, padh), fill=0)
-            mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=0)
+            mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=self.fill)
         # random crop crop_size
         w, h = img.size
         x1 = random.randint(0, w - self.crop_size)

diff --git a/dataloaders/datasets/cityscapes.py b/dataloaders/datasets/cityscapes.py
@@ -81,7 +81,7 @@ def recursive_glob(self, rootdir='.', suffix=''):
     def transform_tr(self, sample):
         composed_transforms = transforms.Compose([
             tr.RandomHorizontalFlip(),
-            tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size),
+            tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size, fill=255),
             tr.RandomGaussianBlur(),
             tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
             tr.ToTensor()])