Repository: azavea/raster-vision

Issue:
Prediction error: Box key error
I encountered an error while running predict as follows:

```
Running predict command...
2021-01-18 09:29:36:rastervision.pytorch_learner.learner: INFO - Loading model weights from: /opt/data/tmp/tmpno2kuo4_/model-bundle/model.pth
2021-01-18 09:29:38:rastervision.core.rv_pipeline.rv_pipeline: INFO - Making predictions for scene
................................................................................................................................................................................................................................................................................................................................................................................................................
2021-01-18 09:29:48:rastervision.core.data.label_store.semantic_segmentation_label_store: INFO - Writing labels to disk.
  [#-----------------------------------]    5%  0d 00:00:48
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/conda/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 248, in <module>
    main()
  File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 722, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 697, in main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 1066, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 895, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.6/site-packages/click/core.py", line 535, in invoke
    return callback(*args, **kwargs)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 240, in run_command
    runner=runner)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 217, in _run_command
    command_fn()
  File "/opt/src/rastervision_core/rastervision/core/rv_pipeline/rv_pipeline.py", line 165, in predict
    for s in dataset.validation_scenes
  File "/opt/src/rastervision_core/rastervision/core/rv_pipeline/rv_pipeline.py", line 161, in _predict
    label_store.save(labels)
  File "/opt/src/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py", line 194, in save
    labels)
  File "/opt/src/rastervision_core/rastervision/core/data/label_store/semantic_segmentation_label_store.py", line 256, in write_discrete_raster_output
    label_arr = labels.get_label_arr(window)
  File "/opt/src/rastervision_core/rastervision/core/data/label/semantic_segmentation_labels.py", line 170, in get_label_arr
    return self.window_to_label_arr[window]
KeyError: Box(0, 4864, 256, 5000)
```

What does `Box(0, 4864, 256, 5000)` imply? 256 is my chip size and 5000 is the size of each tif image but I don't follow what the other numbers mean.


Edit: Adding module code

