Repository: azavea/raster-vision

Issue:
Apply base_transform in the predict stage
Currently, the predict stage uses an entirely different code path than validation, and therefore, the input does not go through the `base_transform` transform (which might include things like resizing and normalization). This means that the input is not guaranteed to be in the form that the model expects.


Discussion:

Resolving patch:
diff --git a/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py b/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py
--- a/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py
+++ b/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py
@@ -239,7 +239,8 @@ def write_discrete_raster_output(
             labels: SemanticSegmentationLabels) -> None:
 
         num_bands = 1 if self.class_transformer is None else 3
-        out_profile.update({'count': num_bands, 'dtype': np.uint8})
+        dtype = np.uint8
+        out_profile.update({'count': num_bands, 'dtype': dtype})
 
         windows = labels.get_windows()
 
@@ -247,7 +248,7 @@ def write_discrete_raster_output(
         with rio.open(path, 'w', **out_profile) as dataset:
             with click.progressbar(windows) as bar:
                 for window in bar:
-                    label_arr = labels.get_label_arr(window)
+                    label_arr = labels.get_label_arr(window).astype(dtype)
                     window, label_arr = self._clip_to_extent(
                         self.extent, window, label_arr)
                     if self.class_transformer is not None:
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py
--- a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py
+++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py
@@ -732,7 +732,7 @@ def to_batch(self, x: Tensor) -> Tensor:
             x with extra batch dimension of length 1 if needed
         """
         if x.ndim == 3:
-            x = x.unsqueeze(0)
+            x = x[None, ...]
         return x
 
     def normalize_input(self, x: np.ndarray) -> np.ndarray:
@@ -797,9 +797,11 @@ def numpy_predict(self, x: np.ndarray,
         Returns:
             predictions using numpy arrays
         """
+        transform, _ = self.get_data_transforms()
         x = self.normalize_input(x)
-        x = torch.tensor(x)
         x = self.to_batch(x)
+        x = np.stack([transform(image=img)['image'] for img in x])
+        x = torch.from_numpy(x)
         x = x.permute((0, 3, 1, 2))
         out = self.predict(x, raw_out=raw_out)
         return self.output_to_numpy(out)
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/object_detection_learner.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/object_detection_learner.py
--- a/rastervision_pytorch_learner/rastervision/pytorch_learner/object_detection_learner.py
+++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/object_detection_learner.py
@@ -7,6 +7,9 @@
 matplotlib.use('Agg')  # noqa
 from albumentations import BboxParams
 
+import numpy as np
+import torch
+
 from rastervision.pytorch_learner.learner import Learner
 from rastervision.pytorch_learner.object_detection_utils import (
     MyFasterRCNN, compute_coco_eval, collate_fn, plot_xyz)
@@ -66,6 +69,20 @@ def validate_end(self, outputs, num_samples):
             metrics = {'map': coco_metrics[0], 'map50': coco_metrics[1]}
         return metrics
 
+    def numpy_predict(self, x: np.ndarray,
+                      raw_out: bool = False) -> np.ndarray:
+        transform, _ = self.get_data_transforms()
+        x = self.normalize_input(x)
+        x = self.to_batch(x)
+        x = np.stack([
+            transform(image=img, bboxes=[], category_id=[])['image']
+            for img in x
+        ])
+        x = torch.from_numpy(x)
+        x = x.permute((0, 3, 1, 2))
+        out = self.predict(x, raw_out=raw_out)
+        return self.output_to_numpy(out)
+
     def output_to_numpy(self, out):
         numpy_out = []
         for boxlist in out:


Test patch: