feat(earthview): Sentinel_2 support

earthview.py

@@ -46,7 +46,10 @@ def load_dataset(subset, dataset="satellogic/EarthView", split="train", shards =
     if shards is None:
         data_files = None
     else:
-        data_files = [f"{path}/{split}-{shard:05d}-of-{nshards:05d}.parquet" for shard in shards]
+        if subset == "sentinel_2":
+            data_files = [f"{path}/sentinel_2-{shard//10}/{split}-{shard % 10:05d}-of-00010.parquet" for shard in shards]
+        else:
+            data_files = [f"{path}/{split}-{shard:05d}-of-{nshards:05d}.parquet" for shard in shards]
         data_files = {split: data_files}
 
     ds = _load_dataset(
@@ -127,12 +130,17 @@ def item_to_images(subset, item):
         ]
         count = len(item["10m"])
     elif subset == "sentinel_2":
-        for channel in ['10m', '20m', '40m', 'rgb']: #, 'scl', 'metadata']:
+        for channel in ['10m', '20m', 'rgb', 'scl']: #, '40m']:
             data = item[channel]
             count = len(data)
-            data = np.asarray(data).astype("uint8")
-            images = [Image.fromarray(data[i,:,:,:].transpose(1,2,0)) for i in range(count)]
+            data = np.asarray(data).astype("uint8").transpose(0,2,3,1)
+            if channel == "20m":
+                data = data[:,:,:,[0,2,4]]
+            mode = "L" if channel in ["10m", "scl"] else "RGB"
+            images = [Image.fromarray(data[i].squeeze(), mode=mode) for i in range(count)]
             item[channel] = images
+        for field in ["solarAngles", "tileGeometry", "viewIncidenceAngles"]:
+            item["metadata"][field] = [json.loads(s) for s in item["metadata"][field]]
     elif subset == "neon":
         item["rgb"] = [
             Image.fromarray(image.transpose(1,2,0))

utils.py

@@ -16,23 +16,27 @@ def lat_lon_mid_pixel(item, subset: str):
     if subset == "satellogic":
         crs = metadata["crs"][0]
         bounds_crs = metadata["bounds"]
+    elif subset == "neon":
+        crs = metadata["epsg"]
+        bounds_crs = metadata["bounds"]
     elif subset == "sentinel_1":
-        crs = metadata["crs"]
-        bounds_crs = metadata["coordinates"][0]
+        geometry = metadata
+    elif subset == "sentinel_2":
+        geometry = metadata["tileGeometry"][0]
+    else:
+        raise ValueError("subset not known")
+
+    if subset.startswith("sentinel_"):
+        crs = geometry["crs"]
+        bounds_crs = geometry["coordinates"][0]
         assert len(bounds_crs) == 5
-        # bounds are a polygon with same first & last vertex
         bounds_crs = (
             bounds_crs[0][0],
             bounds_crs[0][1],
             bounds_crs[2][0],
             bounds_crs[2][1],
         )
-    elif subset == "neon":
-        crs = metadata["epsg"]
-        bounds_crs = metadata["bounds"]
 
-    else:
-        raise ValueError("subset not known")
 
     bounds = pyproj.Transformer.from_crs(crs, "EPSG:4326").transform_bounds(*bounds_crs)
     # dumb average for now