Use downloaded folders

download_files_from_LGL_BW.py

@@ -63,7 +63,7 @@ UTM = 32
 
 SCRIPT_DIR = Path(__file__).parent
 WAIT_BETWEEN_DOWNLOADS = 5  # [s] Be nice to LGL Server.
-GML_GLOB = "LoD2_*/LoD2_*.gml"
+GML_GLOB = "LoD2_*.gml"
 
 if EXTRACT_REGIONS:
     try:
@@ -109,11 +109,12 @@ def wkt_polygon_to_grid_coords(location_name: str, wkt: str) -> tuple[int, int,
     return (x1, x2, y1, y2)
 
 
-def download_whole_region(output_dir: Path, wkt_region: str, x1: int, x2: int, y1: int, y2: int) -> None:
+def download_whole_region(output_dir: Path, wkt_region: str, x1: int, x2: int, y1: int, y2: int) -> list[Path]:
     """Downloads every zip of a given region, to output_dir, and extracts CityGML files."""
 
     wgs84_region = wkt.loads(wkt_region)
     local_region = transform(TO_LOCAL_CRS.transform, wgs84_region)
+    gml_folders = []
     for x in range(x1, x2 + 1, RASTER):
         for y in range(y1, y2 + 1, RASTER):
             tile_center = Point((x + 1) * KILOMETER, (y + 1) * KILOMETER)
@@ -121,14 +122,15 @@ def download_whole_region(output_dir: Path, wkt_region: str, x1: int, x2: int, y
                 continue
             citygml_zip = f"LoD2_{UTM}_{x}_{y}_{RASTER}_{BUNDESLAND}.zip"
             # Oh, zip was deleted, but GML was here
-            citygml_url = f"{CITYGML_SERVER}/{citygml_zip}"
+            remote_zip = f"{CITYGML_SERVER}/{citygml_zip}"
             local_zip = output_dir / citygml_zip
+            gml_folders.append(output_dir / local_zip.stem)
             if local_zip.exists():
                 print(f"  {local_zip.name} already in {output_dir.name}/")
             else:
                 print(f"  Download {citygml_zip} to {output_dir.name}/ ", end="")
                 try:
-                    urllib.request.urlretrieve(citygml_url, local_zip)
+                    urllib.request.urlretrieve(remote_zip, local_zip)
                 except urllib.error.HTTPError as e:
                     print(f"❌ {e}")
                     continue
@@ -141,16 +143,17 @@ def download_whole_region(output_dir: Path, wkt_region: str, x1: int, x2: int, y
                 print("")
                 with zipfile.ZipFile(local_zip, "r") as zip_ref:
                     zip_ref.extractall(output_dir)
+    return gml_folders
 
 
-def extract_region(output_dir: Path, location_name: str, wkt: str) -> None:
+def extract_region(gml_folders: list[Path], output_dir: Path, location_name: str, wkt: str) -> None:
     """Uses RegionChooser to extract a given region from all the CityGML files found in subfolder."""
     output_file = output_dir / (location_name + ".gml")
     if output_file.exists():
         print(f"  {output_file} already exists. Not extracting.")
         return
     region_chooser_libs = Path(SIMSTADT_FOLDER).expanduser() / "lib/*"
-    gml_inputs = list(output_dir.glob(GML_GLOB))
+    gml_inputs = [gml for gml_folder in gml_folders for gml in gml_folder.glob(GML_GLOB)]
     if len(gml_inputs) == 0:
         print("Error: No CityGML found. At least part of the region should be in Baden-Württemberg!")
         return
@@ -215,9 +218,9 @@ def main(regions: dict[str, str]) -> None:
         output_dir.mkdir(parents=True, exist_ok=True)
         wkt = get_wkt(wkt_or_zipcode)
         x1, x2, y1, y2 = wkt_polygon_to_grid_coords(location_name, wkt)
-        download_whole_region(output_dir, wkt, x1, x2, y1, y2)
+        gml_folders = download_whole_region(output_dir, wkt, x1, x2, y1, y2)
         if EXTRACT_REGIONS:
-            extract_region(output_dir, location_name, wkt)
+            extract_region(gml_folders, output_dir, location_name, wkt)
         print()