fix pascal aug test error!

LikeLy-Journey · Mar 24, 2020 · 27845a1 · 27845a1
1 parent 940015d
commit 27845a1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/segmentron/data/dataloader/pascal_aug.py b/segmentron/data/dataloader/pascal_aug.py
@@ -73,6 +73,8 @@ def __getitem__(self, index):
             img, target = self._sync_transform(img, target)
         elif self.mode == 'val':
             img, target = self._val_sync_transform(img, target)
+        elif self.mode == 'testval':
+            img, target = self._val_sync_transform(img, target)
         else:
             raise RuntimeError('unknown mode for dataloader: {}'.format(self.mode))
         # general resize, normalize and toTensor

diff --git a/segmentron/data/dataloader/pascal_voc.py b/segmentron/data/dataloader/pascal_voc.py
@@ -1,6 +1,7 @@
 """Pascal VOC Semantic Segmentation Dataset."""
 import os
 import torch
+import logging
 import numpy as np
 
 from PIL import Image
@@ -84,6 +85,7 @@ def __getitem__(self, index):
             img, mask = self._val_sync_transform(img, mask)
         else:
             assert self.mode == 'testval'
+            logging.warn("Use mode of testval, you should set batch size=1")
             img, mask = self._img_transform(img), self._mask_transform(mask)
         # general resize, normalize and toTensor
         if self.transform is not None: