Skip to content

rslearn.data_sources.local_files

local_files

Data source for raster or vector data in local files.

Importer

Bases: Generic[ItemType]

An abstract base class for importing data from local files.

Source code in rslearn/data_sources/local_files.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class Importer(Generic[ItemType]):
    """An abstract base class for importing data from local files."""

    def list_items(self, src_dir: UPath) -> list[ItemType]:
        """Extract a list of Items from the source directory.

        Args:
            src_dir: the source directory.
        """
        raise NotImplementedError

    def ingest_item(
        self,
        tile_store: TileStoreWithLayer,
        item: ItemType,
        cur_geometries: list[STGeometry],
    ) -> None:
        """Ingest the specified local file item.

        Args:
            tile_store: the tile store to ingest the data into.
            item: the Item to ingest
            cur_geometries: the geometries where the item is needed.
        """
        raise NotImplementedError

list_items 

list_items(src_dir: UPath) -> list[ItemType]

Extract a list of Items from the source directory.

Parameters:

Name Type Description Default
src_dir UPath

the source directory.

 required
Source code in rslearn/data_sources/local_files.py 
62
63
64
65
66
67
68
def list_items(self, src_dir: UPath) -> list[ItemType]:
    """Extract a list of Items from the source directory.

    Args:
        src_dir: the source directory.
    """
    raise NotImplementedError

ingest_item 

ingest_item(tile_store: TileStoreWithLayer, item: ItemType, cur_geometries: list[STGeometry]) -> None

Ingest the specified local file item.

Parameters:

Name Type Description Default
tile_store TileStoreWithLayer

the tile store to ingest the data into.

 required
item ItemType

the Item to ingest

 required
cur_geometries list[STGeometry]

the geometries where the item is needed.

 required
Source code in rslearn/data_sources/local_files.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def ingest_item(
    self,
    tile_store: TileStoreWithLayer,
    item: ItemType,
    cur_geometries: list[STGeometry],
) -> None:
    """Ingest the specified local file item.

    Args:
        tile_store: the tile store to ingest the data into.
        item: the Item to ingest
        cur_geometries: the geometries where the item is needed.
    """
    raise NotImplementedError

RasterItemSpec

Representation of configuration that directly specifies the available items.

Source code in rslearn/data_sources/local_files.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
class RasterItemSpec:
    """Representation of configuration that directly specifies the available items."""

    def __init__(
        self,
        fnames: list[str],
        bands: list[list[str]] | None = None,
        name: str | None = None,
    ):
        """Create a new RasterItemSpec.

        Args:
            fnames: the list of image files in this item.
            bands: the bands provided by each of the image files.
            name: what the item should be named
        """
        self.fnames = fnames
        self.bands = bands
        self.name = name

    def serialize(self) -> dict[str, Any]:
        """Serializes the RasterItemSpec to a JSON-encodable dictionary."""
        return {
            "fnames": [str(path) for path in self.fnames],
            "bands": self.bands,
            "name": self.name,
        }

    @staticmethod
    def deserialize(d: dict[str, Any]) -> "RasterItemSpec":
        """Deserializes a RasterItemSpec from a JSON-decoded dictionary."""
        return RasterItemSpec(
            fnames=[s for s in d["fnames"]],
            bands=d["bands"],
            name=d["name"],
        )

serialize 

serialize() -> dict[str, Any]

Serializes the RasterItemSpec to a JSON-encodable dictionary.

Source code in rslearn/data_sources/local_files.py 
106
107
108
109
110
111
112
def serialize(self) -> dict[str, Any]:
    """Serializes the RasterItemSpec to a JSON-encodable dictionary."""
    return {
        "fnames": [str(path) for path in self.fnames],
        "bands": self.bands,
        "name": self.name,
    }

deserialize staticmethod 

deserialize(d: dict[str, Any]) -> RasterItemSpec

Deserializes a RasterItemSpec from a JSON-decoded dictionary.

Source code in rslearn/data_sources/local_files.py 
114
115
116
117
118
119
120
121
@staticmethod
def deserialize(d: dict[str, Any]) -> "RasterItemSpec":
    """Deserializes a RasterItemSpec from a JSON-decoded dictionary."""
    return RasterItemSpec(
        fnames=[s for s in d["fnames"]],
        bands=d["bands"],
        name=d["name"],
    )

