Fix bugs related to extent-cropping -- contd.

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

@@ -1,14 +1,13 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional, overload, Tuple
+from typing import Any, Optional, overload, Tuple
 
+import numpy as np
 from shapely.ops import transform
 from shapely.geometry.base import BaseGeometry
+from shapely.affinity import translate
 
 from rastervision.core.box import Box
 
-if TYPE_CHECKING:
-    import numpy as np
-
 
 class CRSTransformer(ABC):
     """Transforms map points in some CRS into pixel coordinates.
@@ -53,18 +52,13 @@ def map_to_pixel(self, inp, bbox: Optional[Box] = None):
             inp: (x, y) tuple or Box or shapely geometry in map coordinates.
                 If tuple, x and y can be single values or array-like.
             bbox: If the extent of the associated RasterSource is constrained
-                via a bbox, it can be passed here to get an output Box that is
-                compatible with the RasterSource's get_chip(). In other words,
-                the output Box will be in coordinates of the bbox rather than
-                the full extent of the data source of the RasterSource. Only
-                supported if ``inp`` is a :class:`.Box`. Defaults to None.
+                via a bbox, it can be passed here to get an output that is in
+                coordinates of the bbox rather than the full extent of the data
+                source of the RasterSource. Defaults to None.
 
         Returns:
             Coordinate-transformed input in the same format.
         """
-        if bbox is not None and not isinstance(inp, Box):
-            raise NotImplementedError(
-                'bbox is only supported if inp is a Box.')
         if isinstance(inp, Box):
             box_in = inp
             ymin, xmin, ymax, xmax = box_in
@@ -78,9 +72,19 @@ def map_to_pixel(self, inp, bbox: Optional[Box] = None):
             geom_in = inp
             geom_out = transform(
                 lambda x, y, z=None: self._map_to_pixel((x, y)), geom_in)
+            if bbox is not None:
+                xmin, ymin = bbox.xmin, bbox.ymin
+                geom_out = translate(geom_out, xoff=-xmin, yoff=-ymin)
             return geom_out
         elif len(inp) == 2:
-            return self._map_to_pixel(inp)
+            out = self._map_to_pixel(inp)
+            out_x, out_y = out
+            out = (np.array(out_x), np.array(out_y))
+            if bbox is not None:
+                xmin, ymin = bbox.xmin, bbox.ymin
+                out_x, out_y = out
+                out = (out_x - xmin, out_y - ymin)
+            return out
         else:
             raise TypeError(
                 'Input must be 2-tuple or Box or shapely geometry.')
@@ -114,17 +118,14 @@ def pixel_to_map(self, inp, bbox: Optional[Box] = None):
             inp: (x, y) tuple or Box or shapely geometry in pixel coordinates.
                 If tuple, x and y can be single values or array-like.
             bbox: If the extent of the associated RasterSource is constrained
-                via a bbox, it can be passed here so that the box is
+                via a bbox, it can be passed here so that the input is
                 interpreted to be in coordinates of the bbox rather than the
-                full extent of the data source of the RasterSource. Only
-                supported if ``inp`` is a :class:`.Box`. Defaults to None.
+                full extent of the data source of the RasterSource.
+                Defaults to None.
 
         Returns:
             Coordinate-transformed input in the same format.
         """
-        if bbox is not None and not isinstance(inp, Box):
-            raise NotImplementedError(
-                'bbox is only supported if inp is a Box.')
         if isinstance(inp, Box):
             box_in = inp
             if bbox is not None:
@@ -136,11 +137,21 @@ def pixel_to_map(self, inp, bbox: Optional[Box] = None):
             return box_out
         elif isinstance(inp, BaseGeometry):
             geom_in = inp
+            if bbox is not None:
+                xmin, ymin = bbox.xmin, bbox.ymin
+                geom_in = translate(geom_in, xoff=xmin, yoff=ymin)
             geom_out = transform(
                 lambda x, y, z=None: self._pixel_to_map((x, y)), geom_in)
             return geom_out
         elif len(inp) == 2:
-            return self._pixel_to_map(inp)
+            if bbox is not None:
+                xmin, ymin = bbox.xmin, bbox.ymin
+                inp_x, inp_y = inp
+                inp = (inp_x + xmin, inp_y + ymin)
+            out = self._pixel_to_map(inp)
+            out_x, out_y = out
+            out = (np.array(out_x), np.array(out_y))
+            return out
         else:
             raise TypeError(
                 'Input must be 2-tuple or Box or shapely geometry.')

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

@@ -49,7 +49,7 @@ def __repr__(self) -> str:
         if len(image_crs_str) > 70:
             image_crs_str = image_crs_str[:70] + '...'
 
-        map_crs_str = str(self.image_crs)
+        map_crs_str = str(self.map_crs)
         if len(map_crs_str) > 70:
             map_crs_str = map_crs_str[:70] + '...'
 

rastervision_core/rastervision/core/data/raster_source/raster_source.py

@@ -15,7 +15,7 @@ class ChannelOrderError(Exception):
     def __init__(self, channel_order: List[int], num_channels_raw: int):
         self.channel_order = channel_order
         self.num_channels_raw = num_channels_raw
-        msg = (f'The channel_order ({channel_order}) contains an'
+        msg = (f'The channel_order ({channel_order}) contains an '
                f'index >= num_channels_raw ({num_channels_raw}).')
         super().__init__(msg)
 

rastervision_core/rastervision/core/data/raster_source/rasterio_source.py

@@ -272,7 +272,7 @@ def num_channels(self) -> int:
         return self._num_channels
 
     @property
-    def dtype(self) -> Tuple[int, int, int]:
+    def dtype(self) -> np.dtype:
         if self._dtype is None:
             self._set_info_from_chip()
         return self._dtype
@@ -332,7 +332,8 @@ def get_chip(self,
                 channel_order=[2, 1, 0], then using bands=[0] will return the
                 B-channel. Defaults to None.
             out_shape (Optional[Tuple[int, ...]], optional): (hieght, width) of
-            the output chip. If None, no resizing is done. Defaults to None.
+                the output chip. If None, no resizing is done.
+                Defaults to None.
 
         Returns:
             np.ndarray: A chip of shape (height, width, channels).

rastervision_core/rastervision/core/data/scene.py

@@ -1,8 +1,9 @@
-from typing import TYPE_CHECKING, Any, Optional, Tuple
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple
 
-from rastervision.core.data.utils import match_bboxes
+from rastervision.core.data.utils import match_bboxes, geoms_to_bbox_coords
 
 if TYPE_CHECKING:
+    from shapely.geometry.base import BaseGeometry
     from rastervision.core.box import Box
     from rastervision.core.data import (RasterSource, LabelSource, LabelStore)
 
@@ -15,7 +16,7 @@ def __init__(self,
                  raster_source: 'RasterSource',
                  label_source: Optional['LabelSource'] = None,
                  label_store: Optional['LabelStore'] = None,
-                 aoi_polygons: Optional[list] = None):
+                 aoi_polygons: Optional[List['BaseGeometry']] = None):
         """Constructor.
 
         During initialization, ``Scene`` attempts to set the extents of the
@@ -42,14 +43,26 @@ def __init__(self,
 
         if aoi_polygons is None:
             self.aoi_polygons = []
+            self.aoi_polygons_bbox_coords = []
         else:
-            self.aoi_polygons = aoi_polygons
+            bbox = self.raster_source.bbox
+            bbox_geom = bbox.to_shapely()
+            self.aoi_polygons = [
+                p for p in aoi_polygons if p.intersects(bbox_geom)
+            ]
+            self.aoi_polygons_bbox_coords = list(
+                geoms_to_bbox_coords(self.aoi_polygons, bbox))
 
     @property
     def extent(self) -> 'Box':
         """Extent of the associated :class:`.RasterSource`."""
         return self.raster_source.extent
 
+    @property
+    def bbox(self) -> 'Box':
+        """Bounding box applied to the source data."""
+        return self.raster_source.bbox
+
     def __getitem__(self, key: Any) -> Tuple[Any, Any]:
         x = self.raster_source[key]
         if self.label_source is not None:

rastervision_core/rastervision/core/data/utils/geojson.py

@@ -3,7 +3,11 @@
 from copy import deepcopy
 
 from shapely.geometry import shape, mapping
+from shapely.affinity import translate
 from tqdm.auto import tqdm
+import geopandas as gpd
+
+from rastervision.core.box import Box
 
 if TYPE_CHECKING:
     from rastervision.core.data.crs_transformer import CRSTransformer
@@ -329,17 +333,21 @@ def all_geoms_valid(geojson: dict):
     return all(g.is_valid for g in geoms)
 
 
-def get_polygons_from_uris(
-        uris: Union[str, List[str]],
-        crs_transformer: 'CRSTransformer') -> List['BaseGeometry']:
+def get_polygons_from_uris(uris: Union[str, List[str]],
+                           crs_transformer: 'CRSTransformer',
+                           bbox: Optional['Box'] = None,
+                           map_coords: bool = False) -> List['BaseGeometry']:
     """Load and return polygons (in pixel coords) from one or more URIs."""
 
     # use local imports to avoid circular import problems
     from rastervision.core.data import GeoJSONVectorSource
 
     source = GeoJSONVectorSource(
-        uris=uris, ignore_crs_field=True, crs_transformer=crs_transformer)
-    polygons = source.get_geoms()
+        uris=uris,
+        ignore_crs_field=True,
+        crs_transformer=crs_transformer,
+        bbox=bbox)
+    polygons = source.get_geoms(to_map_coords=map_coords)
     return polygons
 
 
