Repository: azavea/raster-vision

Issue:
Add support for parsing STAC inputs
This has been discussed before in #873 and #962.

It seems like the simplest way would be to add helper functions for reading from STAC catalogs that can be used in the `get_config()` to extract images, AOIs, etc. We have been doing this on other projects, so it should be a simple matter of adding those functions to the repo.


Discussion:
A currently-broken function for reading STACs already exists in RV here: https://github.com/azavea/raster-vision/blob/master/rastervision_core/rastervision/core/utils/stac.py

Resolving patch:
diff --git a/rastervision_core/rastervision/core/utils/stac.py b/rastervision_core/rastervision/core/utils/stac.py
--- a/rastervision_core/rastervision/core/utils/stac.py
+++ b/rastervision_core/rastervision/core/utils/stac.py
@@ -1,13 +1,19 @@
+from typing import List, Optional
 from urllib.parse import urlparse
+import logging
 
 import boto3
-from pystac import STAC_IO
-from pystac import Catalog, LabelItem
+from pystac import StacIO, Catalog, Item
+from pystac.stac_io import DefaultStacIO
 from shapely.geometry import box
 
+from rastervision.pipeline.file_system import file_to_json
 
-def setup_stac_s3():
-    def my_read_method(uri):
+log = logging.getLogger(__name__)
+
+
+def setup_stac_io() -> None:
+    def read_method(uri: str):
         parsed = urlparse(uri)
         if parsed.scheme == 's3':
             bucket = parsed.netloc
@@ -16,9 +22,9 @@ def my_read_method(uri):
             obj = s3.Object(bucket, key)
             return obj.get()['Body'].read().decode('utf-8')
         else:
-            return STAC_IO.default_read_text_method(uri)
+            return DefaultStacIO.read_text_method(uri)
 
-    def my_write_method(uri, txt):
+    def write_method(uri: str, txt: str):
         parsed = urlparse(uri)
         if parsed.scheme == 's3':
             bucket = parsed.netloc
@@ -26,34 +32,93 @@ def my_write_method(uri, txt):
             s3 = boto3.resource('s3')
             s3.Object(bucket, key).put(Body=txt)
         else:
-            STAC_IO.default_write_text_method(uri, txt)
+            DefaultStacIO.write_text_method(uri, txt)
+
+    StacIO.read_text_method = read_method
+    StacIO.write_text_method = write_method
+
+
+def is_label_item(item: Item) -> bool:
+    """Resolve each extension schema into a dict, then check if it has the
+    title of "Label Extension".
+    """
+    for ext_schema_uri in item.stac_extensions:
+        schema = file_to_json(ext_schema_uri)
+        if schema['title'].lower() == 'label extension':
+            return True
+    return False
+
+
+def get_linked_image_item(label_item: Item) -> Optional[Item]:
+    """Find link in the item that has "rel" == "source" and return its
+    "target" item. If no such link, return None. If multiple such links,
+    raise an exception."""
+    links = [l for l in label_item.links if l.rel.lower() == 'source']
+    if len(links) == 0:
+        return None
+    elif len(links) > 1:
+        raise NotImplementedError()
+    image_item = links[0].resolve_stac_object().target
+    return image_item
+
+
+def parse_stac(stac_uri: str) -> List[dict]:
+    """Parse a STAC catalog JSON file to extract label URIs, images URIs,
+    and AOIs.
 
-    STAC_IO.read_text_method = my_read_method
-    STAC_IO.write_text_method = my_write_method
+    Note: This has been tested to be compatible with STAC version 1.0.0 but
+    not any other versions.
 
+    Args:
+        stac_uri (str): Path to the STAC catalog JSON file.
 
-def parse_stac(stac_uri):
-    setup_stac_s3()
+    Returns:
+        List[dict]: A lsit of dicts with keys: "label_uri", "image_uris",
+            "label_bbox", "image_bbox", "bboxes_intersect", and "aoi_geometry".
+            Each dict corresponds to one label item and its associated image
+            assets in the STAC catalog.
+    """
+    setup_stac_io()
     cat = Catalog.from_file(stac_uri)