RasterItem

Bases: Item

An item corresponding to a local file.

Source code in rslearn/data_sources/local_files.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
class RasterItem(Item):
    """An item corresponding to a local file."""

    def __init__(
        self, name: str, geometry: STGeometry, src_dir: str, spec: RasterItemSpec
    ):
        """Creates a new RasterItem.

        Args:
            name: unique name of the item
            geometry: the spatial and temporal extent of the item
            src_dir: the source directory.
            spec: the RasterItemSpec that specifies the filename(s) and bands.
        """
        super().__init__(name, geometry)
        self.src_dir = src_dir
        self.spec = spec

    def serialize(self) -> dict:
        """Serializes the item to a JSON-encodable dictionary."""
        d = super().serialize()
        d["src_dir"] = str(self.src_dir)
        d["spec"] = self.spec.serialize()
        return d

    @staticmethod
    def deserialize(d: dict) -> "RasterItem":
        """Deserializes an item from a JSON-decoded dictionary."""
        item = super(RasterItem, RasterItem).deserialize(d)
        src_dir = d["src_dir"]
        spec = RasterItemSpec.deserialize(d["spec"])
        return RasterItem(
            name=item.name, geometry=item.geometry, src_dir=src_dir, spec=spec
        )

serialize 

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/local_files.py 
142
143
144
145
146
147
def serialize(self) -> dict:
    """Serializes the item to a JSON-encodable dictionary."""
    d = super().serialize()
    d["src_dir"] = str(self.src_dir)
    d["spec"] = self.spec.serialize()
    return d

deserialize staticmethod 

deserialize(d: dict) -> RasterItem

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/local_files.py 
149
150
151
152
153
154
155
156
157
@staticmethod
def deserialize(d: dict) -> "RasterItem":
    """Deserializes an item from a JSON-decoded dictionary."""
    item = super(RasterItem, RasterItem).deserialize(d)
    src_dir = d["src_dir"]
    spec = RasterItemSpec.deserialize(d["spec"])
    return RasterItem(
        name=item.name, geometry=item.geometry, src_dir=src_dir, spec=spec
    )

VectorItem

Bases: Item

An item corresponding to a local file.

Source code in rslearn/data_sources/local_files.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
class VectorItem(Item):
    """An item corresponding to a local file."""

    def __init__(self, name: str, geometry: STGeometry, path_uri: str):
        """Creates a new LocalFileItem.

        Args:
            name: unique name of the item
            geometry: the spatial and temporal extent of the item
            path_uri: URI representation of the path of this file
        """
        super().__init__(name, geometry)
        self.path_uri = path_uri

    def serialize(self) -> dict:
        """Serializes the item to a JSON-encodable dictionary."""
        d = super().serialize()
        d["path_uri"] = self.path_uri
        return d

    @staticmethod
    def deserialize(d: dict) -> "VectorItem":
        """Deserializes an item from a JSON-decoded dictionary."""
        item = super(VectorItem, VectorItem).deserialize(d)
        return VectorItem(
            name=item.name, geometry=item.geometry, path_uri=d["path_uri"]
        )

serialize 

serialize() -> dict

Serializes the item to a JSON-encodable dictionary.

Source code in rslearn/data_sources/local_files.py 
174
175
176
177
178
def serialize(self) -> dict:
    """Serializes the item to a JSON-encodable dictionary."""
    d = super().serialize()
    d["path_uri"] = self.path_uri
    return d

deserialize staticmethod 

deserialize(d: dict) -> VectorItem

Deserializes an item from a JSON-decoded dictionary.

Source code in rslearn/data_sources/local_files.py 
180
181
182
183
184
185
186
@staticmethod
def deserialize(d: dict) -> "VectorItem":
    """Deserializes an item from a JSON-decoded dictionary."""
    item = super(VectorItem, VectorItem).deserialize(d)
    return VectorItem(
        name=item.name, geometry=item.geometry, path_uri=d["path_uri"]
    )

RasterImporter

Bases: Importer

An Importer for raster data.

