don't modify the scripts from main

Author: favyen2

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/cdl.py

@@ -8,6 +8,7 @@
 from rslearn.data_sources import Item
 from rslearn.dataset import Dataset, Window
 from rslearn.utils.mp import star_imap_unordered
+from rslearn.utils.raster_array import RasterArray
 from upath import UPath
 
 from olmoearth_pretrain.data.constants import Modality, TimeSpan
@@ -44,12 +45,12 @@ def convert_cdl(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.CDL.band_sets) == 1
     band_set = Modality.CDL.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = GEOTIFF_RASTER_FORMAT.decode_raster(
+    image = GEOTIFF_RASTER_FORMAT.decode_raster(
         raster_dir, window.projection, window.bounds
-    )
+    ).get_chw_array()
 
     # Skip if there are any background/nodata.
-    if raster.array.min() == 0:
+    if image.min() == 0:
         return
 
     dst_fname = get_modality_fname(
@@ -64,7 +65,7 @@ def convert_cdl(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=RasterArray(chw_array=image),
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/gse.py

@@ -53,7 +53,7 @@ def convert_gse(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.GSE.band_sets) == 1
     band_set = Modality.GSE.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = GEOTIFF_RASTER_FORMAT.decode_raster(
+    image = GEOTIFF_RASTER_FORMAT.decode_raster(
         raster_dir, window.projection, window.bounds
     )
     dst_fname = get_modality_fname(
@@ -68,7 +68,7 @@ def convert_gse(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=image,
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/naip.py

@@ -55,7 +55,7 @@ def convert_naip(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.NAIP.band_sets) == 1
     band_set = Modality.NAIP.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = raster_format.decode_raster(raster_dir, window.projection, window.bounds)
+    image = raster_format.decode_raster(raster_dir, window.projection, window.bounds)
     dst_fname = get_modality_fname(
         olmoearth_path,
         Modality.NAIP,
@@ -68,7 +68,7 @@ def convert_naip(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=image,
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/naip_10.py

@@ -63,7 +63,7 @@ def convert_naip(window: Window, olmoearth_path: UPath) -> None:
         Modality.NAIP_10, band_set, window.projection, window.bounds
     )
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = raster_format.decode_raster(
+    image = raster_format.decode_raster(
         raster_dir, adjusted_projection, adjusted_bounds
     )
     dst_fname = get_modality_fname(
@@ -78,7 +78,7 @@ def convert_naip(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=adjusted_projection,
         bounds=adjusted_bounds,
-        raster=raster,
+        raster=image,
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/srtm.py

@@ -38,7 +38,7 @@ def convert_srtm(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.SRTM.band_sets) == 1
     band_set = Modality.SRTM.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = GEOTIFF_RASTER_FORMAT.decode_raster(
+    image = GEOTIFF_RASTER_FORMAT.decode_raster(
         raster_dir, window.projection, window.bounds
     )
     dst_fname = get_modality_fname(
@@ -53,7 +53,7 @@ def convert_srtm(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=image,
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/worldcover.py

@@ -38,7 +38,7 @@ def convert_worldcover(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.WORLDCOVER.band_sets) == 1
     band_set = Modality.WORLDCOVER.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = GEOTIFF_RASTER_FORMAT.decode_raster(
+    image = GEOTIFF_RASTER_FORMAT.decode_raster(
         raster_dir, window.projection, window.bounds
     )
     dst_fname = get_modality_fname(
@@ -53,7 +53,7 @@ def convert_worldcover(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=image,
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/worldpop.py

@@ -8,6 +8,7 @@
 import tqdm
 from rslearn.dataset import Dataset, Window
 from rslearn.utils.mp import star_imap_unordered
+from rslearn.utils.raster_array import RasterArray
 from upath import UPath
 
 from olmoearth_pretrain.data.constants import Modality, TimeSpan
@@ -38,16 +39,16 @@ def convert_worldpop(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.WORLDPOP.band_sets) == 1
     band_set = Modality.WORLDPOP.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = GEOTIFF_RASTER_FORMAT.decode_raster(
+    image = GEOTIFF_RASTER_FORMAT.decode_raster(
         raster_dir, window.projection, window.bounds
     ).get_chw_array()
 
     # Clip population count to 0. NODATA is -99999 and includes locations that are
     # mapped as "unsettled" but really that is 0 population.
-    raster.array[raster.array < 0] = 0
+    image[image < 0] = 0
     # Skip areas that are fully nodata since it is more likely to be wrong and also is
     # a less useful target.
-    if raster.array.max() == 0:
+    if image.max() == 0:
         return
 
     dst_fname = get_modality_fname(
@@ -62,7 +63,7 @@ def convert_worldpop(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=RasterArray(chw_array=image),
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/wri_canopy_height_map.py

@@ -9,6 +9,7 @@
 import tqdm
 from rslearn.dataset import Dataset, Window
 from rslearn.utils.mp import star_imap_unordered
+from rslearn.utils.raster_array import RasterArray
 from upath import UPath
 
 from olmoearth_pretrain.data.constants import Modality, TimeSpan
@@ -40,15 +41,15 @@ def convert_chm(window: Window, olmoearth_path: UPath) -> None:
     assert len(Modality.WRI_CANOPY_HEIGHT_MAP.band_sets) == 1
     band_set = Modality.WRI_CANOPY_HEIGHT_MAP.band_sets[0]
     raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
-    raster = GEOTIFF_RASTER_FORMAT.decode_raster(
+    image = GEOTIFF_RASTER_FORMAT.decode_raster(
         raster_dir, window.projection, window.bounds
     ).get_chw_array()
 
     # Skip areas with any nodata (255).
-    if raster.array.max() == 255:
+    if image.max() == 255:
         return
     # Also skip if there are not enough positive pixels.
-    if np.count_nonzero(raster.array) / raster.array.size < 0.2:
+    if np.count_nonzero(image) / image.size < 0.2:
         return
 
     dst_fname = get_modality_fname(
@@ -63,7 +64,7 @@ def convert_chm(window: Window, olmoearth_path: UPath) -> None:
         path=dst_fname.parent,
         projection=window.projection,
         bounds=window.bounds,
-        raster=raster,
+        raster=RasterArray(chw_array=image),
         fname=dst_fname.name,
     )
     metadata_fname = get_modality_temp_meta_fname(