Skip to content

Commit 102d851

Browse files
committed
Merge remote-tracking branch 'origin/main' into favyen/20260424-eurocrops
2 parents ede4758 + 8f26741 commit 102d851

16 files changed

Lines changed: 93 additions & 36 deletions

File tree

data/rslearn_dataset_configs/config_worldcereal.json

Lines changed: 56 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,19 @@
66
"bands": [
77
"tc-annual-temporarycrops-classification"
88
],
9-
"dtype": "float32"
9+
"dtype": "float32",
10+
"nodata_vals": [
11+
255.0
12+
]
1013
}
1114
],
1215
"data_source": {
1316
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
1417
"init_args": {
1518
"worldcereal_dir": "source_data/worldcereal/"
19+
},
20+
"query_config": {
21+
"space_mode": "SINGLE_COMPOSITE"
1622
}
1723
},
1824
"resampling_method": "nearest",
@@ -24,13 +30,19 @@
2430
"bands": [
2531
"tc-maize-main-irrigation-classification"
2632
],
27-
"dtype": "float32"
33+
"dtype": "float32",
34+
"nodata_vals": [
35+
255.0
36+
]
2837
}
2938
],
3039
"data_source": {
3140
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
3241
"init_args": {
3342
"worldcereal_dir": "source_data/worldcereal/"
43+
},
44+
"query_config": {
45+
"space_mode": "SINGLE_COMPOSITE"
3446
}
3547
},
3648
"resampling_method": "nearest",
@@ -42,13 +54,19 @@
4254
"bands": [
4355
"tc-maize-main-maize-classification"
4456
],
45-
"dtype": "float32"
57+
"dtype": "float32",
58+
"nodata_vals": [
59+
255.0
60+
]
4661
}
4762
],
4863
"data_source": {
4964
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
5065
"init_args": {
5166
"worldcereal_dir": "source_data/worldcereal/"
67+
},
68+
"query_config": {
69+
"space_mode": "SINGLE_COMPOSITE"
5270
}
5371
},
5472
"resampling_method": "nearest",
@@ -60,13 +78,19 @@
6078
"bands": [
6179
"tc-maize-second-irrigation-classification"
6280
],
63-
"dtype": "float32"
81+
"dtype": "float32",
82+
"nodata_vals": [
83+
255.0
84+
]
6485
}
6586
],
6687
"data_source": {
6788
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
6889
"init_args": {
6990
"worldcereal_dir": "source_data/worldcereal/"
91+
},
92+
"query_config": {
93+
"space_mode": "SINGLE_COMPOSITE"
7094
}
7195
},
7296
"resampling_method": "nearest",
@@ -78,13 +102,19 @@
78102
"bands": [
79103
"tc-maize-second-maize-classification"
80104
],
81-
"dtype": "float32"
105+
"dtype": "float32",
106+
"nodata_vals": [
107+
255.0
108+
]
82109
}
83110
],
84111
"data_source": {
85112
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
86113
"init_args": {
87114
"worldcereal_dir": "source_data/worldcereal/"
115+
},
116+
"query_config": {
117+
"space_mode": "SINGLE_COMPOSITE"
88118
}
89119
},
90120
"resampling_method": "nearest",
@@ -96,13 +126,19 @@
96126
"bands": [
97127
"tc-springcereals-springcereals-classification"
98128
],
99-
"dtype": "float32"
129+
"dtype": "float32",
130+
"nodata_vals": [
131+
255.0
132+
]
100133
}
101134
],
102135
"data_source": {
103136
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
104137
"init_args": {
105138
"worldcereal_dir": "source_data/worldcereal/"
139+
},
140+
"query_config": {
141+
"space_mode": "SINGLE_COMPOSITE"
106142
}
107143
},
108144
"resampling_method": "nearest",
@@ -114,13 +150,19 @@
114150
"bands": [
115151
"tc-wintercereals-irrigation-classification"
116152
],
117-
"dtype": "float32"
153+
"dtype": "float32",
154+
"nodata_vals": [
155+
255.0
156+
]
118157
}
119158
],
120159
"data_source": {
121160
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
122161
"init_args": {
123162
"worldcereal_dir": "source_data/worldcereal/"
163+
},
164+
"query_config": {
165+
"space_mode": "SINGLE_COMPOSITE"
124166
}
125167
},
126168
"resampling_method": "nearest",
@@ -132,13 +174,19 @@
132174
"bands": [
133175
"tc-wintercereals-wintercereals-classification"
134176
],
135-
"dtype": "float32"
177+
"dtype": "float32",
178+
"nodata_vals": [
179+
255.0
180+
]
136181
}
137182
],
138183
"data_source": {
139184
"class_path": "rslearn.data_sources.worldcereal.WorldCereal",
140185
"init_args": {
141186
"worldcereal_dir": "source_data/worldcereal/"
187+
},
188+
"query_config": {
189+
"space_mode": "SINGLE_COMPOSITE"
142190
}
143191
},
144192
"resampling_method": "nearest",

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/cdl.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from rslearn.data_sources import Item
99
from rslearn.dataset import Dataset, Window
1010
from rslearn.utils.mp import star_imap_unordered
11+
from rslearn.utils.raster_array import RasterArray
1112
from upath import UPath
1213