Source code in rslearn/data_sources/local_files.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
class RasterImporter(Importer):
    """An Importer for raster data."""

    def __init__(self, item_specs: list[RasterItemSpec] | None = None):
        """Create a new RasterImporter.

        Args:
            item_specs: the specs to specify the raster items directly. If None, the
                raster items are automatically detected from the files in the source
                directory.
        """
        self.item_specs = item_specs

    def list_items(self, src_dir: UPath) -> list[Item]:
        """Extract a list of Items from the source directory.

        Args:
            src_dir: the source directory.
        """
        item_specs: list[RasterItemSpec]

        # See if user has provided the item specs directly.
        if self.item_specs is not None:
            item_specs = self.item_specs
        else:
            # Otherwise we need to list files and assume each one is separate.
            # And we'll need to autodetect the bands later.
            item_specs = []
            file_paths = src_dir.glob("**/*.*")
            for path in file_paths:
                if _should_skip_source_file(path):
                    continue

                spec = RasterItemSpec(
                    fnames=[get_relative_suffix(src_dir, path)], bands=None
                )
                item_specs.append(spec)

        items: list[Item] = []
        for spec in item_specs:
            # Get geometry from the first raster file.
            # We assume files are readable with rasterio.
            fname = join_upath(src_dir, spec.fnames[0])
            with open_rasterio_upath_reader(fname) as src:
                gcps, gcp_crs = src.gcps
                if src.crs is None and len(gcps) > 0:
                    # Fit an affine transform from the GCPs and project the
                    # image corners to get image's geographic extent.
                    transform = from_gcps(gcps)
                    corners = [
                        transform * (0, 0),
                        transform * (src.width, 0),
                        transform * (src.width, src.height),
                        transform * (0, src.height),
                    ]
                    xs = [c[0] for c in corners]
                    ys = [c[1] for c in corners]
                    shp = shapely.box(min(xs), min(ys), max(xs), max(ys))
                    # We use 1 unit/pixel projection so it is compatible with
                    # the shape above. The actual resolution may differ, but
                    # it is okay since this geometry is only for item-window
                    # matching. Re-projection will be handled by the tile
                    # store, e.g. DefaultTileStore applies WarpedVRT.
                    projection = Projection(gcp_crs, 1, 1)
                    geometry = STGeometry(projection, shp, None)
                else:
                    crs = src.crs
                    left = src.transform.c
                    top = src.transform.f
                    x_resolution = src.transform.a
                    y_resolution = src.transform.e
                    start = (int(left / x_resolution), int(top / y_resolution))
                    shp = shapely.box(
                        start[0],
                        start[1],
                        start[0] + src.width,
                        start[1] + src.height,
                    )
                    projection = Projection(crs, x_resolution, y_resolution)
                    geometry = STGeometry(projection, shp, None)

            if spec.name:
                item_name = spec.name
            else:
                item_name = fname.name.split(".")[0]

            logger.debug(
                "RasterImporter.list_items: got bounds of %s: %s", item_name, geometry
            )
            items.append(RasterItem(item_name, geometry, str(src_dir), spec))

        logger.debug("RasterImporter.list_items: discovered %d items", len(items))
        return items

    def ingest_item(
        self,
        tile_store: TileStoreWithLayer,
        item: Item,
        cur_geometries: list[STGeometry],
    ) -> None:
        """Ingest the specified local file item.

        Args:
            tile_store: the tile store to ingest the data into.
            item: the RasterItem to ingest
            cur_geometries: the geometries where the item is needed.
        """
        assert isinstance(item, RasterItem)
        for file_idx, fname in enumerate(item.spec.fnames):
            fname_upath = join_upath(UPath(item.src_dir), fname)
            with open_rasterio_upath_reader(fname_upath) as src:
                if item.spec.bands:
                    bands = item.spec.bands[file_idx]
                else:
                    bands = [f"B{band_idx + 1}" for band_idx in range(src.count)]

            if tile_store.is_raster_ready(item, bands):
                continue
            tile_store.write_raster_file(
                item, bands, fname_upath, time_range=item.geometry.time_range
            )

list_items 

list_items(src_dir: UPath) -> list[Item]

Extract a list of Items from the source directory.

Parameters:

Name Type Description Default
src_dir UPath

