# Buggy code
def _labels_to_full_label_arr(self, labels: SemanticSegmentationLabels) -> np.ndarray:
    """Get an array of labels covering the full extent."""
    try:
        label_arr = labels.get_label_arr(self.extent)
        return label_arr
    except KeyError:
        pass

    # construct the array from individual windows
    windows = labels.get_windows()

    # value for pixels not convered by any windows
    try:
        default_class_id = self.class_config.get_null_class_id()
    except ValueError:
        # Set it to a high value so that it doesn't match any class's id.
        # assumption: num_classes < 256
        default_class_id = 255

    label_arr = np.full(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)
        label_arr[ymin:ymax, xmin:xmax] = arr
    return label_arr


# Traceback (KeyError)
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)

# Fixed code
def _labels_to_full_label_arr(self, labels: SemanticSegmentationLabels) -> np.ndarray:
    """Get an array of labels covering the full extent."""
    try:
        label_arr = labels.get_label_arr(self.extent)
        return label_arr
    except KeyError:
        pass

    # we will construct the array from individual windows
    windows = labels.get_windows()

    # value for pixels not convered by any windows
    try:
        default_class_id = self.class_config.get_null_class_id()
    except ValueError:
        # Set it to a high value so that it doesn't match any class's id.
        # assumption: num_classes < 256
        default_class_id = 255

    label_arr = np.full(self.extent.size, fill_value=default_class_id, dtype=np.uint8)

    for w in windows:
        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