feat: use ProcessPoolExecutor for download_and_process_tile

Author: Gigaszi

mapswipe_workers/mapswipe_workers/utils/process_mapillary.py

@@ -1,18 +1,11 @@
 import os
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
 
 import mercantile
 import pandas as pd
 import requests
-from shapely import (
-    LineString,
-    MultiLineString,
-    MultiPolygon,
-    Point,
-    Polygon,
-    box,
-    unary_union,
-)
+from shapely import MultiPolygon, Point, Polygon, box, unary_union
 from shapely.geometry import shape
 from vt2geojson import tools as vt2geojson_tools
 
@@ -44,7 +37,7 @@ def create_tiles(polygon, level):
     return tiles
 
 
-def download_and_process_tile(row, attempt_limit=3):
+def download_and_process_tile(row, polygon, kwargs, attempt_limit=3):
     z = row["z"]
     x = row["x"]
     y = row["y"]
@@ -59,31 +52,37 @@ def download_and_process_tile(row, attempt_limit=3):
                 "features", []
             )
             data = []
-            for feature in features:
-                geometry = feature.get("geometry", {})
-                properties = feature.get("properties", {})
-                geometry_type = geometry.get("type", None)
-                coordinates = geometry.get("coordinates", [])
-
-                element_geometry = None
-                if geometry_type == "Point":
-                    element_geometry = Point(coordinates)
-                elif geometry_type == "LineString":
-                    element_geometry = LineString(coordinates)
-                elif geometry_type == "MultiLineString":
-                    element_geometry = MultiLineString(coordinates)
-                elif geometry_type == "Polygon":
-                    element_geometry = Polygon(coordinates[0])
-                elif geometry_type == "MultiPolygon":
-                    element_geometry = MultiPolygon(coordinates)
-
-                # Append the dictionary with geometry and properties
-                row = {"geometry": element_geometry, **properties}
-                data.append(row)
+            data.extend(
+                [
+                    {
+                        "geometry": Point(feature["geometry"]["coordinates"]),
+                        **feature.get("properties", {}),
+                    }
+                    for feature in features
+                    if feature.get("geometry", {}).get("type") == "Point"
+                ]
+            )
 
             data = pd.DataFrame(data)
 
-            if not data.empty:
+            if data.isna().all().all() is False or data.empty is False:
+                data = data[data["geometry"].apply(lambda point: point.within(polygon))]
+                target_columns = [
+                    "id",
+                    "geometry",
+                    "captured_at",
+                    "is_pano",
+                    "compass_angle",
+                    "sequence",
+                    "organization_id",
+                ]
+                for col in target_columns:
+                    if col not in data.columns:
+                        data[col] = None
+
+                if data.isna().all().all() is False or data.empty is False:
+                    data = filter_results(data, **kwargs)
+
                 return data
         except Exception as e:
             print(f"An exception occurred while requesting a tile: {e}")
@@ -94,7 +93,7 @@ def download_and_process_tile(row, attempt_limit=3):
 
 
 def coordinate_download(
-    polygon, level, use_concurrency=True, attempt_limit=3, workers=os.cpu_count() * 4
+    polygon, level, kwargs: dict, use_concurrency=True, workers=os.cpu_count() * 4
 ):
     tiles = create_tiles(polygon, level)
 