1314
from olmoearth_pretrain.data.constants import Modality, TimeSpan
@@ -46,7 +47,7 @@ def convert_cdl(window: Window, olmoearth_path: UPath) -> None:
4647
raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands)
4748
image = GEOTIFF_RASTER_FORMAT.decode_raster(
4849
raster_dir, window.projection, window.bounds
49-
)
50+
).get_chw_array()
5051

5152
# Skip if there are any background/nodata.
5253
if image.min() == 0:
@@ -64,7 +65,7 @@ def convert_cdl(window: Window, olmoearth_path: UPath) -> None:
6465
path=dst_fname.parent,
6566
projection=window.projection,
6667
bounds=window.bounds,
67-
array=image,
68+
raster=RasterArray(chw_array=image),
6869
fname=dst_fname.name,
6970
)
7071
metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/era5.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from rslearn.data_sources import Item
1313
from rslearn.dataset import Dataset, Window
1414
from rslearn.utils.mp import star_imap_unordered
15+
from rslearn.utils.raster_array import RasterArray
1516
from rslearn.utils.raster_format import GeotiffRasterFormat
1617
from upath import UPath
1718

@@ -72,7 +73,7 @@ def convert_era5(window: Window, olmoearth_path: UPath) -> None:
7273
raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands, group_idx)
7374
image = raster_format.decode_raster(
7475
raster_dir, window.projection, window.bounds
75-
)
76+
).get_chw_array()
7677

7778
year_images.append(image)
7879
year_time_ranges.append(time_range)
@@ -119,7 +120,7 @@ def convert_era5(window: Window, olmoearth_path: UPath) -> None:
119120
path=year_dst_fname.parent,
120121
projection=window.projection,
121122
bounds=window.bounds,
122-
array=year_stacked_image,
123+
raster=RasterArray(chw_array=year_stacked_image),
123124
fname=year_dst_fname.name,
124125
)
125126
year_metadata_fname = get_modality_temp_meta_fname(
@@ -155,7 +156,7 @@ def convert_era5(window: Window, olmoearth_path: UPath) -> None:
155156
path=two_week_dst_fname.parent,
156157
projection=window.projection,
157158
bounds=window.bounds,
158-
array=two_week_image,
159+
raster=RasterArray(chw_array=two_week_image),
159160
fname=two_week_dst_fname.name,
160161
)
161162
two_week_metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/era5_10.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from rslearn.data_sources import Item
1313
from rslearn.dataset import Dataset, Window
1414
from rslearn.utils.mp import star_imap_unordered
15+
from rslearn.utils.raster_array import RasterArray
1516
from rslearn.utils.raster_format import GeotiffRasterFormat
1617
from upath import UPath
1718

@@ -80,7 +81,7 @@ def convert_era5(window: Window, olmoearth_path: UPath) -> None:
8081
raster_dir = window.get_raster_dir(LAYER_NAME, band_set.bands, group_idx)
8182
image = raster_format.decode_raster(
8283
raster_dir, adjusted_projection, adjusted_bounds
83-
)
84+
).get_chw_array()
8485

8586
year_images.append(image)
8687
year_time_ranges.append(time_range)
@@ -133,7 +134,7 @@ def convert_era5(window: Window, olmoearth_path: UPath) -> None:
133134
path=year_dst_fname.parent,
134135
projection=adjusted_projection,
135136
bounds=adjusted_bounds,
136-
array=year_stacked_image,
137+
raster=RasterArray(chw_array=year_stacked_image),
137138
fname=year_dst_fname.name,
138139
)
139140
year_metadata_fname = get_modality_temp_meta_fname(
@@ -169,7 +170,7 @@ def convert_era5(window: Window, olmoearth_path: UPath) -> None:
169170
path=two_week_dst_fname.parent,
170171
projection=adjusted_projection,
171172
bounds=adjusted_bounds,
172-
array=two_week_image,
173+
raster=RasterArray(chw_array=two_week_image),
173174
fname=two_week_dst_fname.name,
174175
)
175176
two_week_metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/eurocrops.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from rslearn.dataset import Dataset, Window
2424
from rslearn.utils.geometry import flatten_shape
2525
from rslearn.utils.mp import star_imap_unordered
26+
from rslearn.utils.raster_array import RasterArray
2627
from rslearn.utils.vector_format import GeojsonVectorFormat
2728
from upath import UPath
2829