```
import os
from os.path import join, basename

from rastervision.core.rv_pipeline import *
from rastervision.core.backend import *
from rastervision.core.data import *
from rastervision.core.analyzer import *
from rastervision.pytorch_backend import *
from rastervision.pytorch_learner import *
from rastervision.pytorch_backend.examples.utils import (get_scene_info,
                                                         save_image_crop)
from rastervision.pytorch_backend.examples.semantic_segmentation.utils import (
    example_multiband_transform, example_rgb_transform, imagenet_stats,
    Unnormalize)

# INRIA dataset in total contains 360 images of size 5000 x 5000 across 10 cities
TRAIN_CITIES = ["austin", "chicago", "kitsap", "tyrol-w", "vienna"]
VAL_CITIES = ["vienna", "kitsap"]
# Debug
TRAIN_CITIES = ["austin"]
VAL_CITIES = ["vienna"]

TEST_CITIES = ["bellingham", "bloomington", "innsbruck", "sfo", "tyrol-e"]
NUM_IMAGES_PER_CITY = 1 #36

CLASS_NAMES = ["background", "building"]
CLASS_COLORS = ["white", "red"]

raw_uri = "/opt/data/inria/AerialImageDataset"
root_uri = "/opt/src/code/output"
nochip = False
augment = False
test= False

def get_config(runner,
               raw_uri: str,
               root_uri: str,
               augment: bool = False,
               nochip: bool = False,
               test: bool = False):

    channel_order = [0, 1, 2]
    channel_display_groups = None
    aug_transform = example_rgb_transform

    if augment:
        mu, std = imagenet_stats['mean'], imagenet_stats['std']
        mu, std = mu[channel_order], std[channel_order]

        base_transform = A.Normalize(mean=mu.tolist(), std=std.tolist())
        plot_transform = Unnormalize(mean=mu, std=std)

        aug_transform = A.to_dict(aug_transform)
        base_transform = A.to_dict(base_transform)
        plot_transform = A.to_dict(plot_transform)
    else:
        aug_transform = None
        base_transform = None
        plot_transform = None


    chip_sz = 256
    img_sz = chip_sz
    
    if nochip:
        chip_options = SemanticSegmentationChipOptions()
    else:
        chip_options = SemanticSegmentationChipOptions(
            window_method=SemanticSegmentationWindowMethod.sliding,
            stride=chip_sz)

    class_config = ClassConfig(names=CLASS_NAMES, colors=CLASS_COLORS)
    class_config.ensure_null_class()

    
    def make_scene(id) -> SceneConfig:
        raster_uri = f'{raw_uri}/train/images/{id}.tif'
        label_uri = f'{raw_uri}/train/gt/{id}.tif'


        raster_source = RasterioSourceConfig(
            uris=[raster_uri], channel_order=[0,1,2])

        label_source = SemanticSegmentationLabelSourceConfig(
            raster_source=RasterioSourceConfig(uris=[label_uri],
            transformers=[
                #ReclassTransformerConfig(mapping={255: 1})
                CastTransformerConfig(to_dtype='bool'),  # 255 --> True
                CastTransformerConfig(to_dtype='uint8')  # True --> 1
            ]))

        # URI will be injected by scene config.
        # Using rgb=True because we want prediction TIFFs to be in
        # RGB format.
        label_store = SemanticSegmentationLabelStoreConfig(
            rgb=True, vector_output=[PolygonVectorOutputConfig(class_id=0)])

        scene = SceneConfig(
            id=id,
            raster_source=raster_source,
            label_source=label_source,
            label_store=label_store)

        return scene

    scene_dataset = DatasetConfig(
        class_config=class_config,
        train_scenes=[make_scene(city + str(n)) for city in TRAIN_CITIES for n in range(1, NUM_IMAGES_PER_CITY + 1)],
        validation_scenes=[make_scene(city + str(n)) for city in VAL_CITIES for n in range(1, NUM_IMAGES_PER_CITY + 1)])


    if nochip:
        window_opts = {}
        # set window configs for training scenes
        for s in scene_dataset.train_scenes:
            window_opts[s.id] = GeoDataWindowConfig(
                # method=GeoDataWindowMethod.sliding,
                method=GeoDataWindowMethod.random,
                size=img_sz,
                # size_lims=(200, 300),
                h_lims=(200, 300),
                w_lims=(200, 300),
                max_windows=2209,
            )
        # set window configs for validation scenes
        for s in scene_dataset.validation_scenes:
            window_opts[s.id] = GeoDataWindowConfig(
                method=GeoDataWindowMethod.sliding,
                size=img_sz,
                stride=img_sz // 2)

        data = SemanticSegmentationGeoDataConfig(
            scene_dataset=scene_dataset,
            window_opts=window_opts,
            img_sz=img_sz,
            img_channels=len(channel_order),
            num_workers=4,
            channel_display_groups=channel_display_groups)
    else:
        data = SemanticSegmentationImageDataConfig(
            img_sz=img_sz,
            img_channels=len(channel_order),
            num_workers=4,
            channel_display_groups=channel_display_groups,
            base_transform=base_transform,
            aug_transform=aug_transform,
            plot_options=PlotOptions(transform=plot_transform))

    model = SemanticSegmentationModelConfig(backbone=Backbone.resnet50)

    backend = PyTorchSemanticSegmentationConfig(
        data=data,
        model=model,
        solver=SolverConfig(
            lr=1e-4,
            num_epochs=5,
            test_num_epochs=2,
            batch_sz=8,
            test_batch_sz=2,
            one_cycle=True),
        log_tensorboard=True,
        run_tensorboard=False,
        test_mode=test)

    pipeline = SemanticSegmentationConfig(
        root_uri=root_uri,
        dataset=scene_dataset,
        backend=backend,
        channel_display_groups=channel_display_groups,
        train_chip_sz=chip_sz,
        predict_chip_sz=chip_sz,
        chip_options=chip_options)

    return pipeline
```



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
@@ -1,4 +1,4 @@
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Tuple
 from os.path import join
 import logging
 import click
@@ -173,7 +173,7 @@ def save(self, labels: SemanticSegmentationLabels) -> None:
         local_root = get_local_path(self.root_uri, self.tmp_dir)
         make_dir(local_root)
 
