Misc. small fixes and improvements

rastervision_core/rastervision/core/box.py

@@ -448,29 +448,26 @@ def from_dict(cls, d: dict) -> 'Self':
         return cls(d['ymin'], d['xmin'], d['ymax'], d['xmax'])
 
     @staticmethod
-    def filter_by_aoi(windows: list['Self'],
+    def filter_by_aoi(windows: Sequence['Box'],
                       aoi_polygons: list[Polygon],
-                      within: bool = True) -> list['Self']:
+                      within: bool = True) -> tuple[list['Box'], list[int]]:
         """Filters windows by a list of AOI polygons
 
         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.
         """
-        # 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
-
-        out = [w for w in windows if keep_window(w.to_shapely())]
-        return out
+        aoi: Polygon | MultiPolygon = unary_union(aoi_polygons)
+        keep_window = aoi.contains if within else aoi.intersects
+        inds = [
+            i for i, w in enumerate(windows) if keep_window(w.to_shapely())
+        ]
+        filtered_windows = [windows[i] for i in inds]
+        return filtered_windows, inds
 
     @staticmethod
-    def within_aoi(window: 'Self',
+    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)
@@ -479,7 +476,7 @@ def within_aoi(window: 'Self',
         return out
 
     @staticmethod
-    def intersects_aoi(window: 'Self',
+    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)

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

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Sequence
 import os
+from os.path import join
 import subprocess
 import logging
 
@@ -12,7 +13,7 @@
 
 from rastervision.pipeline.file_system.utils import (
     file_to_json, get_local_path, get_tmp_dir, make_dir, upload_or_copy,
-    download_if_needed)
+    download_if_needed, uri_to_vsi_path)
 from rastervision.core.box import Box
 
 if TYPE_CHECKING:
@@ -141,8 +142,10 @@ def build_vrt(vrt_path: str, image_uris: list[str]) -> None:
     """
     log.info('Building VRT...')
     cmd = ['gdalbuildvrt', vrt_path]
-    cmd.extend(image_uris)
-    subprocess.run(cmd)
+    image_uris_vsi = [uri_to_vsi_path(uri) for uri in image_uris]
+    cmd.extend(image_uris_vsi)
+    make_dir(vrt_path, use_dirname=True)
+    subprocess.run(cmd, env=os.environ)
 
 
 def download_and_build_vrt(image_uris: list[str],
@@ -161,7 +164,7 @@ def download_and_build_vrt(image_uris: list[str],
     """
     if not stream:
         image_uris = [download_if_needed(uri) for uri in image_uris]
-    vrt_path = os.path.join(vrt_dir, 'index.vrt')
+    vrt_path = join(vrt_dir, 'index.vrt')
     build_vrt(vrt_path, image_uris)
     return vrt_path