Misc. fixes and improvements

Dockerfile

@@ -56,6 +56,10 @@ RUN mamba init
 RUN mamba install -y python=${PYTHON_VERSION}
 RUN python -m pip install --upgrade pip
 
+# env variable required by uv
+ENV CONDA_PREFIX=/opt/conda
+RUN pip install uv
+
 # We need to install GDAL first to install Rasterio on non-AMD64 architectures.
 # The Rasterio wheels contain GDAL in them, but they are only built for AMD64 now.
 RUN mamba update mamba -y && mamba install -y -c conda-forge gdal=3.6.3
@@ -111,7 +115,7 @@ WORKDIR /opt/src/
 
 COPY ./rastervision_pytorch_learner/requirements.txt /opt/src/pytorch-requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip cat pytorch-requirements.txt | sort | uniq > all-requirements.txt && \
-    pip install $(grep -ivE "^\s*$|^#|rastervision_*" all-requirements.txt) && \
+    uv pip install $(grep -ivE "^\s*$|^#|rastervision_*" all-requirements.txt) && \
     rm all-requirements.txt
 
 COPY ./rastervision_aws_batch/requirements.txt /opt/src/batch-requirements.txt
@@ -136,7 +140,7 @@ RUN --mount=type=cache,target=/root/.cache/pip cat \
     /opt/src/sagemaker-requirements.txt \
     /opt/src/requirements-dev.txt \
     | sort | uniq > all-requirements.txt && \
-    pip install $(grep -ivE "^\s*$|^#|rastervision_*" all-requirements.txt) && \
+    uv pip install $(grep -ivE "^\s*$|^#|rastervision_*" all-requirements.txt) && \
     rm all-requirements.txt
 
 #########################
@@ -147,7 +151,7 @@ COPY ./docs/requirements.txt /opt/src/docs/pandoc-requirements.txt
 
 # Install pandoc, needed for rendering notebooks
 # Get latest release link from here: https://github.com/jgm/pandoc/releases
-RUN --mount=type=cache,target=/root/.cache/pip pip install -r docs/pandoc-requirements.txt && \
+RUN --mount=type=cache,target=/root/.cache/pip uv pip install -r docs/pandoc-requirements.txt && \
     wget https://github.com/jgm/pandoc/releases/download/3.1.12.2/pandoc-3.1.12.2-1-${TARGETARCH}.deb && \
     dpkg -i pandoc-3.1.12.2-1-${TARGETARCH}.deb && rm pandoc-3.1.12.2-1-${TARGETARCH}.deb
 
@@ -180,13 +184,13 @@ COPY ./rastervision_pytorch_backend/ /opt/src/rastervision_pytorch_backend/
 COPY ./rastervision_pytorch_learner/ /opt/src/rastervision_pytorch_learner/
 
 # needed for this image to be used by the AWS SageMaker PyTorch Estimator
-RUN pip install sagemaker_pytorch_training==2.8.1
+RUN uv pip install sagemaker_pytorch_training==2.8.1
 ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main
 
 # Install a onnxruntime-gpu version compatible with CUDA 12. Specifying
 # --extra-index-url in requirements.txt seems to cause problems with the
 # RTD build.
 RUN if [ "${TARGETARCH}" != "arm64" ]; then \
-    pip install --upgrade onnxruntime-gpu==1.17 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/; fi
+    uv pip install --upgrade onnxruntime-gpu==1.17 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/; fi
 
 CMD ["bash"]

docs/setup/aws.rst

@@ -8,7 +8,7 @@ Running on AWS
 Running on AWS EC2
 ------------------
 
-The simplest way to run Raster Vision on an AWS GPU is by starting a GPU-enabled EC2 instance such as a p3.2xlarge using the `Deep Learning AMI <https://aws.amazon.com/machine-learning/amis/>`_. We have tested this using the "Deep Learning AMI GPU PyTorch 1.11.0 (Ubuntu 20.04)" with id ``ami-0c968d7ef8a4b0c34``. After SSH'ing into the instance, Raster Vision can be installed with ``pip``, and code can be transferred to this instance with a tool such as ``rsync``.
+The simplest way to run Raster Vision on an AWS GPU is by starting a GPU-enabled EC2 instance such as a ``p3.2xlarge`` using a `Deep Learning AMI <https://aws.amazon.com/machine-learning/amis/>`_. After SSH'ing into the instance, Raster Vision can be installed with ``pip``, and code can be transferred to this instance with a tool such as ``rsync``.
 
 .. _aws batch setup:
 

