Merge pull request #1546 from azavea/lf/lightning

Add RV + Lightning notebook
azavea · Nov 22, 2022 · bc6357b · bc6357b
2 parents 0235559 + b40ca44
commit bc6357b
Show file tree

Hide file tree

Showing 7 changed files with 866 additions and 48 deletions.
diff --git a/docs/conf.py b/docs/conf.py
@@ -139,6 +139,7 @@ def setup(app: 'Sphinx') -> None:
     # The _images dir is under build/html. This looks brittle but using the
     # more natural img/tensorboard.png path does not work.
     'tutorials/train': '_images/tensorboard.png',
+    'tutorials/lightning_workflow': '_images/lightning-logo.png',
 }
 nbsphinx_prolog = r"""
 {% set docpath = env.doc2path(env.docname, base=False) %}
@@ -333,7 +334,7 @@ def setup(app: 'Sphinx') -> None:
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ['_static', 'img']
 
 # A list of CSS files. The entry must be a filename string or a tuple
 # containing the filename string and the attributes dictionary. The filename

diff --git a/docs/img/lightning-logo.png b/docs/img/lightning-logo.png
diff --git a/docs/img/tensorboard-lightning.png b/docs/img/tensorboard-lightning.png
diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -16,6 +16,7 @@ The following tutorials highlight many of the things you can do with Raster Visi
     scenes_and_aois
     visualize_data_samples
     train
+    lightning_workflow
     pred_and_eval_ss
     misc
 

diff --git a/docs/tutorials/lightning_workflow.ipynb b/docs/tutorials/lightning_workflow.ipynb
diff --git a/docs/tutorials/sampling_training_data.ipynb b/docs/tutorials/sampling_training_data.ipynb
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/dataset.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/dataset.py
@@ -1,5 +1,6 @@
 from typing import Union, Optional, Tuple, Any, TypeVar
 from typing_extensions import Literal
+import logging
 
 import numpy as np
 import albumentations as A
@@ -14,6 +15,8 @@
                                                             TF_TYPE_TO_TF_FUNC)
 from rastervision.pytorch_learner.dataset.utils import AoiSampler
 
+log = logging.getLogger(__name__)
+
 
 class AlbumentationsDataset(Dataset):
     """An adapter to use arbitrary datasets with albumentations transforms."""
@@ -28,8 +31,15 @@ def __init__(self,
 
         Args:
             orig_dataset (Any): An object with a __getitem__ and __len__.
-            transform (A.BasicTransform, optional): An Albumentations
-                transform. Defaults to None.
+            transform (A.BasicTransform, optional): Albumentations
+                transform to apply to the windows. Defaults to None.
+                Each transform in Albumentations takes images of type uint8, and
+                sometimes other data types. The data type requirements can be
+                seen at https://albumentations.ai/docs/api_reference/augmentations/transforms/ # noqa
+                If there is a mismatch between the data type of imagery and the
+                transform requirements, a RasterTransformer should be set
+                on the RasterSource that converts to uint8, such as
+                MinMaxTransformer or StatsTransformer.
             transform_type (TransformType): The type of transform so that its
                 inputs and outputs can be handled correctly. Defaults to
                 TransformType.noop.
@@ -51,7 +61,15 @@ def __init__(self,
 
     def __getitem__(self, key) -> Tuple[torch.Tensor, torch.Tensor]:
         val = self.orig_dataset[key]
-        x, y = self.transform(val)
+
+        try:
+            x, y = self.transform(val)
+        except Exception as exc:
+            log.warn(
+                'Many albumentations transforms require uint8 input. Therefore, we '
+                'recommend passing a MinMaxTransformer or StatsTransformer to the '
+                'RasterSource so the input will be converted to uint8.')
+            raise exc
 
         if self.normalize and np.issubdtype(x.dtype, np.unsignedinteger):
             max_val = np.iinfo(x.dtype).max
@@ -97,6 +115,13 @@ def __init__(self,
             scene (Scene): A Scene object.
             transform (Optional[A.BasicTransform], optional): Albumentations
                 transform to apply to the windows. Defaults to None.
+                Each transform in Albumentations takes images of type uint8, and
+                sometimes other data types. The data type requirements can be
+                seen at https://albumentations.ai/docs/api_reference/augmentations/transforms/ # noqa
+                If there is a mismatch between the data type of imagery and the
+                transform requirements, a RasterTransformer should be set
+                on the RasterSource that converts to uint8, such as
+                MinMaxTransformer or StatsTransformer.
             transform_type (Optional[TransformType], optional): Type of
                 transform. Defaults to None.
             normalize (bool, optional): If True, x is normalized to [0, 1]
@@ -163,6 +188,13 @@ def __init__(self,
                 paddiong is zero. Defaults to 'end'.
             transform (Optional[A.BasicTransform], optional): Albumentations
                 transform to apply to the windows. Defaults to None.
+                Each transform in Albumentations takes images of type uint8, and
+                sometimes other data types. The data type requirements can be
+                seen at https://albumentations.ai/docs/api_reference/augmentations/transforms/ # noqa
+                If there is a mismatch between the data type of imagery and the
+                transform requirements, a RasterTransformer should be set
+                on the RasterSource that converts to uint8, such as
+                MinMaxTransformer or StatsTransformer.
             transform_type (Optional[TransformType], optional): Type of
                 transform. Defaults to None.
         """
@@ -241,6 +273,13 @@ def __init__(self,
                 np.inf. Defaults to None.
             transform (Optional[A.BasicTransform], optional): Albumentations
                 transform to apply to the windows. Defaults to None.
+                Each transform in Albumentations takes images of type uint8, and
+                sometimes other data types. The data type requirements can be
+                seen at https://albumentations.ai/docs/api_reference/augmentations/transforms/ # noqa
+                If there is a mismatch between the data type of imagery and the
+                transform requirements, a RasterTransformer should be set
+                on the RasterSource that converts to uint8, such as
+                MinMaxTransformer or StatsTransformer.
             transform_type (Optional[TransformType], optional): Type of
                 transform. Defaults to None.
             max_sample_attempts (NonNegInt, optional): Max attempts when trying