@@ -193,7 +194,7 @@ def transform(coords: npt.NDArray) -> npt.NDArray:
193194
path=out_fname.parent,
194195
projection=window.projection,
195196
bounds=window.bounds,
196-
array=array,
197+
raster=RasterArray(chw_array=array),
197198
fname=out_fname.name,
198199
)
199200

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/gse.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def convert_gse(window: Window, olmoearth_path: UPath) -> None:
6868
path=dst_fname.parent,
6969
projection=window.projection,
7070
bounds=window.bounds,
71-
array=image,
71+
raster=image,
7272
fname=dst_fname.name,
7373
)
7474
metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/multitemporal_raster.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from rslearn.data_sources import Item
1010
from rslearn.dataset import Window
1111
from rslearn.utils.geometry import PixelBounds, Projection
12+
from rslearn.utils.raster_array import RasterArray
1213
from upath import UPath
1314

1415
from olmoearth_pretrain.data.constants import BandSet, ModalitySpec, TimeSpan
@@ -157,7 +158,7 @@ def convert_freq(
157158
)
158159
image = GEOTIFF_RASTER_FORMAT.decode_raster(
159160
raster_dir, adjusted_projection, adjusted_bounds
160-
)
161+
).get_chw_array()
161162
expected_image_size = band_set.get_expected_image_size(
162163
window_metadata.get_resolution_factor()
163164
)
@@ -206,7 +207,7 @@ def convert_freq(
206207
path=dst_fname.parent,
207208
projection=adjusted_projection,
208209
bounds=adjusted_bounds,
209-
array=stacked_image,
210+
raster=RasterArray(chw_array=stacked_image),
210211
fname=dst_fname.name,
211212
)
212213

@@ -288,7 +289,7 @@ def convert_monthly(
288289

289290
image = GEOTIFF_RASTER_FORMAT.decode_raster(
290291
raster_dir, adjusted_projection, adjusted_bounds
291-
)
292+
).get_chw_array()
292293
expected_image_size = band_set.get_expected_image_size(
293294
modality.tile_resolution_factor
294295
)
@@ -337,7 +338,7 @@ def convert_monthly(
337338
path=dst_fname.parent,
338339
projection=adjusted_projection,
339340
bounds=adjusted_bounds,
340-
array=stacked_image,
341+
raster=RasterArray(chw_array=stacked_image),
341342
fname=dst_fname.name,
342343
)
343344

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/naip.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def convert_naip(window: Window, olmoearth_path: UPath) -> None:
6868
path=dst_fname.parent,
6969
projection=window.projection,
7070
bounds=window.bounds,
71-
array=image,
71+
raster=image,
7272
fname=dst_fname.name,
7373
)
7474
metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/naip_10.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ def convert_naip(window: Window, olmoearth_path: UPath) -> None:
7878
path=dst_fname.parent,
7979
projection=adjusted_projection,
8080
bounds=adjusted_bounds,
81-
array=image,
81+
raster=image,
8282
fname=dst_fname.name,
8383
)
8484
metadata_fname = get_modality_temp_meta_fname(

olmoearth_pretrain/dataset_creation/rslearn_to_olmoearth/rasterize_openstreetmap.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from rasterio.crs import CRS
1414
from rslearn.utils.geometry import Projection
1515
from rslearn.utils.mp import star_imap_unordered
16+
from rslearn.utils.raster_array import RasterArray
1617
from upath import UPath
1718

1819
from olmoearth_pretrain.data.constants import Modality, TimeSpan
@@ -211,7 +212,7 @@ def transform(coords: npt.NDArray) -> npt.NDArray:
211212
path=out_fname.parent,
212213
projection=Projection(crs, OUTPUT_RESOLUTION, -OUTPUT_RESOLUTION),
213214
bounds=bounds,
214-
array=array,
215+
raster=RasterArray(chw_array=array),
215216
fname=out_fname.name,
216217
)
217218

0 commit comments

Comments
 (0)