docs/setup/index.rst

@@ -19,7 +19,7 @@ Installing via pip
 
 .. currentmodule:: rastervision
 
-You can directly install the library using ``pip`` (or ``pip3`` if you also have Python 2 installed).
+You can directly install the library using ``pip``.
 
 .. code-block:: console
 
@@ -101,15 +101,15 @@ The command above will attempt to install GDAL via ``pip``. If that fails, you c
 Docker Images
 -------------
 
-Using the Docker images published for Raster Vision makes it easy to use a fully set up environment. We have tested this with Docker 20, although you may be able to use a lower version.
+Using the Docker images published for Raster Vision makes it easy to use a fully set up environment.
 
-The images we publish include all plugins and dependencies for using Raster Vision with PyTorch and AWS. These are published to `quay.io/azavea/raster-vision <https://quay.io/repository/azavea/raster-vision>`_.  To run the container for the latest release, run:
+The images we publish include all plugins and dependencies for using Raster Vision with PyTorch and AWS. These are published to `quay.io/azavea/raster-vision <https://quay.io/repository/azavea/raster-vision>`_ (see the *tags* tab).  To run the container for the latest release, run:
 
 .. code-block:: console
 
    > docker run --rm -it quay.io/azavea/raster-vision:pytorch-{{ version }} /bin/bash
 
-There are also images with the `-latest` suffix for the latest commits on the ``master`` branch. You'll likely need to mount volumes and expose ports to make this container fully useful; see the `docker/run <{{ repo }}/docker/run>`_ script for an example usage.
+There are also images with the ``-latest`` suffix for the latest commits on the ``master`` branch. You'll likely need to mount volumes and expose ports to make this container fully useful; see the `docker/run <{{ repo }}/docker/run>`_ script for an example usage.
 
 You can also base your own Dockerfiles off the Raster Vision image to use with your own codebase. See :ref:`bootstrap` for more information.
 
@@ -140,16 +140,16 @@ To run a Bash console in the PyTorch Docker container use:
 
 This will mount the ``$RASTER_VISION_DATA_DIR`` local directory to to ``/opt/data/`` inside the container.
 
+.. note::
+
+    If you have built an ARM64 image, you should pass the ``--arm64`` flag to ``docker/run``.
+
 .. warning::
 
     Users running under WSL2 in Windows will need to unset the ``NAME`` environment variable. For example, instead of
     ``docker/run``, you would run ``NAME='' docker/run``. By default, WSL2 sets a ``NAME`` variable that matches the network
     name of your computer. This environment variable collides with a variable in the ``docker/run`` script.
 
-.. warning::
-
-    If you have built an ARM64 image, you should pass the ``--arm64`` flag to ``docker/run``.
-
 This script also has options for forwarding AWS credentials, and running Jupyter notebooks which can be seen below.
 
 .. code-block:: console

rastervision_core/rastervision/core/box.py

@@ -6,14 +6,15 @@
 
 import numpy as np
 from shapely.geometry import Polygon
+from shapely.ops import unary_union
 from rasterio.windows import Window as RioWindow
 
 from rastervision.pipeline.utils import repr_with_args
 
 NonNegInt = conint(ge=0)
 
 if TYPE_CHECKING:
-    pass
+    from shapely.geometry import MultiPolygon
 
 
 class BoxSizeError(ValueError):
