# Buggy code
def write_smooth_raster_output(
    self,
    out_smooth_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(
        {
            "count": labels.num_classes,
            "dtype": dtype,
        }
    )
    if chip_sz is None:
        windows = [self.extent]
    else:
        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 click.progressbar(windows) as bar:
            for window in bar:
                window = window.intersection(self.extent)
                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)
    # save pixel hits too
    np.save(hits_path, labels.pixel_hits)


# 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 write_smooth_raster_output(
    self,
    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_profile.update(
        {
            "count": labels.num_classes,
            "dtype": dtype,
        }
    )
    if chip_sz is None:
        windows = [self.extent]
    else:
        windows = labels.get_windows(chip_sz=chip_sz)

    log.info("Writing smooth labels to disk.")
    with rio.open(scores_path, "w", **out_profile) as dataset:
        with click.progressbar(windows) as bar:
            for window in bar:
                window, _ = self._clip_to_extent(self.extent, window)
                score_arr = labels.get_score_arr(window)
                if self.smooth_as_uint8:
                    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)