the source directory.

 required
Source code in rslearn/data_sources/local_files.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
def list_items(self, src_dir: UPath) -> list[Item]:
    """Extract a list of Items from the source directory.

    Args:
        src_dir: the source directory.
    """
    item_specs: list[RasterItemSpec]

    # See if user has provided the item specs directly.
    if self.item_specs is not None:
        item_specs = self.item_specs
    else:
        # Otherwise we need to list files and assume each one is separate.
        # And we'll need to autodetect the bands later.
        item_specs = []
        file_paths = src_dir.glob("**/*.*")
        for path in file_paths:
            if _should_skip_source_file(path):
                continue

            spec = RasterItemSpec(
                fnames=[get_relative_suffix(src_dir, path)], bands=None
            )
            item_specs.append(spec)

    items: list[Item] = []
    for spec in item_specs:
        # Get geometry from the first raster file.
        # We assume files are readable with rasterio.
        fname = join_upath(src_dir, spec.fnames[0])
        with open_rasterio_upath_reader(fname) as src:
            gcps, gcp_crs = src.gcps
            if src.crs is None and len(gcps) > 0:
                # Fit an affine transform from the GCPs and project the
                # image corners to get image's geographic extent.
                transform = from_gcps(gcps)
                corners = [
                    transform * (0, 0),
                    transform * (src.width, 0),
                    transform * (src.width, src.height),
                    transform * (0, src.height),
                ]
                xs = [c[0] for c in corners]
                ys = [c[1] for c in corners]
                shp = shapely.box(min(xs), min(ys), max(xs), max(ys))
                # We use 1 unit/pixel projection so it is compatible with
                # the shape above. The actual resolution may differ, but
                # it is okay since this geometry is only for item-window
                # matching. Re-projection will be handled by the tile
                # store, e.g. DefaultTileStore applies WarpedVRT.
                projection = Projection(gcp_crs, 1, 1)
                geometry = STGeometry(projection, shp, None)
            else:
                crs = src.crs
                left = src.transform.c
                top = src.transform.f
                x_resolution = src.transform.a
                y_resolution = src.transform.e
                start = (int(left / x_resolution), int(top / y_resolution))
                shp = shapely.box(
                    start[0],
                    start[1],
                    start[0] + src.width,
                    start[1] + src.height,
                )
                projection = Projection(crs, x_resolution, y_resolution)
                geometry = STGeometry(projection, shp, None)

        if spec.name:
            item_name = spec.name
        else:
            item_name = fname.name.split(".")[0]

        logger.debug(
            "RasterImporter.list_items: got bounds of %s: %s", item_name, geometry
        )
        items.append(RasterItem(item_name, geometry, str(src_dir), spec))

    logger.debug("RasterImporter.list_items: discovered %d items", len(items))
    return items

ingest_item 

ingest_item(tile_store: TileStoreWithLayer, item: Item, cur_geometries: list[STGeometry]) -> None

Ingest the specified local file item.

Parameters:

Name Type Description Default
tile_store TileStoreWithLayer

the tile store to ingest the data into.

 required
item Item

the RasterItem to ingest

 required
cur_geometries list[STGeometry]

the geometries where the item is needed.

 required
Source code in rslearn/data_sources/local_files.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
def ingest_item(
    self,
    tile_store: TileStoreWithLayer,
    item: Item,
    cur_geometries: list[STGeometry],
) -> None:
    """Ingest the specified local file item.

    Args:
        tile_store: the tile store to ingest the data into.
        item: the RasterItem to ingest
        cur_geometries: the geometries where the item is needed.
    """
    assert isinstance(item, RasterItem)
    for file_idx, fname in enumerate(item.spec.fnames):
        fname_upath = join_upath(UPath(item.src_dir), fname)
        with open_rasterio_upath_reader(fname_upath) as src:
            if item.spec.bands:
                bands = item.spec.bands[file_idx]
            else:
                bands = [f"B{band_idx + 1}" for band_idx in range(src.count)]

        if tile_store.is_raster_ready(item, bands):
            continue
        tile_store.write_raster_file(
            item, bands, fname_upath, time_range=item.geometry.time_range
        )

VectorImporter

Bases: Importer