@@ -104,45 +103,22 @@ def coordinate_download(
         if not use_concurrency:
             workers = 1
 
-        futures = []
-        with ThreadPoolExecutor(max_workers=workers) as executor:
-            for index, row in tiles.iterrows():
-                futures.append(
-                    executor.submit(download_and_process_tile, row, attempt_limit)
-                )
-
-            for future in as_completed(futures):
-                if future is not None:
-                    df = future.result()
+        process_tile_with_args = partial(
+            download_and_process_tile, polygon=polygon, kwargs=kwargs
+        )
+        with ProcessPoolExecutor(max_workers=workers) as executor:
+            futures = list(
+                executor.map(process_tile_with_args, tiles.to_dict(orient="records"))
+            )
 
-                    if df is not None and not df.empty:
-                        downloaded_metadata.append(df)
+            for df in futures:
+                if df is not None and not df.empty:
+                    downloaded_metadata.append(df)
         if len(downloaded_metadata):
             downloaded_metadata = pd.concat(downloaded_metadata, ignore_index=True)
         else:
             return pd.DataFrame(downloaded_metadata)
 
-        target_columns = [
-            "id",
-            "geometry",
-            "captured_at",
-            "is_pano",
-            "compass_angle",
-            "sequence",
-            "organization_id",
-        ]
-        for col in target_columns:
-            if col not in downloaded_metadata.columns:
-                downloaded_metadata[col] = None
-        if (
-            downloaded_metadata.isna().all().all() is False
-            or downloaded_metadata.empty is False
-        ):
-            downloaded_metadata = downloaded_metadata[
-                downloaded_metadata["geometry"].apply(
-                    lambda point: point.within(polygon)
-                )
-            ]
         return downloaded_metadata
 
 
@@ -198,74 +174,72 @@ def filter_results(
             )
             return None
         df = df[df["creator_id"] == creator_id]
-
     if is_pano is not None:
         if df["is_pano"].isna().all():
+            print(df)
             logger.exception("No Mapillary Feature in the AoI has a 'is_pano' value.")
             return None
         df = df[df["is_pano"] == is_pano]
-
     if organization_id is not None:
         if df["organization_id"].isna().all():
             logger.exception(
                 "No Mapillary Feature in the AoI has an 'organization_id' value."
             )
             return None
         df = df[df["organization_id"] == organization_id]
-
     if start_time is not None:
         if df["captured_at"].isna().all():
             logger.exception(
                 "No Mapillary Feature in the AoI has a 'captured_at' value."
             )
             return None
         df = filter_by_timerange(df, start_time, end_time)
-
     return df
 
 
 def get_image_metadata(
     aoi_geojson,
     level=14,
-    attempt_limit=3,
     is_pano: bool = None,
     creator_id: int = None,
     organization_id: str = None,
     start_time: str = None,
     end_time: str = None,
     sampling_threshold=None,
 ):
+    kwargs = {
+        "is_pano": is_pano,
+        "creator_id": creator_id,
+        "organization_id": organization_id,
+        "start_time": start_time,
+        "end_time": end_time,
+    }
     aoi_polygon = geojson_to_polygon(aoi_geojson)
-    downloaded_metadata = coordinate_download(aoi_polygon, level, attempt_limit)
-
+    downloaded_metadata = coordinate_download(aoi_polygon, level, kwargs)
     if downloaded_metadata.empty or downloaded_metadata.isna().all().all():
         raise ValueError("No Mapillary Features in the AoI.")
 
     downloaded_metadata = downloaded_metadata[
         downloaded_metadata["geometry"].apply(lambda geom: isinstance(geom, Point))
     ]
 
-    filtered_metadata = filter_results(
-        downloaded_metadata, creator_id, is_pano, organization_id, start_time, end_time
-    )
-
     if (
-        filtered_metadata is None
-        or filtered_metadata.empty
-        or filtered_metadata.isna().all().all()
+        downloaded_metadata is None
+        or downloaded_metadata.empty
+        or downloaded_metadata.isna().all().all()
     ):
         raise ValueError("No Mapillary Features in the AoI match the filter criteria.")
 
     if sampling_threshold is not None:
-        filtered_metadata = spatial_sampling(filtered_metadata, sampling_threshold)
+        downloaded_metadata = spatial_sampling(downloaded_metadata, sampling_threshold)
 
-    total_images = len(filtered_metadata)
+    total_images = len(downloaded_metadata)
     if total_images > 100000:
         raise ValueError(
             f"Too many Images with selected filter options for the AoI: {total_images}"
         )
 
     return {
-        "ids": filtered_metadata["id"].tolist(),
-        "geometries": filtered_metadata["geometry"].tolist(),
+        "ids": downloaded_metadata["id"].tolist(),
+        "geometries": downloaded_metadata["geometry"].tolist(),
     }