We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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>


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py`
Content:
```
1 from typing import (TYPE_CHECKING, Sequence, Optional, List, Dict, Union,
2                     Tuple, Any)
3 from abc import ABC, abstractmethod
4 
5 import numpy as np
6 import torch
7 from torch import Tensor
8 import albumentations as A
9 from torch.utils.data import DataLoader
10 import matplotlib.pyplot as plt
11 
12 from rastervision.pipeline.file_system import make_dir
13 from rastervision.pytorch_learner.utils import (
14     deserialize_albumentation_transform, validate_albumentation_transform,
15     MinMaxNormalize)
16 from rastervision.pytorch_learner.learner_config import (
17     RGBTuple,
18     ChannelInds,
19     ensure_class_colors,
20     validate_channel_display_groups,
21     get_default_channel_display_groups,
22 )
23 
24 if TYPE_CHECKING:
25     from torch.utils.data import Dataset
26     from matplotlib.figure import Figure
27 
28 
29 class Visualizer(ABC):
30     """Base class for plotting samples from computer vision PyTorch Datasets."""
31 
32     scale: float = 3.
33 
34     def __init__(self,
35                  class_names: List[str],
36                  class_colors: Optional[List[Union[str, RGBTuple]]] = None,
37                  transform: Optional[Dict] = A.to_dict(MinMaxNormalize()),
38                  channel_display_groups: Optional[Union[Dict[
39                      str, ChannelInds], Sequence[ChannelInds]]] = None):
40         """Constructor.
41 
42         Args:
43             class_names: names of classes
44             class_colors: Colors used to display classes. Can be color 3-tuples
45                 in list form.
46             transform: An Albumentations transform serialized as a dict that
47                 will be applied to each image before it is plotted. Mainly useful
48                 for undoing any data transformation that you do not want included in
49                 the plot, such as normalization. The default value will shift and scale
50                 the image so the values range from 0.0 to 1.0 which is the expected range
51                 for the plotting function. This default is useful for cases where the
52                 values after normalization are close to zero which makes the plot
53                 difficult to see.
54             channel_display_groups: Groups of image channels to display together as a
55                 subplot when plotting the data and predictions.
56                 Can be a list or tuple of groups (e.g. [(0, 1, 2), (3,)]) or a
57                 dict containing title-to-group mappings
58                 (e.g. {"RGB": [0, 1, 2], "IR": [3]}),
59                 where each group is a list or tuple of channel indices and
60                 title is a string that will be used as the title of the subplot
61                 for that group.
62         """
63         self.class_names = class_names
64         self.class_colors = ensure_class_colors(self.class_names, class_colors)
65         self.transform = validate_albumentation_transform(transform)
66         self._channel_display_groups = validate_channel_display_groups(
67             channel_display_groups)
68 
69     @abstractmethod
70     def plot_xyz(self,
71                  axs,
72                  x: Tensor,
73                  y: Sequence,
74                  z: Optional[Sequence] = None,
75                  plot_title: bool = True):
76         """Plot image, ground truth labels, and predicted labels.
77 
78         Args:
79             axs: matplotlib axes on which to plot
80             x: image
81             y: ground truth labels
82             z: optional predicted labels
83         """
84         pass
85 
86     def plot_batch(self,
87                    x: Tensor,
88                    y: Sequence,
89                    output_path: Optional[str] = None,
90                    z: Optional[Sequence] = None,
91                    batch_limit: Optional[int] = None,
92                    show: bool = False):
93         """Plot a whole batch in a grid using plot_xyz.
94 
95         Args:
96             x: batch of images
97             y: ground truth labels
98             output_path: local path where to save plot image
99             z: optional predicted labels
100             batch_limit: optional limit on (rendered) batch size
101         """
102         params = self.get_plot_params(
103             x=x, y=y, z=z, output_path=output_path, batch_limit=batch_limit)
104         if params['subplot_args']['nrows'] == 0:
105             return
106 
107         if x.ndim == 4:
108             fig, axs = plt.subplots(**params['fig_args'],
109                                     **params['subplot_args'])
110             plot_xyz_args = params['plot_xyz_args']
111             self._plot_batch(fig, axs, plot_xyz_args, x, y=y, z=z)
112         elif x.ndim == 5:
113             # If a temporal dimension is present, we divide the figure into
114             # multiple subfigures--one for each batch. Then, in each subfigure,
115             # we plot all timesteps as if they were a single batch. To
116             # delineate the boundary b/w batch items, we adopt the convention
117             # of only displaying subplot titles once per batch (above the first
118             # row in each batch).
119             batch_sz, T, *_ = x.shape
120             params['fig_args']['figsize'][1] *= T
121             fig = plt.figure(**params['fig_args'])
122             subfigs = fig.subfigures(nrows=batch_sz, ncols=1, hspace=0.0)
123             subfig_axs = [
124                 subfig.subplots(
125                     nrows=T, ncols=params['subplot_args']['ncols'])
126                 for subfig in subfigs.flat
127             ]
128             for i, axs in enumerate(subfig_axs):
129                 plot_xyz_args = [
130                     dict(params['plot_xyz_args'][i]) for _ in range(T)
131                 ]
132                 plot_xyz_args[0]['plot_title'] = True
133                 for args in plot_xyz_args[1:]:
134                     args['plot_title'] = False
135                 _x = x[i]
136                 _y = [y[i]] * T
137                 _z = None if z is None else [z[i]] * T
138                 self._plot_batch(fig, axs, plot_xyz_args, _x, y=_y, z=_z)
139         else:
140             raise ValueError('Expected x to have 4 or 5 dims, but found '
141                              f'x.shape: {x.shape}')
142 
143         if show:
144             plt.show()
145         if output_path is not None:
146             make_dir(output_path, use_dirname=True)
147             fig.savefig(output_path, bbox_inches='tight', pad_inches=0.2)
148 
149         plt.close(fig)
150 
151     def _plot_batch(
152             self,
153             fig: 'Figure',
154             axs: Sequence,
155             plot_xyz_args: List[dict],
156             x: Tensor,
157             y: Optional[Sequence] = None,
158             z: Optional[Sequence] = None,
159     ):
160         # (N, c, h, w) --> (N, h, w, c)
161         x = x.permute(0, 2, 3, 1)
162 
163         # apply transform, if given
164         if self.transform is not None:
165             tf = deserialize_albumentation_transform(self.transform)
166             imgs = [tf(image=img)['image'] for img in x.numpy()]
167             x = torch.from_numpy(np.stack(imgs))
168 
169         for i, row_axs in enumerate(axs):
170             _z = None if z is None else z[i]
171             self.plot_xyz(row_axs, x[i], y[i], z=_z, **plot_xyz_args[i])
172 
173     def get_channel_display_groups(
174             self, nb_img_channels: int
175     ) -> Union[Dict[str, ChannelInds], Sequence[ChannelInds]]:
176         # The default channel_display_groups object depends on the number of
177         # channels in the image. This number is not known when the Visualizer
178         # is constructed which is why it needs to be created later.
179         if self._channel_display_groups is not None:
180             return self._channel_display_groups
181         return get_default_channel_display_groups(nb_img_channels)
182 
183     def get_collate_fn(self) -> Optional[callable]:
184         """Returns a custom collate_fn to use in DataLoader.
185 
186         None is returned if default collate_fn should be used.
187 
188         See https://pytorch.org/docs/stable/data.html#working-with-collate-fn
189         """
190         return None
191 
192     def get_batch(self, dataset: 'Dataset', batch_sz: int = 4,
193                   **kwargs) -> Tuple[Tensor, Any]:
194         """Generate a batch from a dataset.
195 
196         This is a convenience method for generating a batch of data to plot.
197 
198         Args:
199             dataset (Dataset): A Pytorch Datset.
200             batch_sz (int): Batch size. Defaults to 4.
201             **kwargs: Extra args for :class:`~torch.utils.data.DataLoader`.
202 
203         Returns:
204             Tuple[Tensor, Any]: (x, y) tuple where x is images and y is labels.
205         """
206         collate_fn = self.get_collate_fn()
207         dl = DataLoader(dataset, batch_sz, collate_fn=collate_fn, **kwargs)
208         try:
209             x, y = next(iter(dl))
210         except StopIteration:
211             raise ValueError('dataset did not return a batch')
212 
213         return x, y
214 
215     def get_plot_nrows(self, **kwargs) -> int:
216         x = kwargs['x']
217         batch_limit = kwargs.get('batch_limit')
218         batch_sz = x.shape[0]
219         nrows = min(batch_sz,
220                     batch_limit) if batch_limit is not None else batch_sz
221         return nrows
222 
223     def get_plot_ncols(self, **kwargs) -> int:
224         x = kwargs['x']
225         nb_img_channels = x.shape[1]
226         ncols = len(self.get_channel_display_groups(nb_img_channels))
227         return ncols
228 
229     def get_plot_params(self, **kwargs) -> dict:
230         nrows = self.get_plot_nrows(**kwargs)
231         ncols = self.get_plot_ncols(**kwargs)
232         params = {
233             'fig_args': {
234                 'constrained_layout': True,
235                 'figsize': np.array((self.scale * ncols, self.scale * nrows)),
236             },
237             'subplot_args': {
238                 'nrows': nrows,
239                 'ncols': ncols,
240                 'squeeze': False
241             },
242             'plot_xyz_args': [{} for _ in range(nrows)]
243         }
244         return params
245
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
             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'])

