Merge pull request #783 from azavea/lf/ignore-class

Handle "ignore" class for semantic segmentation

docs/changelog.rst

@@ -7,6 +7,11 @@ Raster Vision 0.10
 Raster Vision 0.10.0
 ~~~~~~~~~~~~~~~~~~~
 
+Features
+^^^^^^^^^^^^
+
+- Handle "ignore" class for semantic segmentation `#783 <https://github.com/azavea/raster-vision/pull/783>`__
+
 Bug Fixes
 ^^^^^^^^^
 

rastervision/core/class_map.py

@@ -40,7 +40,12 @@ def from_proto(msg):
 
 
 class ClassMap(object):
-    """A map from class_id to ClassItem."""
+    """A map from class_id to ClassItem.
+
+    The class ids should be integers >= 1. For semantic segmentation,
+    the class 0 is reserved for use as an "ignore" class, which denotes
+    pixels that should not be used for training the model or evaluating it.
+    """
 
     def __init__(self, class_items):
         """Construct a new ClassMap.

rastervision/evaluation/semantic_segmentation_evaluation.py

@@ -78,9 +78,10 @@ def compute(self, gt_labels, pred_labels):
 
             eval_items = []
             for class_id in self.class_map.get_keys():
-                eval_item = get_class_eval_item(gt_arr, pred_arr, class_id,
-                                                self.class_map)
-                eval_items.append(eval_item)
+                if class_id > 0:
+                    eval_item = get_class_eval_item(gt_arr, pred_arr, class_id,
+                                                    self.class_map)
+                    eval_items.append(eval_item)
 
             # Treat each window as if it was a small Scene.
             window_eval = SemanticSegmentationEvaluation(self.class_map)

rastervision/task/semantic_segmentation.py

@@ -7,6 +7,10 @@
 from rastervision.core.box import Box
 from rastervision.data.scene import Scene
 from rastervision.data.label import SemanticSegmentationLabels
+from rastervision.core.training_data import TrainingData
+
+TRAIN = 'train'
+VALIDATION = 'validation'
 
 log = logging.getLogger(__name__)
 
@@ -56,22 +60,31 @@ def get_train_windows(self, scene: Scene) -> List[Box]:
              A list of windows, list(Box)
 
         """
+        raster_source = scene.raster_source
+        extent = raster_source.get_extent()
+        label_source = scene.ground_truth_label_source
+        chip_size = self.config.chip_size
+        chip_options = self.config.chip_options
 
         def filter_windows(windows):
             if scene.aoi_polygons:
                 windows = Box.filter_by_aoi(windows, scene.aoi_polygons)
-            return windows
 
-        raster_source = scene.raster_source
-        extent = raster_source.get_extent()
-        label_store = scene.ground_truth_label_source
-        chip_size = self.config.chip_size
-
-        chip_options = self.config.chip_options
+            if 0 in self.config.class_map.get_keys():
+                filt_windows = []
+                for w in windows:
+                    label_arr = label_source.get_labels(w).get_label_arr(w)
+                    ignore_inds = label_arr.ravel() == 0
+                    if np.all(ignore_inds):
+                        pass
+                    else:
+                        filt_windows.append(w)
+                windows = filt_windows
+            return windows
 
         if chip_options.window_method == 'random_sample':
             return get_random_sample_train_windows(
-                label_store, chip_size, self.config.class_map, extent,
+                label_source, chip_size, self.config.class_map, extent,
                 chip_options, filter_windows)
         elif chip_options.window_method == 'sliding':
             stride = chip_options.stride
@@ -82,6 +95,63 @@ def filter_windows(windows):
             return list(
                 filter_windows((extent.get_windows(chip_size, stride))))
 
+    def make_chips(self, train_scenes, validation_scenes, augmentors, tmp_dir):
+        """Make training chips.
+
+        Convert Scenes with a ground_truth_label_store into training
+        chips in MLBackend-specific format, and write to URI specified in
+        options.
+
+        Args:
+            train_scenes: list of Scenes
+            validation_scenes: list of Scenes
+                (that is disjoint from train_scenes)
+            augmentors: Augmentors used to augment training data
+        """
+
+        def _process_scene(scene, type_, augment):
+            with scene.activate():
+                data = TrainingData()
+                log.info('Making {} chips for scene: {}'.format(
+                    type_, scene.id))
+                windows = self.get_train_windows(scene)
+                for window in windows:
+                    chip = scene.raster_source.get_chip(window)
+                    labels = self.get_train_labels(window, scene)
+
+                    # If chip has ignore labels, fill in those pixels with
+                    # nodata.
+                    label_arr = labels.get_label_arr(window)
+                    zero_inds = label_arr.ravel() == 0
+                    chip_shape = chip.shape
+                    if np.any(zero_inds):
+                        chip = np.reshape(chip, (-1, chip.shape[2]))
+                        chip[zero_inds, :] = 0
+                        chip = np.reshape(chip, chip_shape)
+
+                    data.append(chip, window, labels)
+                # Shuffle data so the first N samples which are displayed in
+                # Tensorboard are more diverse.
+                data.shuffle()
+
+                # Process augmentation
+                if augment:
+                    for augmentor in augmentors:
+                        data = augmentor.process(data, tmp_dir)
+
+                return self.backend.process_scene_data(scene, data, tmp_dir)
+
+        def _process_scenes(scenes, type_, augment):
+            return [_process_scene(scene, type_, augment) for scene in scenes]
+
+        processed_training_results = _process_scenes(
+            train_scenes, TRAIN, augment=True)
+        processed_validation_results = _process_scenes(
+            validation_scenes, VALIDATION, augment=False)
+
+        self.backend.process_sceneset_results(
+            processed_training_results, processed_validation_results, tmp_dir)
+
     def get_train_labels(self, window: Box, scene: Scene) -> np.ndarray:
         """Get the training labels for the given window in the given scene.
 

rastervision/task/task.py

@@ -123,7 +123,6 @@ def _process_scene(scene, type_, augment):
         def _process_scenes(scenes, type_, augment):
             return [_process_scene(scene, type_, augment) for scene in scenes]
 
-        # TODO: parallel processing!
         processed_training_results = _process_scenes(
             train_scenes, TRAIN, augment=True)
         processed_validation_results = _process_scenes(

tests/evaluation/test_semantic_segmentation_evaluation.py

@@ -53,6 +53,27 @@ def test_compute(self):
         self.assertEqual(precision2, eval.class_to_eval_item[2].precision)
         self.assertAlmostEqual(recall2, eval.class_to_eval_item[2].recall)
 
+    def test_compute_ignore_class(self):
+        gt_array = np.ones((4, 4, 1), dtype=np.uint8)
+        gt_array[0, 0, 0] = 0
+        gt_raster = MockRasterSource([0], 1)
+        gt_raster.set_raster(gt_array)
+        gt_label_source = SemanticSegmentationLabelSource(source=gt_raster)
+
+        pred_array = np.ones((4, 4, 1), dtype=np.uint8)
+        pred_raster = MockRasterSource([0], 1)
+        pred_raster.set_raster(pred_array)
+        pred_label_source = SemanticSegmentationLabelSource(source=pred_raster)
+
+        class_map = ClassMap(
+            [ClassItem(id=0, name='ignore'),
+             ClassItem(id=1, name='one')])
+        eval = SemanticSegmentationEvaluation(class_map)
+        eval.compute(gt_label_source.get_labels(),
+                     pred_label_source.get_labels())
+        self.assertAlmostEqual(1, len(eval.class_to_eval_item))
+        self.assertAlmostEqual(1.0, eval.class_to_eval_item[0].f1)
+
     def test_vector_compute(self):
         class_map = ClassMap([ClassItem(id=1, name='one', color='#000021')])
         gt_uri = data_file_path('3-gt-polygons.geojson')