Repository: azavea/raster-vision

Issue:
Multi-temporal raster source visualizer fails when batch size is 1
https://github.com/azavea/raster-vision/blob/e4e10ad04313bbe5355693ef96f3854f7963f2b1/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py#L122-L127

This code fails when the batch size is 1 because the created `subfigs` doesn't have a `flat` property if there's only one row and one column. Not sure whether this should be fixed upstream in `matplotlib`...

Matplotlib version 3.7.1, rastervision version 0.21.2

```python
import matplotlib.pyplot as plt
fig = plt.figure()
subfigs = fig.subfigures(nrows=2, ncols=1, hspace=0)
subfigs.flat
#> <numpy.flatiter object at 0x5575c63e68f0>

subfigs = fig.subfigures(nrows=1, ncols=1, hspace=0)
subfigs.flat
#> Traceback (most recent call last):
#>   File "<string>", line 1, in <module>
#> AttributeError: 'SubFigure' object has no attribute 'flat'
```

<sup>Created at 2023-10-11 17:12:31 CDT by [reprexlite](https://github.com/jayqi/reprexlite) v0.5.0</sup>



Discussion:
Looks like `fig.subfigures()` needs `squeeze=False` to act more consistently, just like `plt.subplots()`. Thanks for catching that!

Resolving patch:
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py
--- a/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py
+++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py
@@ -119,7 +119,8 @@ def plot_batch(self,
             batch_sz, T, *_ = x.shape
             params['fig_args']['figsize'][1] *= T
             fig = plt.figure(**params['fig_args'])
-            subfigs = fig.subfigures(nrows=batch_sz, ncols=1, hspace=0.0)
+            subfigs = fig.subfigures(
+                nrows=batch_sz, ncols=1, hspace=0.0, squeeze=False)
             subfig_axs = [
                 subfig.subplots(
                     nrows=T, ncols=params['subplot_args']['ncols'])


Test patch:
diff --git a/tests/pytorch_learner/dataset/visualizer/test_semantic_segmentation_visualizer.py b/tests/pytorch_learner/dataset/visualizer/test_semantic_segmentation_visualizer.py
--- a/tests/pytorch_learner/dataset/visualizer/test_semantic_segmentation_visualizer.py
+++ b/tests/pytorch_learner/dataset/visualizer/test_semantic_segmentation_visualizer.py
@@ -40,6 +40,8 @@ def test_plot_batch_temporal(self):
         x = torch.randn(size=(2, 3, 4, 256, 256))
         y = (torch.randn(size=(2, 256, 256)) > 0).long()
         self.assertNoError(lambda: viz.plot_batch(x, y))
+        # w/o z, batch size = 1
+        self.assertNoError(lambda: viz.plot_batch(x[[0]], y[[0]]))
 
         # w/ z
         viz = SemanticSegmentationVisualizer(
@@ -50,3 +52,5 @@ def test_plot_batch_temporal(self):
         y = (torch.randn(size=(2, 256, 256)) > 0).long()
         z = torch.randn(size=(2, num_classes, 256, 256)).softmax(dim=-3)
         self.assertNoError(lambda: viz.plot_batch(x, y, z=z))
+        # w/ z, batch size = 1
+        self.assertNoError(lambda: viz.plot_batch(x[[0]], y[[0]]))