{"golden_diff": "diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py\n--- a/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py\n+++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py\n@@ -119,7 +119,8 @@\n             batch_sz, T, *_ = x.shape\n             params['fig_args']['figsize'][1] *= T\n             fig = plt.figure(**params['fig_args'])\n-            subfigs = fig.subfigures(nrows=batch_sz, ncols=1, hspace=0.0)\n+            subfigs = fig.subfigures(\n+                nrows=batch_sz, ncols=1, hspace=0.0, squeeze=False)\n             subfig_axs = [\n                 subfig.subplots(\n                     nrows=T, ncols=params['subplot_args']['ncols'])\n", "issue": "Multi-temporal raster source visualizer fails when batch size is 1\nhttps://github.com/azavea/raster-vision/blob/e4e10ad04313bbe5355693ef96f3854f7963f2b1/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py#L122-L127\r\n\r\nThis 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`...\r\n\r\nMatplotlib version 3.7.1, rastervision version 0.21.2\r\n\r\n```python\r\nimport matplotlib.pyplot as plt\r\nfig = plt.figure()\r\nsubfigs = fig.subfigures(nrows=2, ncols=1, hspace=0)\r\nsubfigs.flat\r\n#> <numpy.flatiter object at 0x5575c63e68f0>\r\n\r\nsubfigs = fig.subfigures(nrows=1, ncols=1, hspace=0)\r\nsubfigs.flat\r\n#> Traceback (most recent call last):\r\n#>   File \"<string>\", line 1, in <module>\r\n#> AttributeError: 'SubFigure' object has no attribute 'flat'\r\n```\r\n\r\n<sup>Created at 2023-10-11 17:12:31 CDT by [reprexlite](https://github.com/jayqi/reprexlite) v0.5.0</sup>\r\n\n", "before_files": [{"content": "from typing import (TYPE_CHECKING, Sequence, Optional, List, Dict, Union,\n                    Tuple, Any)\nfrom abc import ABC, abstractmethod\n\nimport numpy as np\nimport torch\nfrom torch import Tensor\nimport albumentations as A\nfrom torch.utils.data import DataLoader\nimport matplotlib.pyplot as plt\n\nfrom rastervision.pipeline.file_system import make_dir\nfrom rastervision.pytorch_learner.utils import (\n    deserialize_albumentation_transform, validate_albumentation_transform,\n    MinMaxNormalize)\nfrom rastervision.pytorch_learner.learner_config import (\n    RGBTuple,\n    ChannelInds,\n    ensure_class_colors,\n    validate_channel_display_groups,\n    get_default_channel_display_groups,\n)\n\nif TYPE_CHECKING:\n    from torch.utils.data import Dataset\n    from matplotlib.figure import Figure\n\n\nclass Visualizer(ABC):\n    \"\"\"Base class for plotting samples from computer vision PyTorch Datasets.\"\"\"\n\n    scale: float = 3.\n\n    def __init__(self,\n                 class_names: List[str],\n                 class_colors: Optional[List[Union[str, RGBTuple]]] = None,\n                 transform: Optional[Dict] = A.to_dict(MinMaxNormalize()),\n                 channel_display_groups: Optional[Union[Dict[\n                     str, ChannelInds], Sequence[ChannelInds]]] = None):\n        \"\"\"Constructor.\n\n        Args:\n            class_names: names of classes\n            class_colors: Colors used to display classes. Can be color 3-tuples\n                in list form.\n            transform: An Albumentations transform serialized as a dict that\n                will be applied to each image before it is plotted. Mainly useful\n                for undoing any data transformation that you do not want included in\n                the plot, such as normalization. The default value will shift and scale\n                the image so the values range from 0.0 to 1.0 which is the expected range\n                for the plotting function. This default is useful for cases where the\n                values after normalization are close to zero which makes the plot\n                difficult to see.\n            channel_display_groups: Groups of image channels to display together as a\n                subplot when plotting the data and predictions.\n                Can be a list or tuple of groups (e.g. [(0, 1, 2), (3,)]) or a\n                dict containing title-to-group mappings\n                (e.g. {\"RGB\": [0, 1, 2], \"IR\": [3]}),\n                where each group is a list or tuple of channel indices and\n                title is a string that will be used as the title of the subplot\n                for that group.\n        \"\"\"\n        self.class_names = class_names\n        self.class_colors = ensure_class_colors(self.class_names, class_colors)\n        self.transform = validate_albumentation_transform(transform)\n        self._channel_display_groups = validate_channel_display_groups(\n            channel_display_groups)\n\n    @abstractmethod\n    def plot_xyz(self,\n                 axs,\n                 x: Tensor,\n                 y: Sequence,\n                 z: Optional[Sequence] = None,\n                 plot_title: bool = True):\n        \"\"\"Plot image, ground truth labels, and predicted labels.\n\n        Args:\n            axs: matplotlib axes on which to plot\n            x: image\n            y: ground truth labels\n            z: optional predicted labels\n        \"\"\"\n        pass\n\n    def plot_batch(self,\n                   x: Tensor,\n                   y: Sequence,\n                   output_path: Optional[str] = None,\n                   z: Optional[Sequence] = None,\n                   batch_limit: Optional[int] = None,\n                   show: bool = False):\n        \"\"\"Plot a whole batch in a grid using plot_xyz.\n\n        Args:\n            x: batch of images\n            y: ground truth labels\n            output_path: local path where to save plot image\n            z: optional predicted labels\n            batch_limit: optional limit on (rendered) batch size\n        \"\"\"\n        params = self.get_plot_params(\n            x=x, y=y, z=z, output_path=output_path, batch_limit=batch_limit)\n        if params['subplot_args']['nrows'] == 0:\n            return\n\n        if x.ndim == 4:\n            fig, axs = plt.subplots(**params['fig_args'],\n                                    **params['subplot_args'])\n            plot_xyz_args = params['plot_xyz_args']\n            self._plot_batch(fig, axs, plot_xyz_args, x, y=y, z=z)\n        elif x.ndim == 5:\n            # If a temporal dimension is present, we divide the figure into\n            # multiple subfigures--one for each batch. Then, in each subfigure,\n            # we plot all timesteps as if they were a single batch. To\n            # delineate the boundary b/w batch items, we adopt the convention\n            # of only displaying subplot titles once per batch (above the first\n            # row in each batch).\n            batch_sz, T, *_ = x.shape\n            params['fig_args']['figsize'][1] *= T\n            fig = plt.figure(**params['fig_args'])\n            subfigs = fig.subfigures(nrows=batch_sz, ncols=1, hspace=0.0)\n            subfig_axs = [\n                subfig.subplots(\n                    nrows=T, ncols=params['subplot_args']['ncols'])\n                for subfig in subfigs.flat\n            ]\n            for i, axs in enumerate(subfig_axs):\n                plot_xyz_args = [\n                    dict(params['plot_xyz_args'][i]) for _ in range(T)\n                ]\n                plot_xyz_args[0]['plot_title'] = True\n                for args in plot_xyz_args[1:]:\n                    args['plot_title'] = False\n                _x = x[i]\n                _y = [y[i]] * T\n                _z = None if z is None else [z[i]] * T\n                self._plot_batch(fig, axs, plot_xyz_args, _x, y=_y, z=_z)\n        else:\n            raise ValueError('Expected x to have 4 or 5 dims, but found '\n                             f'x.shape: {x.shape}')\n\n        if show:\n            plt.show()\n        if output_path is not None:\n            make_dir(output_path, use_dirname=True)\n            fig.savefig(output_path, bbox_inches='tight', pad_inches=0.2)\n\n        plt.close(fig)\n\n    def _plot_batch(\n            self,\n            fig: 'Figure',\n            axs: Sequence,\n            plot_xyz_args: List[dict],\n            x: Tensor,\n            y: Optional[Sequence] = None,\n            z: Optional[Sequence] = None,\n    ):\n        # (N, c, h, w) --> (N, h, w, c)\n        x = x.permute(0, 2, 3, 1)\n\n        # apply transform, if given\n        if self.transform is not None:\n            tf = deserialize_albumentation_transform(self.transform)\n            imgs = [tf(image=img)['image'] for img in x.numpy()]\n            x = torch.from_numpy(np.stack(imgs))\n\n        for i, row_axs in enumerate(axs):\n            _z = None if z is None else z[i]\n            self.plot_xyz(row_axs, x[i], y[i], z=_z, **plot_xyz_args[i])\n\n    def get_channel_display_groups(\n            self, nb_img_channels: int\n    ) -> Union[Dict[str, ChannelInds], Sequence[ChannelInds]]:\n        # The default channel_display_groups object depends on the number of\n        # channels in the image. This number is not known when the Visualizer\n        # is constructed which is why it needs to be created later.\n        if self._channel_display_groups is not None:\n            return self._channel_display_groups\n        return get_default_channel_display_groups(nb_img_channels)\n\n    def get_collate_fn(self) -> Optional[callable]:\n        \"\"\"Returns a custom collate_fn to use in DataLoader.\n\n        None is returned if default collate_fn should be used.\n\n        See https://pytorch.org/docs/stable/data.html#working-with-collate-fn\n        \"\"\"\n        return None\n\n    def get_batch(self, dataset: 'Dataset', batch_sz: int = 4,\n                  **kwargs) -> Tuple[Tensor, Any]:\n        \"\"\"Generate a batch from a dataset.\n\n        This is a convenience method for generating a batch of data to plot.\n\n        Args:\n            dataset (Dataset): A Pytorch Datset.\n            batch_sz (int): Batch size. Defaults to 4.\n            **kwargs: Extra args for :class:`~torch.utils.data.DataLoader`.\n\n        Returns:\n            Tuple[Tensor, Any]: (x, y) tuple where x is images and y is labels.\n        \"\"\"\n        collate_fn = self.get_collate_fn()\n        dl = DataLoader(dataset, batch_sz, collate_fn=collate_fn, **kwargs)\n        try:\n            x, y = next(iter(dl))\n        except StopIteration:\n            raise ValueError('dataset did not return a batch')\n\n        return x, y\n\n    def get_plot_nrows(self, **kwargs) -> int:\n        x = kwargs['x']\n        batch_limit = kwargs.get('batch_limit')\n        batch_sz = x.shape[0]\n        nrows = min(batch_sz,\n                    batch_limit) if batch_limit is not None else batch_sz\n        return nrows\n\n    def get_plot_ncols(self, **kwargs) -> int:\n        x = kwargs['x']\n        nb_img_channels = x.shape[1]\n        ncols = len(self.get_channel_display_groups(nb_img_channels))\n        return ncols\n\n    def get_plot_params(self, **kwargs) -> dict:\n        nrows = self.get_plot_nrows(**kwargs)\n        ncols = self.get_plot_ncols(**kwargs)\n        params = {\n            'fig_args': {\n                'constrained_layout': True,\n                'figsize': np.array((self.scale * ncols, self.scale * nrows)),\n            },\n            'subplot_args': {\n                'nrows': nrows,\n                'ncols': ncols,\n                'squeeze': False\n            },\n            'plot_xyz_args': [{} for _ in range(nrows)]\n        }\n        return params\n", "path": "rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py"}], "after_files": [{"content": "from typing import (TYPE_CHECKING, Sequence, Optional, List, Dict, Union,\n                    Tuple, Any)\nfrom abc import ABC, abstractmethod\n\nimport numpy as np\nimport torch\nfrom torch import Tensor\nimport albumentations as A\nfrom torch.utils.data import DataLoader\nimport matplotlib.pyplot as plt\n\nfrom rastervision.pipeline.file_system import make_dir\nfrom rastervision.pytorch_learner.utils import (\n    deserialize_albumentation_transform, validate_albumentation_transform,\n    MinMaxNormalize)\nfrom rastervision.pytorch_learner.learner_config import (\n    RGBTuple,\n    ChannelInds,\n    ensure_class_colors,\n    validate_channel_display_groups,\n    get_default_channel_display_groups,\n)\n\nif TYPE_CHECKING:\n    from torch.utils.data import Dataset\n    from matplotlib.figure import Figure\n\n\nclass Visualizer(ABC):\n    \"\"\"Base class for plotting samples from computer vision PyTorch Datasets.\"\"\"\n\n    scale: float = 3.\n\n    def __init__(self,\n                 class_names: List[str],\n                 class_colors: Optional[List[Union[str, RGBTuple]]] = None,\n                 transform: Optional[Dict] = A.to_dict(MinMaxNormalize()),\n                 channel_display_groups: Optional[Union[Dict[\n                     str, ChannelInds], Sequence[ChannelInds]]] = None):\n        \"\"\"Constructor.\n\n        Args:\n            class_names: names of classes\n            class_colors: Colors used to display classes. Can be color 3-tuples\n                in list form.\n            transform: An Albumentations transform serialized as a dict that\n                will be applied to each image before it is plotted. Mainly useful\n                for undoing any data transformation that you do not want included in\n                the plot, such as normalization. The default value will shift and scale\n                the image so the values range from 0.0 to 1.0 which is the expected range\n                for the plotting function. This default is useful for cases where the\n                values after normalization are close to zero which makes the plot\n                difficult to see.\n            channel_display_groups: Groups of image channels to display together as a\n                subplot when plotting the data and predictions.\n                Can be a list or tuple of groups (e.g. [(0, 1, 2), (3,)]) or a\n                dict containing title-to-group mappings\n                (e.g. {\"RGB\": [0, 1, 2], \"IR\": [3]}),\n                where each group is a list or tuple of channel indices and\n                title is a string that will be used as the title of the subplot\n                for that group.\n        \"\"\"\n        self.class_names = class_names\n        self.class_colors = ensure_class_colors(self.class_names, class_colors)\n        self.transform = validate_albumentation_transform(transform)\n        self._channel_display_groups = validate_channel_display_groups(\n            channel_display_groups)\n\n    @abstractmethod\n    def plot_xyz(self,\n                 axs,\n                 x: Tensor,\n                 y: Sequence,\n                 z: Optional[Sequence] = None,\n                 plot_title: bool = True):\n        \"\"\"Plot image, ground truth labels, and predicted labels.\n\n        Args:\n            axs: matplotlib axes on which to plot\n            x: image\n            y: ground truth labels\n            z: optional predicted labels\n        \"\"\"\n        pass\n\n    def plot_batch(self,\n                   x: Tensor,\n                   y: Sequence,\n                   output_path: Optional[str] = None,\n                   z: Optional[Sequence] = None,\n                   batch_limit: Optional[int] = None,\n                   show: bool = False):\n        \"\"\"Plot a whole batch in a grid using plot_xyz.\n\n        Args:\n            x: batch of images\n            y: ground truth labels\n            output_path: local path where to save plot image\n            z: optional predicted labels\n            batch_limit: optional limit on (rendered) batch size\n        \"\"\"\n        params = self.get_plot_params(\n            x=x, y=y, z=z, output_path=output_path, batch_limit=batch_limit)\n        if params['subplot_args']['nrows'] == 0:\n            return\n\n        if x.ndim == 4:\n            fig, axs = plt.subplots(**params['fig_args'],\n                                    **params['subplot_args'])\n            plot_xyz_args = params['plot_xyz_args']\n            self._plot_batch(fig, axs, plot_xyz_args, x, y=y, z=z)\n        elif x.ndim == 5:\n            # If a temporal dimension is present, we divide the figure into\n            # multiple subfigures--one for each batch. Then, in each subfigure,\n            # we plot all timesteps as if they were a single batch. To\n            # delineate the boundary b/w batch items, we adopt the convention\n            # of only displaying subplot titles once per batch (above the first\n            # row in each batch).\n            batch_sz, T, *_ = x.shape\n            params['fig_args']['figsize'][1] *= T\n            fig = plt.figure(**params['fig_args'])\n            subfigs = fig.subfigures(\n                nrows=batch_sz, ncols=1, hspace=0.0, squeeze=False)\n            subfig_axs = [\n                subfig.subplots(\n                    nrows=T, ncols=params['subplot_args']['ncols'])\n                for subfig in subfigs.flat\n            ]\n            for i, axs in enumerate(subfig_axs):\n                plot_xyz_args = [\n                    dict(params['plot_xyz_args'][i]) for _ in range(T)\n                ]\n                plot_xyz_args[0]['plot_title'] = True\n                for args in plot_xyz_args[1:]:\n                    args['plot_title'] = False\n                _x = x[i]\n                _y = [y[i]] * T\n                _z = None if z is None else [z[i]] * T\n                self._plot_batch(fig, axs, plot_xyz_args, _x, y=_y, z=_z)\n        else:\n            raise ValueError('Expected x to have 4 or 5 dims, but found '\n                             f'x.shape: {x.shape}')\n\n        if show:\n            plt.show()\n        if output_path is not None:\n            make_dir(output_path, use_dirname=True)\n            fig.savefig(output_path, bbox_inches='tight', pad_inches=0.2)\n\n        plt.close(fig)\n\n    def _plot_batch(\n            self,\n            fig: 'Figure',\n            axs: Sequence,\n            plot_xyz_args: List[dict],\n            x: Tensor,\n            y: Optional[Sequence] = None,\n            z: Optional[Sequence] = None,\n    ):\n        # (N, c, h, w) --> (N, h, w, c)\n        x = x.permute(0, 2, 3, 1)\n\n        # apply transform, if given\n        if self.transform is not None:\n            tf = deserialize_albumentation_transform(self.transform)\n            imgs = [tf(image=img)['image'] for img in x.numpy()]\n            x = torch.from_numpy(np.stack(imgs))\n\n        for i, row_axs in enumerate(axs):\n            _z = None if z is None else z[i]\n            self.plot_xyz(row_axs, x[i], y[i], z=_z, **plot_xyz_args[i])\n\n    def get_channel_display_groups(\n            self, nb_img_channels: int\n    ) -> Union[Dict[str, ChannelInds], Sequence[ChannelInds]]:\n        # The default channel_display_groups object depends on the number of\n        # channels in the image. This number is not known when the Visualizer\n        # is constructed which is why it needs to be created later.\n        if self._channel_display_groups is not None:\n            return self._channel_display_groups\n        return get_default_channel_display_groups(nb_img_channels)\n\n    def get_collate_fn(self) -> Optional[callable]:\n        \"\"\"Returns a custom collate_fn to use in DataLoader.\n\n        None is returned if default collate_fn should be used.\n\n        See https://pytorch.org/docs/stable/data.html#working-with-collate-fn\n        \"\"\"\n        return None\n\n    def get_batch(self, dataset: 'Dataset', batch_sz: int = 4,\n                  **kwargs) -> Tuple[Tensor, Any]:\n        \"\"\"Generate a batch from a dataset.\n\n        This is a convenience method for generating a batch of data to plot.\n\n        Args:\n            dataset (Dataset): A Pytorch Datset.\n            batch_sz (int): Batch size. Defaults to 4.\n            **kwargs: Extra args for :class:`~torch.utils.data.DataLoader`.\n\n        Returns:\n            Tuple[Tensor, Any]: (x, y) tuple where x is images and y is labels.\n        \"\"\"\n        collate_fn = self.get_collate_fn()\n        dl = DataLoader(dataset, batch_sz, collate_fn=collate_fn, **kwargs)\n        try:\n            x, y = next(iter(dl))\n        except StopIteration:\n            raise ValueError('dataset did not return a batch')\n\n        return x, y\n\n    def get_plot_nrows(self, **kwargs) -> int:\n        x = kwargs['x']\n        batch_limit = kwargs.get('batch_limit')\n        batch_sz = x.shape[0]\n        nrows = min(batch_sz,\n                    batch_limit) if batch_limit is not None else batch_sz\n        return nrows\n\n    def get_plot_ncols(self, **kwargs) -> int:\n        x = kwargs['x']\n        nb_img_channels = x.shape[1]\n        ncols = len(self.get_channel_display_groups(nb_img_channels))\n        return ncols\n\n    def get_plot_params(self, **kwargs) -> dict:\n        nrows = self.get_plot_nrows(**kwargs)\n        ncols = self.get_plot_ncols(**kwargs)\n        params = {\n            'fig_args': {\n                'constrained_layout': True,\n                'figsize': np.array((self.scale * ncols, self.scale * nrows)),\n            },\n            'subplot_args': {\n                'nrows': nrows,\n                'ncols': ncols,\n                'squeeze': False\n            },\n            'plot_xyz_args': [{} for _ in range(nrows)]\n        }\n        return params\n", "path": "rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py"}]}