@@ -459,28 +460,37 @@ def filter_by_aoi(windows: List['Box'],
 
         Args:
             within: if True, windows are only kept if they lie fully within an
-                AOI polygon. Otherwise, windows are kept if they intersect an AOI
-                polygon.
+                AOI polygon. Otherwise, windows are kept if they intersect an
+                AOI polygon.
         """
-        result = []
-        for window in windows:
-            w = window.to_shapely()
-            for polygon in aoi_polygons:
-                if ((within and w.within(polygon))
-                        or ((not within) and w.intersects(polygon))):
-                    result.append(window)
-                    break
+        # merge overlapping polygons, if any
+        aoi_polygons: Polygon | MultiPolygon = unary_union(aoi_polygons)
+
+        if within:
+            keep_window = aoi_polygons.contains
+        else:
+            keep_window = aoi_polygons.intersects
 
-        return result
+        out = [w for w in windows if keep_window(w.to_shapely())]
+        return out
+
+    @staticmethod
+    def within_aoi(window: 'Box',
+                   aoi_polygons: Polygon | List[Polygon]) -> bool:
+        """Check if window is within the union of given AOI polygons."""
+        aoi_polygons: Polygon | MultiPolygon = unary_union(aoi_polygons)
+        w = window.to_shapely()
+        out = aoi_polygons.contains(w)
+        return out
 
     @staticmethod
-    def within_aoi(window: 'Box', aoi_polygons: List[Polygon]) -> bool:
-        """Check if window is within a list of AOI polygons."""
+    def intersects_aoi(window: 'Box',
+                       aoi_polygons: Polygon | List[Polygon]) -> bool:
+        """Check if window intersects with the union of given AOI polygons."""
+        aoi_polygons: Polygon | MultiPolygon = unary_union(aoi_polygons)
         w = window.to_shapely()
-        for polygon in aoi_polygons:
-            if w.within(polygon):
-                return True
-        return False
+        out = aoi_polygons.intersects(w)
+        return out
 
     def __contains__(self, query: Union['Box', Sequence]) -> bool:
         """Check if box or point is contained within this box.

rastervision_core/rastervision/core/data/crs_transformer/rasterio_crs_transformer.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 from pyproj import Transformer
 
 import numpy as np
@@ -9,6 +9,9 @@
 from rastervision.core.data.crs_transformer import (CRSTransformer,
                                                     IdentityCRSTransformer)
 
+if TYPE_CHECKING:
+    from typing import Self
+
 
 class RasterioCRSTransformer(CRSTransformer):
     """Transformer for a RasterioRasterSource."""
@@ -124,8 +127,8 @@ def from_dataset(
         return cls(transform, image_crs, map_crs, **kwargs)
 
     @classmethod
-    def from_uri(cls, uri: str, map_crs: Optional[str] = 'epsg:4326', **kwargs
-                 ) -> Union[IdentityCRSTransformer, 'RasterioCRSTransformer']:
+    def from_uri(cls, uri: str, map_crs: Optional[str] = 'epsg:4326',
+                 **kwargs) -> 'Self':
         """Build from raster URI.
 
         Args:

rastervision_core/rastervision/core/data/label/chip_classification_labels.py

@@ -62,7 +62,7 @@ def __setitem__(self, window: Box,
 
     @classmethod
     def from_predictions(cls, windows: Iterable['Box'],
-                         predictions: Iterable[Any]) -> 'Labels':
+                         predictions: Iterable[Any]):
         """Overrid to convert predictions to (class_id, scores) pairs."""
         predictions = ((np.argmax(p), p) for p in predictions)
         return super().from_predictions(windows, predictions)

rastervision_core/rastervision/core/data/label/labels.py

@@ -1,9 +1,10 @@
 """Defines the abstract Labels class."""
 
 from typing import TYPE_CHECKING, Any, Iterable, List
-from abc import (ABC, abstractclassmethod, abstractmethod)
+from abc import ABC, abstractmethod
 
 if TYPE_CHECKING:
+    from typing import Self
     from shapely.geometry import Polygon
     from rastervision.core.box import Box
 
@@ -16,14 +17,14 @@ class Labels(ABC):
     """
 
     @abstractmethod
-    def __add__(self, other: 'Labels'):
+    def __add__(self, other: 'Self'):
         """Add labels to these labels.
 
         Returns a concatenation of this and the other labels.
         """
 
     @abstractmethod
-    def filter_by_aoi(self, aoi_polygons: List['Polygon']) -> 'Labels':
+    def filter_by_aoi(self, aoi_polygons: List['Polygon']) -> 'Self':
         """Return a copy of these labels filtered by given AOI polygons.
 
         Args:
@@ -39,8 +40,9 @@ def __eq__(self, other: 'Labels'):
     def __setitem__(self, key, value):
         pass
 
-    @abstractclassmethod
-    def make_empty(cls) -> 'Labels':
+    @classmethod
+    @abstractmethod
+    def make_empty(cls) -> 'Self':
         """Instantiate an empty instance of this class.
 
         Returns:
@@ -50,7 +52,7 @@ def make_empty(cls) -> 'Labels':
 
     @classmethod
     def from_predictions(cls, windows: Iterable['Box'],
-                         predictions: Iterable[Any]) -> 'Labels':
+                         predictions: Iterable[Any]) -> 'Self':
         """Instantiate from windows and their corresponding predictions.
 
         This makes no assumptions about the type or format of the predictions.