BERATools 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

beratools/tools/centerline.py

@@ -1,176 +1,136 @@
-import logging
-import time
-
-import sys
-from pathlib import Path
-from inspect import getsourcefile
-
-if __name__ == "__main__":
-    current_file = Path(getsourcefile(lambda: 0)).resolve()
-    btool_dir = current_file.parents[2]
-    sys.path.insert(0, btool_dir.as_posix())
-
-from beratools.core.logger import Logger
-from beratools.core.algo_centerline import *
-from beratools.core.dijkstra_algorithm import *
-from beratools.core.constants import *
-from common import *
-
-log = Logger('centerline', file_level=logging.INFO)
-logger = log.get_logger()
-print = log.print
-
-
-def process_single_line(line_args):
-    line = line_args[0][0]
-    prop = line_args[0][1]
-    line_radius = line_args[1]
-    in_raster = line_args[2]
-    line_id = line_args[3]
-    seed_line = shape(line)  # LineString
-    line_radius = float(line_radius)
-
-    default_return = (seed_line, seed_line, prop, None)
-
-    cost_clip, out_meta = clip_raster(in_raster, seed_line, line_radius)
-
-    if not HAS_COST_RASTER:
-        cost_clip, _ = cost_raster(cost_clip, out_meta)
-
-    lc_path = line
-    try:
-        if CL_USE_SKIMAGE_GRAPH:
-            # skimage shortest path
-            lc_path = find_least_cost_path_skimage(cost_clip, out_meta, seed_line)
-        else:
-            lc_path = find_least_cost_path(cost_clip, out_meta, seed_line)
-    except Exception as e:
-        print(e)
-        return default_return
-
-    if lc_path:
-        lc_path_coords = lc_path.coords
-    else:
-        lc_path_coords = []
-
-    # search for centerline
-    if len(lc_path_coords) < 2:
-        print('No least cost path detected, use input line.')
-        prop['status'] = CenterlineStatus.FAILED.value
-        return default_return
-
-    # get corridor raster
-    lc_path = LineString(lc_path_coords)
-    cost_clip, out_meta = clip_raster(in_raster, lc_path, line_radius * 0.9)
-    if not HAS_COST_RASTER:
-        cost_clip, _ = cost_raster(cost_clip, out_meta)
-
-    out_transform = out_meta['transform']
-    transformer = rasterio.transform.AffineTransformer(out_transform)
-    cell_size = (out_transform[0], -out_transform[4])
-
-    x1, y1 = lc_path_coords[0]
-    x2, y2 = lc_path_coords[-1]
-    source = [transformer.rowcol(x1, y1)]
-    destination = [transformer.rowcol(x2, y2)]
-    corridor_thresh_cl = corridor_raster(cost_clip, out_meta, source, destination,
-                                         cell_size, FP_CORRIDOR_THRESHOLD)
-
-    # find contiguous corridor polygon and extract centerline
-    df = gpd.GeoDataFrame(geometry=[seed_line], crs=out_meta['crs'])
-    corridor_poly_gpd = find_corridor_polygon(corridor_thresh_cl, out_transform, df)
-    center_line, status = find_centerline(corridor_poly_gpd.geometry.iloc[0], lc_path)
-    prop['status'] = status.value
-
-    print(" Searching centerline: line {} ".format(line_id), flush=True)
-    return center_line, lc_path, prop, corridor_poly_gpd
-
-
-def centerline(callback, in_line, in_raster, line_radius,
-               proc_segments, out_line, processes, verbose):
-    if not compare_crs(vector_crs(in_line), raster_crs(in_raster)):
-        print("Line and CHM have different spatial references, please check.")
-        return
-
-    # Read input line features
-    layer_crs = None
-    schema = None
-    input_lines = []
-
-    with fiona.open(in_line) as open_line_file:
-        layer_crs = open_line_file.crs
-        schema = open_line_file.meta['schema']
-        for line in open_line_file:
-            if line.geometry:
-                if line.geometry['type'] != 'MultiLineString':
-                    input_lines.append([line.geometry, line.properties])
-                else:
-                    print('MultiLineString found.')
-                    geoms = shape(line['geometry']).geoms
-                    for item in geoms:
-                        line_part = fiona.Geometry.from_dict(item)
-                        if line_part:
-                            input_lines.append([line_part, line.properties])
-            else:
-                print(f'Line {line.id} has empty geometry.')
-
-    if proc_segments:
-        # split line segments at vertices
-        input_lines_temp = []
-        for line in input_lines:
-            line_seg = line[0]
-            line_prop = line[1]
-            line_segs = segments(line_seg.coordinates)
-            line_feats = [(line, line_prop) for line in line_segs]
-            if line_segs:
-                input_lines_temp.extend(line_feats)
-
-        input_lines = input_lines_temp
-
-    # Process lines
-    all_lines = []
-    i = 0
-    for line in input_lines:
-        all_lines.append((line, line_radius, in_raster, i))
-        i += 1
-
-    print('{} lines to be processed.'.format(len(all_lines)))
-
-    lc_path_geoms = []
-    feat_props = []
-    center_line_geoms = []
-    corridor_poly_list = []
-    result = execute_multiprocessing(process_single_line, all_lines, 'Centerline',
-                                     processes, 1, verbose=verbose)
-
-    for item in result:
-        center_line = item[0]
-        lc_path = item[1]
-        prop = item[2]
-        corridor_poly = item[3]
-
-        if lc_path and prop:
-            lc_path_geoms.append(lc_path)
-            feat_props.append(prop)
-            center_line_geoms.append(center_line)
-            corridor_poly_list.append(corridor_poly)
-
-    out_centerline_path = Path(out_line)
-    schema['properties']['status'] = 'int'
-    save_features_to_shapefile(out_centerline_path.as_posix(), layer_crs, center_line_geoms, feat_props, schema)
-
-    out_least_cost_path = out_centerline_path.with_stem(out_centerline_path.stem + '_least_cost_path')
-    save_features_to_shapefile(out_least_cost_path.as_posix(), layer_crs, lc_path_geoms, feat_props, schema)
-
-    # save corridor polygons
-    corridor_polys = pd.concat(corridor_poly_list)
-    out_corridor_poly_path = Path(out_line)
-    out_corridor_poly_path = out_corridor_poly_path.with_stem(out_corridor_poly_path.stem + '_corridor_poly')
-    corridor_polys.to_file(out_corridor_poly_path.as_posix())
-
-
-if __name__ == '__main__':
-    in_args, in_verbose = check_arguments()
-    start_time = time.time()
-    centerline(print, **in_args.input, processes=int(in_args.processes), verbose=in_verbose)
-    print('Elapsed time: {}'.format(time.time() - start_time))
+"""
+Copyright (C) 2025 Applied Geospatial Research Group.
+
+This script is licensed under the GNU General Public License v3.0.
+See <https://gnu.org/licenses/gpl-3.0> for full license details.
+
+Author: Richard Zeng
+
+Description:
+    This script is part of the BERA Tools.
+    Webpage: https://github.com/appliedgrg/beratools
+
+    The purpose of this script is to provide main interface for centerline tool.
+"""
+
+import logging
+import time
+from pathlib import Path
+
+import pandas as pd
+
+import beratools.core.algo_centerline as algo_centerline
+import beratools.core.algo_common as algo_common
+import beratools.core.constants as bt_const
+import beratools.tools.common as bt_common
+from beratools.core.logger import Logger
+from beratools.core.tool_base import execute_multiprocessing
+
+log = Logger("centerline", file_level=logging.INFO)
+logger = log.get_logger()
+print = log.print
+
+def generate_line_class_list(
+    in_vector, in_raster, line_radius, layer=None, proc_segments=True
+) -> list:
+    line_classes = []
+    line_list = algo_common.prepare_lines_gdf(in_vector, layer, proc_segments)
+
+    for item in line_list:
+        line_classes.append(
+            algo_centerline.SeedLine(item, in_raster, proc_segments, line_radius)
+        )
+
+    return line_classes
+
+
+def process_single_line_class(seed_line):
+    seed_line.compute()
+    return seed_line
+
+
+def centerline(
+    in_line,
+    in_raster,
+    line_radius,
+    proc_segments,
+    out_line,
+    processes,
+    verbose,
+    in_layer=None,
+    out_layer=None,
+    parallel_mode=bt_const.ParallelMode.MULTIPROCESSING
+):
+    if not bt_common.compare_crs(
+        bt_common.vector_crs(in_line), bt_common.raster_crs(in_raster)
+    ):
+        print("Line and CHM have different spatial references, please check.")
+        return
+
+    line_class_list = generate_line_class_list(
+        in_line,
+        in_raster,
+        line_radius=float(line_radius),
+        layer=in_layer,
+        proc_segments=proc_segments,
+    )
+
+    print("{} lines to be processed.".format(len(line_class_list)))
+
+    lc_path_list = []
+    centerline_list = []
+    corridor_poly_list = []
+    result = execute_multiprocessing(
+        process_single_line_class,
+        line_class_list,
+        "Centerline",
+        processes,
+        verbose=verbose,
+        mode=parallel_mode,
+    )
+    if not result:
+        print("No centerlines found.")
+        return
+
+    for item in result:
+        lc_path_list.append(item.lc_path)
+        centerline_list.append(item.centerline)
+        corridor_poly_list.append(item.corridor_poly_gpd)
+
+    # Concatenate the lists of GeoDataFrames into single GeoDataFrames
+    if (
+        len(lc_path_list) == 0
+        or len(centerline_list) == 0
+        or len(corridor_poly_list) == 0
+    ):
+        print("No centerline generated.")
+        return 1
+
+    lc_path_list = pd.concat(lc_path_list, ignore_index=True)
+    centerline_list = pd.concat(centerline_list, ignore_index=True)
+    corridor_polys = pd.concat(corridor_poly_list, ignore_index=True)
+
+    # Save the concatenated GeoDataFrames to the shapefile/gpkg
+    centerline_list.to_file(out_line, layer=out_layer)
+
+    # Check if the output file is a shapefile
+    out_line_path = Path(out_line)
+
+    if out_line_path.suffix == ".shp":
+        # Generate the new file name for the GeoPackage with '_aux' appended
+        aux_file = out_line_path.with_name(out_line_path.stem + "_aux.gpkg")
+        print(f"Saved auxiliary data to: {aux_file}")
+    else:
+        aux_file = out_line  # continue using out_line (gpkg)
+
+    # Save lc_path_list and corridor_polys to the new GeoPackage with '_aux' suffix
+    lc_path_list.to_file(aux_file, layer="least_cost_path")
+    corridor_polys.to_file(aux_file, layer="corridor_polygon")
+
+
+# TODO: fix geometries when job done
+if __name__ == "__main__":
+    in_args, in_verbose = bt_common.check_arguments()
+    start_time = time.time()
+    centerline(**in_args.input, processes=int(in_args.processes), verbose=in_verbose)
+    print("Elapsed time: {}".format(time.time() - start_time))