Merge pull request #716 from yuyu2172/cub-fail-test

Fix failing CUB tests

Author: Hakuyume

chainercv/datasets/cub/cub_label_dataset.py

@@ -2,6 +2,7 @@
 import os
 
 from chainercv.datasets.cub.cub_utils import CUBDatasetBase
+from chainercv import utils
 
 
 class CUBLabelDataset(CUBDatasetBase):
@@ -54,6 +55,7 @@ def __init__(self, data_dir='auto', return_bb=False,
                   d_label in open(image_class_labels_file)]
         self._labels = np.array(labels, dtype=np.int32)
 
+        self.add_getter('img', self._get_image)
         self.add_getter('label', self._get_label)
 
         keys = ('img', 'label')
@@ -63,5 +65,11 @@ def __init__(self, data_dir='auto', return_bb=False,
             keys += ('prob_map',)
         self.keys = keys
 
+    def _get_image(self, i):
+        img = utils.read_image(
+            os.path.join(self.data_dir, 'images', self.paths[i]),
+            color=True)
+        return img
+
     def _get_label(self, i):
         return self._labels[i]

chainercv/datasets/cub/cub_point_dataset.py

@@ -3,6 +3,7 @@
 import os
 
 from chainercv.datasets.cub.cub_utils import CUBDatasetBase
+from chainercv import utils
 
 
 class CUBPointDataset(CUBDatasetBase):
@@ -66,7 +67,8 @@ def __init__(self, data_dir='auto', return_bb=False,
             self._point_dict[id_].append(point)
             self._mask_dict[id_].append(mask)
 
-        self.add_getter(('point', 'mask'), self._get_annotations)
+        self.add_getter(('img', 'point', 'mask'),
+                        self._get_img_and_annotations)
 
         keys = ('img', 'point', 'mask')
         if return_bb:
@@ -75,7 +77,17 @@ def __init__(self, data_dir='auto', return_bb=False,
             keys += ('prob_map',)
         self.keys = keys
 
-    def _get_annotations(self, i):
+    def _get_img_and_annotations(self, i):
+        img = utils.read_image(
+            os.path.join(self.data_dir, 'images', self.paths[i]),
+            color=True)
+
         point = np.array(self._point_dict[i], dtype=np.float32)
         mask = np.array(self._mask_dict[i], dtype=np.bool)
-        return point, mask
+
+        _, H, W = img.shape
+        invalid = np.logical_or(
+            np.logical_or(point[:, 0] > H, point[:, 1] > W),
+            np.any(point < 0, axis=1))
+        mask[invalid] = False
+        return img, point, mask

chainercv/datasets/cub/cub_utils.py

@@ -77,52 +77,16 @@ def __init__(self, data_dir='auto', prob_map_dir='auto'):
             os.path.join(self.prob_map_dir, os.path.splitext(path)[0] + '.png')
             for path in self.paths]
 
-        self.add_getter('img', self.get_image)
-        self.add_getter('bb', self.get_bb)
-        self.add_getter('prob_map', self.get_prob_map)
+        self.add_getter('bb', self._get_bb)
+        self.add_getter('prob_map', self._get_prob_map)
 
     def __len__(self):
         return len(self.paths)
 
-    def get_image(self, i):
-        """Returns the i-th image.
-
-        Args:
-            i (int): The index of the example.
-
-        Returns:
-            An image.
-            The image is in CHW format and its color channel is ordered in
-            RGB.
-
-        """
-        img = utils.read_image(
-            os.path.join(self.data_dir, 'images', self.paths[i]),
-            color=True)
-        return img
-
-    def get_bb(self, i):
-        """Returns the bounding box of the i-th example.
-
-        Args:
-            i (int): The index of the example.
-
-        Returns:
-            A bounding box.
-
-        """
+    def _get_bb(self, i):
         return self.bbs[i]
 
-    def get_prob_map(self, i):
-        """Returns the probability map of the i-th example.
-
-        Args:
-            i (int): The index of the example.
-
-        Returns:
-            A probability map.
-
-        """
+    def _get_prob_map(self, i):
         prob_map = utils.read_image(self.prob_map_paths[i],
                                     dtype=np.uint8, color=False)
         prob_map = prob_map.astype(np.float32) / 255  # [0, 255] -> [0, 1]

tests/datasets_tests/cub_tests/test_cub_point_dataset.py

@@ -21,7 +21,7 @@ def setUp(self):
                                        return_prob_map=self.return_prob_map)
 
     @attr.slow
-    def test_camvid_dataset(self):
+    def test_cub_point_dataset(self):
         assert_is_point_dataset(
             self.dataset, n_point=15, n_example=10)