+    version: str = cat.to_dict()['stac_version']
+
+    if not version.startswith('1.0'):
+        log.warning(f'Parsing is not guaranteed to work correctly for '
+                    f'STAC version != 1.0.*. Found version: {version}.')
+
     cat.make_all_asset_hrefs_absolute()
-    labels_uri = None
-    geotiff_uris = []
-    for item in cat.get_all_items():
-        if isinstance(item, LabelItem):
-            labels_uri = list(item.assets.values())[0].href
-            labels_box = box(*item.bbox)
-
-    # only use geotiffs that intersect with bbox of labels
-    for item in cat.get_all_items():
-        if not isinstance(item, LabelItem):
-            geotiff_uri = list(item.assets.values())[0].href
-            geotiff_box = box(*item.bbox)
-            if labels_box.intersects(geotiff_box):
-                geotiff_uri = geotiff_uri.replace('%7C', '|')
-                geotiff_uris.append(geotiff_uri)
-
-    if not labels_uri:
-        raise ValueError('Unable to read labels URI from STAC.')
-    if not geotiff_uris:
-        raise ValueError('Unable to read GeoTIFF URIs from STAC.')
-    return labels_uri, labels_box, geotiff_uris
+
+    label_items = [item for item in cat.get_all_items() if is_label_item(item)]
+    image_items = [get_linked_image_item(item) for item in label_items]
+
+    if len(label_items) == 0:
+        raise ValueError('Unable to find any label items in STAC catalog.')
+
+    out = []
+    for label_item, image_item in zip(label_items, image_items):
+        label_uri: str = list(label_item.assets.values())[0].href
+        label_bbox = box(*label_item.bbox)
+        aoi_geometry: Optional[dict] = label_item.geometry
+
+        if image_item is not None:
+            image_assets = [
+                asset for asset in image_item.get_assets().values()
+                if 'image' in asset.media_type
+            ]
+            image_uris = [asset.href for asset in image_assets]
+            image_bbox = box(*image_item.bbox)
+            bboxes_intersect = label_bbox.intersects(image_bbox)
+        else:
+            image_uris = []
+            image_bbox = None
+            bboxes_intersect = False
+
+        out.append({
+            'label_uri': label_uri,
+            'image_uris': image_uris,
+            'label_bbox': label_bbox,
+            'image_bbox': image_bbox,
+            'bboxes_intersect': bboxes_intersect,
+            'aoi_geometry': aoi_geometry
+        })
+    return out
diff --git a/rastervision_pytorch_backend/rastervision/pytorch_backend/examples/utils.py b/rastervision_pytorch_backend/rastervision/pytorch_backend/examples/utils.py
--- a/rastervision_pytorch_backend/rastervision/pytorch_backend/examples/utils.py
+++ b/rastervision_pytorch_backend/rastervision/pytorch_backend/examples/utils.py
@@ -1,7 +1,10 @@
+from typing import List
 import csv
 from io import StringIO
-import tempfile
 import os
+from os.path import join
+from pathlib import Path
+from tempfile import TemporaryDirectory
 
 import rasterio
 from shapely.strtree import STRtree
@@ -11,9 +14,10 @@
 from rastervision.core import Box
 from rastervision.core.data import (RasterioCRSTransformer,
                                     GeoJSONVectorSourceConfig)
-from rastervision.pipeline.file_system import (file_to_str, file_exists,
-                                               get_local_path, upload_or_copy,
-                                               make_dir, json_to_file)
+from rastervision.core.utils.stac import parse_stac
+from rastervision.pipeline.file_system import (
+    file_to_str, file_exists, get_local_path, upload_or_copy, make_dir,
+    json_to_file, download_if_needed, unzip)
 from rastervision.aws_s3 import S3FileSystem
 
 
@@ -39,7 +43,7 @@ def crop_image(image_uri, window, crop_uri):
     rasterio_window = window.rasterio_format()
     im = im_dataset.read(window=rasterio_window)
 
-    with tempfile.TemporaryDirectory() as tmp_dir:
+    with TemporaryDirectory() as tmp_dir:
         crop_path = get_local_path(crop_uri, tmp_dir)
         make_dir(crop_path, use_dirname=True)
 