An Importer for vector data.

Source code in rslearn/data_sources/local_files.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
class VectorImporter(Importer):
    """An Importer for vector data."""

    # We need some buffer around GeoJSON bounds in case it just contains one point.
    item_buffer_epsilon = 1e-4

    def list_items(self, src_dir: UPath) -> list[Item]:
        """Extract a list of Items from the source directory.

        Args:
            src_dir: the source directory.
        """
        file_paths = src_dir.glob("**/*.*")
        items: list[Item] = []

        for path in file_paths:
            if _should_skip_source_file(path):
                continue

            # Get the bounds of the features in the vector file, which we assume fiona can
            # read.
            # For shapefile, to open it we need to copy all the aux files.
            aux_files: list[UPath] = []
            if path.name.endswith(".shp"):
                prefix = ".".join(path.name.split(".")[:-1])
                for ext in SHAPEFILE_AUX_EXTENSIONS:
                    aux_files.append(path.parent / (prefix + ext))

            with get_upath_local(path, extra_paths=aux_files) as local_fname:
                with fiona.open(local_fname) as src:
                    crs = CRS.from_wkt(src.crs.to_wkt())
                    bounds = None
                    for feat in src:
                        shp = shapely.geometry.shape(feat.geometry)
                        cur_bounds = shp.bounds
                        if bounds is None:
                            bounds = list(cur_bounds)
                        else:
                            bounds[0] = min(bounds[0], cur_bounds[0])
                            bounds[1] = min(bounds[1], cur_bounds[1])
                            bounds[2] = max(bounds[2], cur_bounds[2])
                            bounds[3] = max(bounds[3], cur_bounds[3])

                    # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
                    # should be 1 projection unit/pixel.
                    projection = Projection(crs, 1, 1)
                    geometry = STGeometry(
                        projection,
                        shapely.box(*bounds).buffer(self.item_buffer_epsilon),
                        None,
                    )

            logger.debug(
                "VectorImporter.list_items: got bounds of %s: %s", path, geometry
            )
            items.append(
                VectorItem(path.name.split(".")[0], geometry, path.absolute().as_uri())
            )

        logger.debug("VectorImporter.list_items: discovered %d items", len(items))
        return items

    def ingest_item(
        self,
        tile_store: TileStoreWithLayer,
        item: Item,
        cur_geometries: list[STGeometry],
    ) -> None:
        """Ingest the specified local file item.

        Args:
            tile_store: the TileStore to ingest the data into.
            item: the Item to ingest
            cur_geometries: the geometries where the item is needed.
        """
        if tile_store.is_vector_ready(item):
            return

        assert isinstance(item, VectorItem)
        path = UPath(item.path_uri)

        aux_files: list[UPath] = []
        if path.name.split(".")[-1] == "shp":
            prefix = ".".join(path.name.split(".")[:-1])
            for ext in SHAPEFILE_AUX_EXTENSIONS:
                aux_files.append(path.parent / (prefix + ext))

        # TODO: move converting fiona file to list[Feature] to utility function.
        with get_upath_local(path, extra_paths=aux_files) as local_fname:
            with fiona.open(local_fname) as src:
                crs = CRS.from_wkt(src.crs.to_wkt())
                # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
                # should be 1 projection unit/pixel.
                projection = Projection(crs, 1, 1)

                features = []
                for feat in src:
                    features.append(
                        Feature.from_geojson(
                            projection,
                            {
                                "type": "Feature",
                                "geometry": dict(feat.geometry),
                                "properties": dict(feat.properties),
                            },
                        )
                    )

                tile_store.write_vector(item, features)

list_items 

list_items(src_dir: UPath) -> list[Item]

Extract a list of Items from the source directory.

Parameters:

Name Type Description Default
src_dir UPath

the source directory.

 required