-        out_smooth_profile = {
+        out_profile = {
             'driver': 'GTiff',
             'height': self.extent.ymax,
             'width': self.extent.xmax,
@@ -190,12 +190,11 @@ def save(self, labels: SemanticSegmentationLabels) -> None:
             labels += old_labels
 
         self.write_discrete_raster_output(
-            out_smooth_profile, get_local_path(self.label_uri, self.tmp_dir),
-            labels)
+            out_profile, get_local_path(self.label_uri, self.tmp_dir), labels)
 
         if self.smooth_output:
             self.write_smooth_raster_output(
-                out_smooth_profile,
+                out_profile,
                 get_local_path(self.score_uri, self.tmp_dir),
                 get_local_path(self.hits_uri, self.tmp_dir),
                 labels,
@@ -207,14 +206,14 @@ def save(self, labels: SemanticSegmentationLabels) -> None:
         sync_to_dir(local_root, self.root_uri)
 
     def write_smooth_raster_output(self,
-                                   out_smooth_profile: dict,
+                                   out_profile: dict,
                                    scores_path: str,
                                    hits_path: str,
                                    labels: SemanticSegmentationLabels,
                                    chip_sz: Optional[int] = None) -> None:
         dtype = np.uint8 if self.smooth_as_uint8 else np.float32
 
-        out_smooth_profile.update({
+        out_profile.update({
             'count': labels.num_classes,
             'dtype': dtype,
         })
@@ -224,45 +223,38 @@ def write_smooth_raster_output(self,
             windows = labels.get_windows(chip_sz=chip_sz)
 
         log.info('Writing smooth labels to disk.')
-        with rio.open(scores_path, 'w', **out_smooth_profile) as dataset:
+        with rio.open(scores_path, 'w', **out_profile) as dataset:
             with click.progressbar(windows) as bar:
                 for window in bar:
-                    window = window.intersection(self.extent)
+                    window, _ = self._clip_to_extent(self.extent, window)
                     score_arr = labels.get_score_arr(window)
                     if self.smooth_as_uint8:
-                        score_arr *= 255
-                        score_arr = np.around(score_arr, out=score_arr)
-                        score_arr = score_arr.astype(dtype)
-                    window = window.rasterio_format()
-                    for i, class_scores in enumerate(score_arr, start=1):
-                        dataset.write_band(i, class_scores, window=window)
+                        score_arr = self._scores_to_uint8(score_arr)
+                    self._write_array(dataset, window, score_arr)
         # save pixel hits too
         np.save(hits_path, labels.pixel_hits)
 
     def write_discrete_raster_output(
-            self, out_smooth_profile: dict, path: str,
+            self, out_profile: dict, path: str,
             labels: SemanticSegmentationLabels) -> None:
 
         num_bands = 1 if self.class_transformer is None else 3
-        out_smooth_profile.update({'count': num_bands, 'dtype': np.uint8})
+        out_profile.update({'count': num_bands, 'dtype': np.uint8})
 
         windows = labels.get_windows()
 
         log.info('Writing labels to disk.')
-        with rio.open(path, 'w', **out_smooth_profile) as dataset:
+        with rio.open(path, 'w', **out_profile) as dataset:
             with click.progressbar(windows) as bar:
                 for window in bar:
-                    window = window.intersection(self.extent)
                     label_arr = labels.get_label_arr(window)
-                    window = window.rasterio_format()
-                    if self.class_transformer is None:
-                        dataset.write_band(1, label_arr, window=window)
-                    else:
-                        rgb_labels = self.class_transformer.class_to_rgb(
+                    window, label_arr = self._clip_to_extent(
+                        self.extent, window, label_arr)
+                    if self.class_transformer is not None:
+                        label_arr = self.class_transformer.class_to_rgb(
                             label_arr)
-                        rgb_labels = rgb_labels.transpose(2, 0, 1)
-                        for i, band in enumerate(rgb_labels, start=1):
-                            dataset.write_band(i, band, window=window)
+                        label_arr = label_arr.transpose(2, 0, 1)
+                    self._write_array(dataset, window, label_arr)
 
     def _labels_to_full_label_arr(
             self, labels: SemanticSegmentationLabels) -> np.ndarray:
@@ -273,7 +265,7 @@ def _labels_to_full_label_arr(
         except KeyError:
             pass
 
-        # construct the array from individual windows
+        # we will construct the array from individual windows
         windows = labels.get_windows()
 
         # value for pixels not convered by any windows
@@ -288,9 +280,9 @@ def _labels_to_full_label_arr(
             self.extent.size, fill_value=default_class_id, dtype=np.uint8)
 
         for w in windows:
-            w = w.intersection(self.extent)
             ymin, xmin, ymax, xmax = w
             arr = labels.get_label_arr(w)
+            w, arr = self._clip_to_extent(self.extent, w, arr)
             label_arr[ymin:ymax, xmin:xmax] = arr
         return label_arr
 
@@ -342,3 +334,33 @@ def empty_labels(self) -> SemanticSegmentationLabels:
             extent=self.extent,
             num_classes=len(self.class_config))
         return labels
+
+    def _write_array(self, dataset: rio.DatasetReader, window: Box,
+                     arr: np.ndarray) -> None:
+        """Write array out to a rasterio dataset. Array must be of shape
+        (C, H, W).
+        """
+        window = window.rasterio_format()
+        if len(arr.shape) == 2:
+            dataset.write_band(1, arr, window=window)
+        else:
+            for i, band in enumerate(arr, start=1):
+                dataset.write_band(i, band, window=window)
+
+    def _clip_to_extent(self,
+                        extent: Box,
+                        window: Box,
+                        arr: Optional[np.ndarray] = None
+                        ) -> Tuple[Box, Optional[np.ndarray]]:
+        clipped_window = window.intersection(extent)
+        if arr is not None:
+            h, w = clipped_window.size
+            arr = arr[:h, :w]
+        return clipped_window, arr
+
+    def _scores_to_uint8(self, score_arr: np.ndarray) -> np.ndarray:
+        """Quantize scores to uint8 (0-255)."""
+        score_arr *= 255
+        score_arr = np.around(score_arr, out=score_arr)
+        score_arr = score_arr.astype(np.uint8)
+        return score_arr


Test patch: