Repository: azavea/raster-vision

Issue:
Fix raster stats bug
If you run compute_raster_stats on 4-channel imagery (yielding stats for 4 channels), and use a `channel_order` of [0, 1, 2] in your raster_transformer, and then switch to using 3-channel imagery, it leads to an error because currently the `means` do not have the `channel_order` applied to them before being subtracted from the raster. In other words, 4 channel means is subtracted from a 3 channel raster.


Discussion:

Resolving patch:
diff --git a/src/rastervision/builders/raster_transformer_builder.py b/src/rastervision/builders/raster_transformer_builder.py
--- a/src/rastervision/builders/raster_transformer_builder.py
+++ b/src/rastervision/builders/raster_transformer_builder.py
@@ -1,5 +1,12 @@
 from rastervision.core.raster_transformer import RasterTransformer
+from rastervision.core.raster_stats import RasterStats
 
 
 def build(config):
-    return RasterTransformer(config)
+    raster_stats = None
+    if config.stats_uri:
+        raster_stats = RasterStats()
+        raster_stats.load(config.stats_uri)
+
+    return RasterTransformer(
+        channel_order=config.channel_order, raster_stats=raster_stats)
diff --git a/src/rastervision/core/raster_transformer.py b/src/rastervision/core/raster_transformer.py
--- a/src/rastervision/core/raster_transformer.py
+++ b/src/rastervision/core/raster_transformer.py
@@ -1,43 +1,50 @@
 import numpy as np
 
-from rastervision.core.raster_stats import RasterStats
-
 
 class RasterTransformer(object):
-    """Transforms chips according to a config."""
+    """Transforms raw chips to be input to a neural network."""
 
-    def __init__(self, options):
+    def __init__(self, channel_order=None, raster_stats=None):
         """Construct a new RasterTransformer.
 
         Args:
-            options: protos.raster_transformer_pb2.RasterTransformer
+            channel_order: numpy array of length n where n is the number of
+                channels to use and the values are channel indices
+            raster_stats: (RasterStats) used to transform chip to have
+                desired statistics
         """
-        self.options = options
-        self.raster_stats = None
-        if options.stats_uri:
-            self.raster_stats = RasterStats()
-            self.raster_stats.load(options.stats_uri)
+        self.channel_order = channel_order
+        self.raster_stats = raster_stats
 
     def transform(self, chip):
         """Transform a chip.
 
         Selects a subset of the channels and transforms non-uint8 to
-        uint8 values using options.stats_uri
+        uint8 values using raster_stats.
 
         Args:
             chip: [height, width, channels] numpy array
 
         Returns:
             [height, width, channels] uint8 numpy array where channels is equal
-                to len(self.options.channel_order)
+                to len(channel_order)
         """
+        if self.channel_order is None:
+            channel_order = np.arange(chip.shape[2])
+        else:
+            channel_order = self.channel_order
+
+        chip = chip[:, :, channel_order]
+
         if chip.dtype != np.uint8:
             if self.raster_stats:
                 # Subtract mean and divide by std to get zscores.
                 means = np.array(self.raster_stats.means)
-                means = means[np.newaxis, np.newaxis, :].astype(np.float)
+                means = means[np.newaxis, np.newaxis, channel_order].astype(
+                    np.float)
                 stds = np.array(self.raster_stats.stds)
-                stds = stds[np.newaxis, np.newaxis, :].astype(np.float)
+                stds = stds[np.newaxis, np.newaxis, channel_order].astype(
+                    np.float)
 
                 # Don't transform NODATA zero values.
                 nodata = chip == 0
@@ -56,8 +63,6 @@ def transform(self, chip):
                 chip[nodata] = 0
             else:
                 raise ValueError(
-                    'Need to provide stats_uri for non-uint8 rasters.')
+                    'Need to provide raster_stats for non-uint8 rasters.')
 
-        if self.options.channel_order:
-            return chip[:, :, self.options.channel_order]
         return chip


Test patch:
diff --git a/src/rastervision/core/raster_transformer_test.py b/src/rastervision/core/raster_transformer_test.py
new file mode 100644
--- /dev/null
+++ b/src/rastervision/core/raster_transformer_test.py
@@ -0,0 +1,66 @@
+import unittest
+
+import numpy as np
+
+from rastervision.core.raster_transformer import RasterTransformer
+from rastervision.core.raster_stats import RasterStats
+
+
+class TestRasterTransformer(unittest.TestCase):
+    def test_no_channel_order_no_stats(self):
+        transformer = RasterTransformer()
+        chip = np.ones((2, 2, 3)).astype(np.uint8)
+        out_chip = transformer.transform(chip)
+        np.testing.assert_equal(chip, out_chip)
+
+        # Need to supply raster_stats for non-uint8 chips.
+        chip = np.ones((2, 2, 3))
+        with self.assertRaises(ValueError):
+            out_chip = transformer.transform(chip)
+
+    def test_no_channel_order_has_stats(self):
+        raster_stats = RasterStats()
+        raster_stats.means = np.ones((4, ))
+        raster_stats.stds = np.ones((4, )) * 2
+
+        # All values have z-score of 1, which translates to
+        # uint8 value of 170.
+        transformer = RasterTransformer(raster_stats=raster_stats)
+        chip = np.ones((2, 2, 4)) * 3
+        out_chip = transformer.transform(chip)
+        expected_out_chip = np.ones((2, 2, 4)) * 170
+        np.testing.assert_equal(out_chip, expected_out_chip)
+
+    def test_has_channel_order_no_stats(self):
+        channel_order = [0, 1, 2]
+        transformer = RasterTransformer(channel_order=channel_order)
+        chip = np.ones((2, 2, 4)).astype(np.uint8)
+        chip[:, :, :] *= np.array([0, 1, 2, 3]).astype(np.uint8)
+        out_chip = transformer.transform(chip)
+        expected_out_chip = np.ones((2, 2, 3)).astype(np.uint8)
+        expected_out_chip[:, :, :] *= np.array([0, 1, 2]).astype(np.uint8)
+        np.testing.assert_equal(out_chip, expected_out_chip)
+
+    def test_has_channel_order_has_stats(self):
+        raster_stats = RasterStats()
+        raster_stats.means = np.ones((4, ))
+        raster_stats.stds = np.ones((4, )) * 2
+        channel_order = [0, 1, 2]
+        transformer = RasterTransformer(
+            raster_stats=raster_stats, channel_order=channel_order)
+
+        chip = np.ones((2, 2, 4)) * [3, 3, 3, 0]
+        out_chip = transformer.transform(chip)
+        expected_out_chip = np.ones((2, 2, 3)) * 170
+        np.testing.assert_equal(out_chip, expected_out_chip)
+
+        # Also test when chip has same number of channels as channel_order
+        # but different number of channels than stats.
+        chip = np.ones((2, 2, 3)) * [3, 3, 3]
+        out_chip = transformer.transform(chip)
+        expected_out_chip = np.ones((2, 2, 3)) * 170
+        np.testing.assert_equal(out_chip, expected_out_chip)
+
+
+if __name__ == '__main__':
+    unittest.main()