Source code in rslearn/data_sources/local_files.py 
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
def list_items(self, src_dir: UPath) -> list[Item]:
    """Extract a list of Items from the source directory.

    Args:
        src_dir: the source directory.
    """
    file_paths = src_dir.glob("**/*.*")
    items: list[Item] = []

    for path in file_paths:
        if _should_skip_source_file(path):
            continue

        # Get the bounds of the features in the vector file, which we assume fiona can
        # read.
        # For shapefile, to open it we need to copy all the aux files.
        aux_files: list[UPath] = []
        if path.name.endswith(".shp"):
            prefix = ".".join(path.name.split(".")[:-1])
            for ext in SHAPEFILE_AUX_EXTENSIONS:
                aux_files.append(path.parent / (prefix + ext))

        with get_upath_local(path, extra_paths=aux_files) as local_fname:
            with fiona.open(local_fname) as src:
                crs = CRS.from_wkt(src.crs.to_wkt())
                bounds = None
                for feat in src:
                    shp = shapely.geometry.shape(feat.geometry)
                    cur_bounds = shp.bounds
                    if bounds is None:
                        bounds = list(cur_bounds)
                    else:
                        bounds[0] = min(bounds[0], cur_bounds[0])
                        bounds[1] = min(bounds[1], cur_bounds[1])
                        bounds[2] = max(bounds[2], cur_bounds[2])
                        bounds[3] = max(bounds[3], cur_bounds[3])

                # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
                # should be 1 projection unit/pixel.
                projection = Projection(crs, 1, 1)
                geometry = STGeometry(
                    projection,
                    shapely.box(*bounds).buffer(self.item_buffer_epsilon),
                    None,
                )

        logger.debug(
            "VectorImporter.list_items: got bounds of %s: %s", path, geometry
        )
        items.append(
            VectorItem(path.name.split(".")[0], geometry, path.absolute().as_uri())
        )

    logger.debug("VectorImporter.list_items: discovered %d items", len(items))
    return items

ingest_item 

ingest_item(tile_store: TileStoreWithLayer, item: Item, cur_geometries: list[STGeometry]) -> None

Ingest the specified local file item.

Parameters:

Name Type Description Default
tile_store TileStoreWithLayer

the TileStore to ingest the data into.

 required
item Item

the Item to ingest

 required
cur_geometries list[STGeometry]

the geometries where the item is needed.

 required
Source code in rslearn/data_sources/local_files.py 
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
def ingest_item(
    self,
    tile_store: TileStoreWithLayer,
    item: Item,
    cur_geometries: list[STGeometry],
) -> None:
    """Ingest the specified local file item.

    Args:
        tile_store: the TileStore to ingest the data into.
        item: the Item to ingest
        cur_geometries: the geometries where the item is needed.
    """
    if tile_store.is_vector_ready(item):
        return

    assert isinstance(item, VectorItem)
    path = UPath(item.path_uri)

    aux_files: list[UPath] = []
    if path.name.split(".")[-1] == "shp":
        prefix = ".".join(path.name.split(".")[:-1])
        for ext in SHAPEFILE_AUX_EXTENSIONS:
            aux_files.append(path.parent / (prefix + ext))

    # TODO: move converting fiona file to list[Feature] to utility function.
    with get_upath_local(path, extra_paths=aux_files) as local_fname:
        with fiona.open(local_fname) as src:
            crs = CRS.from_wkt(src.crs.to_wkt())
            # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
            # should be 1 projection unit/pixel.
            projection = Projection(crs, 1, 1)

            features = []
            for feat in src:
                features.append(
                    Feature.from_geojson(
                        projection,
                        {
                            "type": "Feature",
                            "geometry": dict(feat.geometry),
                            "properties": dict(feat.properties),
                        },
                    )
                )

            tile_store.write_vector(item, features)

LocalFiles

Bases: DataSource

A data source for ingesting data from local files.

