Merge pull request #250 from 3DBAG/249-fix-ahn5-laz-tiles-input-for-roofer

249 fix ahn5 laz tiles input for roofer

Author: balazsdukai

packages/core/src/bag3d/core/assets/ahn/tile.py

@@ -1,11 +1,13 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
 
 from dagster import asset, AssetKey, Output, Field, get_dagster_logger
 
 from pgutils import PostgresTableIdentifier, inject_parameters
 from psycopg.sql import SQL
 from psycopg import Connection
 
+from bag3d.common.types import LocalPath
 from bag3d.common.utils.database import create_schema, load_sql
 from bag3d.common.utils.geodata import wkt_from_bbox
 from bag3d.core.assets.ahn.core import (
@@ -92,9 +94,9 @@ def regular_grid_200m(context):
     deps={AssetKey(["ahn", "metadata_ahn3"])},
     required_resource_keys={"file_store", "lastools", "db_connection"},
 )
-def laz_tiles_ahn3_200m(context, regular_grid_200m, metadata_table_ahn3):
+def laz_tiles_ahn3_200m(context, regular_grid_200m, metadata_table_ahn3) -> LocalPath:
     """AHN3 partitioned by a grid of 200m cells on the extent of the AHN PDOK tiles."""
-    partition_laz_with_grid(
+    return partition_laz_with_grid(
         context,
         metadata_table_ahn3,
         regular_grid_200m,
@@ -123,9 +125,9 @@ def laz_tiles_ahn3_200m(context, regular_grid_200m, metadata_table_ahn3):
     deps={AssetKey(["ahn", "metadata_ahn4"])},
     required_resource_keys={"file_store", "lastools", "db_connection"},
 )
-def laz_tiles_ahn4_200m(context, regular_grid_200m, metadata_table_ahn4):
+def laz_tiles_ahn4_200m(context, regular_grid_200m, metadata_table_ahn4) -> LocalPath:
     """AHN4 partitioned by a grid of 200m cells on the extent of the AHN PDOK tiles."""
-    partition_laz_with_grid(
+    return partition_laz_with_grid(
         context,
         metadata_table_ahn4,
         regular_grid_200m,
@@ -154,9 +156,9 @@ def laz_tiles_ahn4_200m(context, regular_grid_200m, metadata_table_ahn4):
     deps={AssetKey(["ahn", "metadata_ahn5"])},
     required_resource_keys={"file_store", "lastools", "db_connection"},
 )
-def laz_tiles_ahn5_200m(context, regular_grid_200m, metadata_table_ahn5):
+def laz_tiles_ahn5_200m(context, regular_grid_200m, metadata_table_ahn5) -> LocalPath:
     """AHN5 partitioned by a grid of 200m cells on the extent of the AHN PDOK tiles."""
-    partition_laz_with_grid(
+    return partition_laz_with_grid(
         context,
         metadata_table_ahn5,
         regular_grid_200m,
@@ -175,7 +177,7 @@ def partition_laz_with_grid(
     cellsize,
     max_workers,
     verbose: bool = False,
-):
+) -> Path:
     """
     Partitions LAZ files into tiles based on a regular grid.
     Queries the PostgreSQL database to find the grid cells that overlap the AHN tile.
@@ -192,7 +194,7 @@ def partition_laz_with_grid(
         verbose: Whether to suppress output from the subprocesses.
 
     Returns:
-        None
+        The Path to the directory with the output 200m tiles
     """
     conn = context.resources.db_connection.connect
     query_params = {
@@ -262,3 +264,4 @@ def partition_laz_with_grid(
                 logger.warning(f"Tile {tile} raised an exception: {e}")
     if len(failed) > 0:
         logger.error(f"Failed {len(failed)} tiles. Failed tiles: {failed}")
+    return out_dir

packages/core/src/bag3d/core/assets/reconstruction/reconstruction.py

@@ -56,6 +56,9 @@ def __init__(self, schema: str, table_tiles: str, wkt: str = None):
     ),
     ins={
         "regular_grid_200m": AssetIn(key_prefix="ahn"),
+        "laz_tiles_ahn3_200m": AssetIn(key_prefix="ahn"),
+        "laz_tiles_ahn4_200m": AssetIn(key_prefix="ahn"),
+        "laz_tiles_ahn5_200m": AssetIn(key_prefix="ahn"),
         "tiles": AssetIn(key_prefix="input"),
         "index": AssetIn(key_prefix="input"),
         "reconstruction_input": AssetIn(key_prefix="input"),
@@ -80,25 +83,17 @@ def __init__(self, schema: str, table_tiles: str, wkt: str = None):
             is_required=False,
             default_value="info",
         ),
-        "dir_tiles_200m_ahn3": Field(
-            str,
-            description="Directory of the 200m tiles of AHN3. Used if the tiles are stored in a non-standard location.",
-            is_required=False,
-        ),
-        "dir_tiles_200m_ahn4": Field(
-            str,
-            description="Directory of the 200m tiles of AHN4. Used if the tiles are stored in a non-standard location.",
-            is_required=False,
-        ),
-        "dir_tiles_200m_ahn5": Field(
-            str,
-            description="Directory of the 200m tiles of AHN5. Used if the tiles are stored in a non-standard location.",
-            is_required=False,
-        ),
     },
 )
 def reconstructed_building_models_nl(
-    context, regular_grid_200m, tiles, index, reconstruction_input
+    context,
+    regular_grid_200m,
+    tiles,
+    index,
+    reconstruction_input,
+    laz_tiles_ahn3_200m,
+    laz_tiles_ahn4_200m,
+    laz_tiles_ahn5_200m,
 ):
     """Generate the 3D building models by running the reconstruction sequentially
     within one partition.
@@ -110,15 +105,9 @@ def reconstructed_building_models_nl(
         reconstruction_input,
         regular_grid_200m,
         tiles,
-        dir_tiles_200m_ahn3=context.op_execution_context.op_config.get(
-            "dir_tiles_200m_ahn3"
-        ),
-        dir_tiles_200m_ahn4=context.op_execution_context.op_config.get(
-            "dir_tiles_200m_ahn4"
-        ),
-        dir_tiles_200m_ahn5=context.op_execution_context.op_config.get(
-            "dir_tiles_200m_ahn5"
-        ),
+        dir_tiles_200m_ahn3=laz_tiles_ahn3_200m,
+        dir_tiles_200m_ahn4=laz_tiles_ahn4_200m,
+        dir_tiles_200m_ahn5=laz_tiles_ahn5_200m,
     )
 
     context.log.info(f"{roofer_toml=}")
@@ -152,6 +141,19 @@ def create_roofer_config(
     dir_tiles_200m_ahn4=None,
     dir_tiles_200m_ahn5=None,
 ):
+    def laz_filepaths_generator(
+        file_store_dir: Path, resultset: list[tuple], dir_200m_laz=None
+    ):
+        """Generator for the full 200m laz file paths that only emits existing paths."""
+        if dir_200m_laz is not None:
+            lazdir = Path(dir_200m_laz)
+        else:
+            lazdir = ahn_dir(file_store_dir, ahn_version=5).joinpath("tiles_200m")
+        for tile_id_ahn in resultset:
+            p = lazdir / f"t_{tile_id_ahn[0]}.laz"
+            if p.is_file():
+                yield str(p)
+
     toml_template = """
     polygon-source = "{footprint_file}"
     id-attribute = "identificatie"
@@ -237,26 +239,30 @@ def create_roofer_config(
             "tile_id": tile_id,
         },
     )
-    if dir_tiles_200m_ahn3 is not None:
-        out_dir_ahn3 = Path(dir_tiles_200m_ahn3)
-    else:
-        out_dir_ahn3 = ahn_dir(
-            context.resources.file_store.file_store.data_dir, ahn_version=3
-        ).joinpath("tiles_200m")
-    laz_files_ahn3 = [
-        str(out_dir_ahn3 / f"t_{tile_id_ahn[0]}.laz") for tile_id_ahn in res
-    ]
-    # TODO: probably should take the tiles_200m directory from the asset output
-    if dir_tiles_200m_ahn4 is not None:
-        out_dir_ahn4 = Path(dir_tiles_200m_ahn4)
-    else:
-        out_dir_ahn4 = ahn_dir(
-            context.resources.file_store.file_store.data_dir, ahn_version=4
-        ).joinpath("tiles_200m")
-    # TODO: same with the laz filename pattern
-    laz_files_ahn4 = [
-        str(out_dir_ahn4 / f"t_{tile_id_ahn[0]}.laz") for tile_id_ahn in res
-    ]
+    laz_files_ahn3 = list(
+        laz_filepaths_generator(
+            file_store_dir=context.resources.file_store.file_store.data_dir,
+            resultset=res,
+            dir_200m_laz=dir_tiles_200m_ahn3,
+        )
+    )
+
+    laz_files_ahn4 = list(
+        laz_filepaths_generator(
+            file_store_dir=context.resources.file_store.file_store.data_dir,
+            resultset=res,
+            dir_200m_laz=dir_tiles_200m_ahn4,
+        )
+    )
+
+    laz_files_ahn5 = list(
+        laz_filepaths_generator(
+            file_store_dir=context.resources.file_store.file_store.data_dir,
+            resultset=res,
+            dir_200m_laz=dir_tiles_200m_ahn5,
+        )
+    )
+
     # Would be neater if we could use -sql in the OGR connection to do this query,
     # instead of creating a view.
     tile_view = PostgresTableIdentifier(tiles.schema, f"t_{tile_id}")
@@ -284,13 +290,7 @@ def create_roofer_config(
         footprint_file=f"PG:{context.resources.db_connection.connect.dsn} tables={tile_view}",
         ahn3_files=laz_files_ahn3,
         ahn4_files=laz_files_ahn4,
-        ahn5_files=[
-            str(
-                ahn_dir(
-                    context.resources.file_store.file_store.data_dir, ahn_version=5
-                ).joinpath("as_downloaded/LAZ")
-            )
-        ],
+        ahn5_files=laz_files_ahn5,
         output_path=output_dir,
     )
     path_toml = output_dir / "roofer.toml"

packages/core/tests/conftest.py

@@ -329,3 +329,66 @@ def handle_output(self, context, obj):  # pragma: no cover
         key=AssetKey(["ahn", "regular_grid_200m"]),
         io_manager_def=MockIOManager(),
     )