@@ -348,3 +356,37 @@ def merge_geojsons(geojsons: Iterable[dict]) -> dict:
     features = sum([g.get('features', []) for g in geojsons], [])
     geojson_merged = features_to_geojson(features)
     return geojson_merged
+
+
+def geojson_to_geodataframe(geojson: dict) -> gpd.GeoDataFrame:
+    df = gpd.GeoDataFrame.from_features(geojson)
+    if len(df) == 0 and 'geometry' not in df.columns:
+        df.loc[:, 'geometry'] = []
+    return df
+
+
+def get_geodataframe_extent(gdf: gpd.GeoDataFrame) -> Box:
+    envelope = gdf.unary_union.envelope
+    extent = Box.from_shapely(envelope).to_int()
+    return extent
+
+
+def get_geojson_extent(geojson: dict) -> Box:
+    gdf = geojson_to_geodataframe(geojson)
+    extent = get_geodataframe_extent(gdf)
+    return extent
+
+
+def filter_geojson_to_window(geojson: dict, window: Box) -> dict:
+    gdf = geojson_to_geodataframe(geojson)
+    window_geom = window.to_shapely()
+    gdf: gpd.GeoDataFrame = gdf[gdf.intersects(window_geom)]
+    out_geojson = gdf._to_geo()
+    return out_geojson
+
+
+def geoms_to_bbox_coords(geoms: Iterable['BaseGeometry'],
+                         bbox: Box) -> Iterator['BaseGeometry']:
+    xmin, ymin = bbox.xmin, bbox.ymin
+    out = (translate(p, xoff=-xmin, yoff=-ymin) for p in geoms)
+    return out

rastervision_core/rastervision/core/data/vector_source/geojson_vector_source.py

@@ -1,6 +1,7 @@
-from typing import TYPE_CHECKING, List, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from rastervision.pipeline.file_system import download_if_needed, file_to_json
+from rastervision.core.box import Box
 from rastervision.core.data.vector_source.vector_source import VectorSource
 from rastervision.core.data.utils import listify_uris, merge_geojsons
 
@@ -15,7 +16,8 @@ def __init__(self,
                  uris: Union[str, List[str]],
                  crs_transformer: 'CRSTransformer',
                  vector_transformers: List['VectorTransformer'] = [],
-                 ignore_crs_field: bool = False):
+                 ignore_crs_field: bool = False,
+                 bbox: Optional[Box] = None):
         """Constructor.
 
         Args:
@@ -29,11 +31,15 @@ def __init__(self,
                 assume WGS84 (EPSG:4326) coords. Only WGS84 is supported
                 currently. If False, and the file contains a CRS, will throw an
                 exception on read. Defaults to False.
+            bbox (Optional[Box]): User-specified crop of the extent. If None,
+                the full extent available in the source file is used.
         """
         self.uris = listify_uris(uris)
         self.ignore_crs_field = ignore_crs_field
         super().__init__(
-            crs_transformer, vector_transformers=vector_transformers)
+            crs_transformer,
+            vector_transformers=vector_transformers,
+            bbox=bbox)
 
     def _get_geojson(self) -> dict:
         geojsons = [self._get_geojson_single(uri) for uri in self.uris]