Source code in rslearn/data_sources/local_files.py 
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
class LocalFiles(DataSource):
    """A data source for ingesting data from local files."""

    def __init__(
        self,
        src_dir: str,
        raster_item_specs: list[RasterItemSpec] | None = None,
        layer_type: LayerType | None = None,
        context: DataSourceContext = DataSourceContext(),
    ) -> None:
        """Initialize a new LocalFiles instance.

        Args:
            src_dir: source directory to ingest
            raster_item_specs: the specs to specify the raster items directly. If None,
                the raster items are automatically detected from the files in the
                source directory.
            layer_type: the layer type. It only needs to be set if the layer_config is
                missing from the context.
            context: the data source context. The layer config must be in the context.
        """
        if context.ds_path is not None:
            self.src_dir = join_upath(context.ds_path, src_dir)
        else:
            self.src_dir = UPath(src_dir)

        # Determine layer type.
        if context.layer_config is not None:
            self.layer_type = context.layer_config.type
        elif layer_type is not None:
            self.layer_type = layer_type
        else:
            raise ValueError(
                "layer type must be specified if the layer config is not in the context"
            )

        self.importer: Importer
        if self.layer_type == LayerType.RASTER:
            self.importer = RasterImporter(item_specs=raster_item_specs)
        elif self.layer_type == LayerType.VECTOR:
            self.importer = VectorImporter()
        else:
            raise ValueError(f"unknown layer type {self.layer_type}")

    def get_cache_fname(self) -> UPath:
        """Returns the file where the item list is cached."""
        return self.src_dir / "summary.json"

    def _read_cached_items(self, cache_fname: UPath) -> list[dict] | None:
        """Read the serialized item list from the cache file.

        Args:
            cache_fname: the file where the item list is cached.

        Returns:
            the serialized items, or None if the cache is missing or unusable and the
            items must be listed again. The cache can be unusable if it was truncated
            by an interrupted write, or if a concurrent writer is filling it in on a
            filesystem where the write is not atomic.
        """
        try:
            with cache_fname.open() as f:
                return json.load(f)
        except FileNotFoundError:
            logger.debug("cache at %s does not exist, listing items", cache_fname)
            return None
        except json.JSONDecodeError:
            logger.warning(
                "cache at %s is incomplete or corrupt, listing items instead",
                cache_fname,
            )
            return None

    @functools.cache
    def list_items(self) -> list[Item]:
        """Lists items from the source directory while maintaining a cache file."""
        cache_fname = self.get_cache_fname()

        logger.debug("loading item list from cache at %s", cache_fname)
        serialized_items = self._read_cached_items(cache_fname)
        if serialized_items is not None:
            return [
                self.deserialize_item(serialized_item)
                for serialized_item in serialized_items
            ]

        items = self.importer.list_items(self.src_dir)
        serialized_items = [item.serialize() for item in items]
        # Write the cache atomically. Several workers may be preparing windows against
        # the same source directory at once, and a reader must never see a cache file
        # that another worker has created but not finished writing.
        with open_atomic(cache_fname, "w") as f:
            json.dump(serialized_items, f)
        return items

    @functools.cache
    def _get_spatial_index(self) -> GridIndex:
        """Build an in-memory GridIndex over items in WGS84 for fast spatial lookup."""
        grid_index = GridIndex(1.0)
        for item in self.list_items():
            wgs84_geom = item.geometry.to_wgs84()
            grid_index.insert(wgs84_geom.shp.bounds, item)
        return grid_index

    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[MatchedItemGroup[Item]]]:
        """Get a list of items in the data source intersecting the given geometries.

        Args:
            geometries: the spatiotemporal geometries
            query_config: the query configuration

        Returns:
            List of groups of items that should be retrieved for each geometry.
        """
        grid_index = self._get_spatial_index()

        groups = []
        for geometry in geometries:
            wgs84_geometry = geometry.to_wgs84()
            cur_items = []
            for item in grid_index.query(wgs84_geometry.shp.bounds):
                if not item.geometry.intersects(geometry):
                    continue
                cur_items.append(item)

            cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
                geometry, cur_items, query_config
            )
            groups.append(cur_groups)
        return groups

    def deserialize_item(self, serialized_item: dict) -> RasterItem | VectorItem:
        """Deserializes an item from JSON-decoded data."""
        if self.layer_type == LayerType.RASTER:
            return RasterItem.deserialize(serialized_item)
        elif self.layer_type == LayerType.VECTOR:
            return VectorItem.deserialize(serialized_item)
        else:
            raise ValueError(f"Unknown layer type: {self.layer_type}")

    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[Item],
        geometries: list[list[STGeometry]],
    ) -> None:
        """Ingest items into the given tile store.

        Args:
            tile_store: the tile store to ingest into
            items: the items to ingest
            geometries: a list of geometries needed for each item
        """
        for item, cur_geometries in zip(items, geometries):
            self.importer.ingest_item(tile_store, item, cur_geometries)