+
+
+@pytest.fixture(scope="session")
+def mock_asset_laz_tiles_ahn3_200m(test_data_dir):
+    class MockIOManager(IOManager):
+        def load_input(self, context):
+            return (
+                test_data_dir
+                / "reconstruction_input"
+                / "pointcloud"
+                / "AHN3"
+                / "tiles_200m"
+            )
+
+        def handle_output(self, context, obj):  # pragma: no cover
+            raise NotImplementedError()
+
+    return SourceAsset(
+        key=AssetKey(["ahn", "laz_tiles_ahn3_200m"]),
+        io_manager_def=MockIOManager(),
+    )
+
+
+@pytest.fixture(scope="session")
+def mock_asset_laz_tiles_ahn4_200m(test_data_dir):
+    class MockIOManager(IOManager):
+        def load_input(self, context):
+            return (
+                test_data_dir
+                / "reconstruction_input"
+                / "pointcloud"
+                / "AHN4"
+                / "tiles_200m"
+            )
+
+        def handle_output(self, context, obj):  # pragma: no cover
+            raise NotImplementedError()
+
+    return SourceAsset(
+        key=AssetKey(["ahn", "laz_tiles_ahn4_200m"]),
+        io_manager_def=MockIOManager(),
+    )
+
+
+@pytest.fixture(scope="session")
+def mock_asset_laz_tiles_ahn5_200m(test_data_dir):
+    class MockIOManager(IOManager):
+        def load_input(self, context):
+            return (
+                test_data_dir
+                / "reconstruction_input"
+                / "pointcloud"
+                / "AHN5"
+                / "tiles_200m"
+            )
+
+        def handle_output(self, context, obj):  # pragma: no cover
+            raise NotImplementedError()
+
+    return SourceAsset(
+        key=AssetKey(["ahn", "laz_tiles_ahn5_200m"]),
+        io_manager_def=MockIOManager(),
+    )

packages/core/tests/test_integration.py

@@ -35,6 +35,9 @@ def test_integration_reconstruction_and_export(
     mock_asset_reconstruction_input,
     mock_asset_tiles,
     mock_asset_index,
+    mock_asset_laz_tiles_ahn3_200m,
+    mock_asset_laz_tiles_ahn4_200m,
+    mock_asset_laz_tiles_ahn5_200m,
 ):
     # update quadtree
     og_quadtree = test_data_dir / "quadtree.tsv"
@@ -104,6 +107,9 @@ def test_integration_reconstruction_and_export(
             mock_asset_reconstruction_input,
             mock_asset_tiles,
             mock_asset_index,
+            mock_asset_laz_tiles_ahn3_200m,
+            mock_asset_laz_tiles_ahn4_200m,
+            mock_asset_laz_tiles_ahn5_200m,
             *reconstruction_assets,
             *all_export_assets,
         ],