Download from tileserver concurrently

georeader/readers/tileserver.py

@@ -13,6 +13,7 @@
 import rasterio.windows
 import rasterio.transform
 from shapely.geometry import box
+from concurrent.futures import ThreadPoolExecutor
 
 def read_from_tileserver(tile_server:str, geometry:Union[Polygon, MultiPolygon],
                          zoom:int=16, crs_geometry:Any="EPSG:4326") -> GeoTensor:
@@ -33,11 +34,21 @@ def read_from_tileserver(tile_server:str, geometry:Union[Polygon, MultiPolygon],
 
     min_lon, min_lat, max_lon, max_lat = geometry.bounds
     tiles = mercantile.tiles(min_lon, min_lat, max_lon, max_lat, zooms=zoom)
-    geotensors = []
-    for tile in tiles:
-        if not box(*mercantile.bounds(tile)).intersects(geometry):
-            continue
+    tiles = [tile for tile in tiles if box(*mercantile.bounds(tile)).intersects(geometry)]
+    # geotensors = []
+    # for tile in tiles:
 
+    #     rsp = requests.get(tile_server.format(x=tile.x, y=tile.y, z=tile.z))
+    #     img = Image.open(BytesIO(rsp.content))
+    #     xmin, ymin, xmax, ymax = window_utils.normalize_bounds(mercantile.xy_bounds(tile))
+    #     img_np = np.array(img).transpose(2,0,1)
+
+    #     transform = rasterio.transform.from_bounds(west=xmin, south=ymin, east=xmax, north=ymax,
+    #                                                width=img_np.shape[2], height=img_np.shape[1])
+
+    #     geotensors.append(GeoTensor(img_np, transform=transform, crs="EPSG:3857"))
+
+    def read_tile(tile):
         rsp = requests.get(tile_server.format(x=tile.x, y=tile.y, z=tile.z))
         img = Image.open(BytesIO(rsp.content))
         xmin, ymin, xmax, ymax = window_utils.normalize_bounds(mercantile.xy_bounds(tile))
@@ -46,7 +57,10 @@ def read_from_tileserver(tile_server:str, geometry:Union[Polygon, MultiPolygon],
         transform = rasterio.transform.from_bounds(west=xmin, south=ymin, east=xmax, north=ymax,
                                                    width=img_np.shape[2], height=img_np.shape[1])
 
-        geotensors.append(GeoTensor(img_np, transform=transform, crs="EPSG:3857"))
+        return GeoTensor(img_np, transform=transform, crs="EPSG:3857")
+
+    with ThreadPoolExecutor() as executor:
+        geotensors = list(executor.map(lambda tile: read_tile(tile), tiles))
 
     if len(geotensors) == 1:
         return read.read_from_polygon(geotensors[0], polygon=geometry, crs_polygon="EPSG:4326")