get_cache_fname 

get_cache_fname() -> UPath

Returns the file where the item list is cached.

Source code in rslearn/data_sources/local_files.py 
467
468
469
def get_cache_fname(self) -> UPath:
    """Returns the file where the item list is cached."""
    return self.src_dir / "summary.json"

list_items cached 

list_items() -> list[Item]

Lists items from the source directory while maintaining a cache file.

Source code in rslearn/data_sources/local_files.py 
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
@functools.cache
def list_items(self) -> list[Item]:
    """Lists items from the source directory while maintaining a cache file."""
    cache_fname = self.get_cache_fname()

    logger.debug("loading item list from cache at %s", cache_fname)
    serialized_items = self._read_cached_items(cache_fname)
    if serialized_items is not None:
        return [
            self.deserialize_item(serialized_item)
            for serialized_item in serialized_items
        ]

    items = self.importer.list_items(self.src_dir)
    serialized_items = [item.serialize() for item in items]
    # Write the cache atomically. Several workers may be preparing windows against
    # the same source directory at once, and a reader must never see a cache file
    # that another worker has created but not finished writing.
    with open_atomic(cache_fname, "w") as f:
        json.dump(serialized_items, f)
    return items

get_items 

get_items(geometries: list[STGeometry], query_config: QueryConfig) -> list[list[MatchedItemGroup[Item]]]

Get a list of items in the data source intersecting the given geometries.

Parameters:

Name Type Description Default
geometries list[STGeometry]

the spatiotemporal geometries

 required
query_config QueryConfig

the query configuration

 required

Returns:

Type Description
list[list[MatchedItemGroup[Item]]]

List of groups of items that should be retrieved for each geometry.
Source code in rslearn/data_sources/local_files.py 
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
def get_items(
    self, geometries: list[STGeometry], query_config: QueryConfig
) -> list[list[MatchedItemGroup[Item]]]:
    """Get a list of items in the data source intersecting the given geometries.

    Args:
        geometries: the spatiotemporal geometries
        query_config: the query configuration

    Returns:
        List of groups of items that should be retrieved for each geometry.
    """
    grid_index = self._get_spatial_index()

    groups = []
    for geometry in geometries:
        wgs84_geometry = geometry.to_wgs84()
        cur_items = []
        for item in grid_index.query(wgs84_geometry.shp.bounds):
            if not item.geometry.intersects(geometry):
                continue
            cur_items.append(item)

        cur_groups = rslearn.data_sources.utils.match_candidate_items_to_window(
            geometry, cur_items, query_config
        )
        groups.append(cur_groups)
    return groups

deserialize_item 

deserialize_item(serialized_item: dict) -> RasterItem | VectorItem

Deserializes an item from JSON-decoded data.

Source code in rslearn/data_sources/local_files.py 
556
557
558
559
560
561
562
563
def deserialize_item(self, serialized_item: dict) -> RasterItem | VectorItem:
    """Deserializes an item from JSON-decoded data."""
    if self.layer_type == LayerType.RASTER:
        return RasterItem.deserialize(serialized_item)
    elif self.layer_type == LayerType.VECTOR:
        return VectorItem.deserialize(serialized_item)
    else:
        raise ValueError(f"Unknown layer type: {self.layer_type}")

ingest 

ingest(tile_store: TileStoreWithLayer, items: list[Item], geometries: list[list[STGeometry]]) -> None

Ingest items into the given tile store.

Parameters:

Name Type Description Default
tile_store TileStoreWithLayer

the tile store to ingest into

 required
items list[Item]

the items to ingest

 required
geometries list[list[STGeometry]]

a list of geometries needed for each item

 required
Source code in rslearn/data_sources/local_files.py 
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
def ingest(
    self,
    tile_store: TileStoreWithLayer,
    items: list[Item],
    geometries: list[list[STGeometry]],
) -> None:
    """Ingest items into the given tile store.

    Args:
        tile_store: the tile store to ingest into
        items: the items to ingest
        geometries: a list of geometries needed for each item
    """
    for item, cur_geometries in zip(items, geometries):
        self.importer.ingest_item(tile_store, item, cur_geometries)