@@ -139,3 +143,42 @@ def p2m(x, y, z=None):
         finally:
             os.environ.clear()
             os.environ.update(old_environ)
+
+
+def read_stac(uri: str) -> List[dict]:
+    """Parse the contents of a STAC catalog (downloading it first, if
+    remote). If the uri is a zip file, unzip it, find catalog.json inside it
+    and parse that.
+
+    Args:
+        uri (str): Either a URI to a STAC catalog JSON file or a URI to a zip
+            file containing a STAC catalog JSON file.
+
+    Raises:
+        FileNotFoundError: If catalog.json is not found inside the zip file.
+        Exception: If multiple catalog.json's are found inside the zip file.
+
+    Returns:
+        List[dict]: A lsit of dicts with keys: "label_uri", "image_uris",
+            "label_bbox", "image_bbox", "bboxes_intersect", and "aoi_geometry".
+            Each dict corresponds to one label item and its associated image
+            assets in the STAC catalog.
+    """
+    uri_path = Path(uri)
+    is_zip = uri_path.suffix.lower() == '.zip'
+
+    with TemporaryDirectory() as tmp_dir:
+        catalog_path = download_if_needed(uri, tmp_dir)
+        if not is_zip:
+            return parse_stac(catalog_path)
+        zip_path = catalog_path
+        unzip_path = join(tmp_dir, uri_path.stem)
+        unzip(zip_path, target_dir=unzip_path)
+        catalog_paths = list(Path(unzip_path).glob('**/catalog.json'))
+        if len(catalog_paths) == 0:
+            raise FileNotFoundError(f'Unable to find "catalog.json" in {uri}.')
+        elif len(catalog_paths) > 1:
+            raise Exception(f'More than one "catalog.json" found in '
+                            f'{uri}.')
+        catalog_path = str(catalog_paths[0])
+        return parse_stac(catalog_path)


Test patch:
diff --git a/tests/core/utils/test_stac.py b/tests/core/utils/test_stac.py
new file mode 100644
--- /dev/null
+++ b/tests/core/utils/test_stac.py
@@ -0,0 +1,19 @@
+from typing import Callable
+import unittest
+
+from rastervision.core.utils.stac import setup_stac_io
+
+
+class TestStac(unittest.TestCase):
+    def assertNoError(self, fn: Callable, msg: str = ''):
+        try:
+            fn()
+        except Exception:
+            self.fail(msg)
+
+    def test_setup_stac_io(self):
+        self.assertNoError(setup_stac_io)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/data_files/catalog.zip b/tests/data_files/catalog.zip
new file mode 100644
Binary files /dev/null and b/tests/data_files/catalog.zip differ
diff --git a/tests/pytorch_backend/examples/test_utils.py b/tests/pytorch_backend/examples/test_utils.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_backend/examples/test_utils.py
@@ -0,0 +1,38 @@
+import unittest
+
+from shapely.geometry import Polygon
+from rastervision.pytorch_backend.examples.utils import read_stac
+
+from tests import data_file_path
+
+
+class TestUtils(unittest.TestCase):
+    def test_read_stac(self):
+        expected_keys = {
+            'label_uri': str,
+            'image_uris': list,
+            'label_bbox': Polygon,
+            'image_bbox': (type(None), Polygon),
+            'bboxes_intersect': bool,
+            'aoi_geometry': (type(None), dict)
+        }
+        zip_path = data_file_path('catalog.zip')
+        out = read_stac(zip_path)
+
+        # check for correctness of format
+        self.assertIsInstance(out, list)
+        for o in out:
+            self.assertIsInstance(o, dict)
+            for k, v in o.items():
+                self.assertTrue(k in expected_keys)
+                self.assertIsInstance(v, expected_keys[k])
+            for uri in o['image_uris']:
+                self.assertIsInstance(uri, str)
+
+        # check for correctness of content (WRT the test catalog)
+        self.assertEqual(len(out), 1)
+        self.assertEqual(len(out[0]['image_uris']), 1)
+
+
+if __name__ == '__main__':
+    unittest.main()