Spaces:
Running
Running
| """Stage 2: vectorize a probability raster into fracture centerlines (GeoPackage). | |
| Ported from FraXtex's `fraxtex_tools/vectorize.py`, which is canonical for this | |
| pipeline (see CLAUDE.md). Kept line-comparable with the plugin; the only changes are | |
| the QGIS parameter plumbing and the output path, which the caller supplies. | |
| Expects the inference convention: **high values mean fracture**. | |
| All distances are given in pixels and multiplied by the raster resolution, so the | |
| same parameter values behave the same on a georeferenced raster and on a plain PNG | |
| (where the resolution is 1 and the units are pixels). | |
| """ | |
| import geopandas as gpd | |
| import numpy as np | |
| import rasterio | |
| import shapely | |
| import skimage.filters as skf | |
| import skimage.morphology as skmp | |
| from rasterio import features | |
| from shapely.affinity import affine_transform | |
| # Keep in the same order as the UI labels. | |
| THRESH_METHODS = ["mean", "adaptive", "otsu", "manual"] | |
| SKEL_METHODS = ["zhang", "lee", "medial_axis"] | |
| def vectorize( | |
| raster_path, | |
| output_path, | |
| feedback, | |
| thresh_method="mean", | |
| manual_val=127, | |
| min_size=142, | |
| skel_method="zhang", | |
| buffer_px=3, | |
| min_branch=10, | |
| simplify_px=10, | |
| ): | |
| """Threshold -> skeletonize -> polygonize -> centerline -> GeoPackage. | |
| Returns `(output_path, feature_count)`. | |
| """ | |
| with rasterio.open(raster_path) as src: | |
| arr = src.read(1) | |
| crs = src.crs | |
| transform = src.transform | |
| res_x = abs(transform.a) | |
| feedback.pushInfo(f"Shape: {arr.shape}, res: {res_x:.6f}") | |
| # 1. Threshold | |
| feedback.setProgress(5) | |
| if thresh_method == "otsu": | |
| threshold = skf.threshold_otsu(arr) | |
| elif thresh_method == "adaptive": | |
| threshold = skf.threshold_local(arr, block_size=51, method="gaussian") | |
| elif thresh_method == "mean": | |
| threshold = skf.threshold_mean(arr) | |
| else: | |
| threshold = manual_val | |
| # An adaptive threshold is a whole array, so summarise rather than print it. | |
| if np.ndim(threshold) == 0: | |
| feedback.pushInfo(f"Method: {thresh_method}, threshold: {threshold:.3f}") | |
| else: | |
| feedback.pushInfo( | |
| f"Method: {thresh_method}, per-pixel threshold " | |
| f"(mean {float(np.mean(threshold)):.3f})" | |
| ) | |
| binary = (arr > threshold).astype(np.uint8) | |
| # 2. Remove small objects | |
| feedback.setProgress(15) | |
| binary = skmp.remove_small_objects(binary.astype(bool), min_size=min_size).astype(np.uint8) | |
| # 3. Skeletonize | |
| feedback.setProgress(25) | |
| if skel_method == "zhang": | |
| skel = skmp.skeletonize(binary, method="zhang") | |
| elif skel_method == "lee": | |
| skel = skmp.skeletonize(binary, method="lee") | |
| else: | |
| skel, _ = skmp.medial_axis(binary, return_distance=True) | |
| skel = skel.astype(np.uint8) | |
| # 4. Raster -> vector | |
| feedback.setProgress(40) | |
| polys = [shapely.geometry.shape(c) for c, v in features.shapes(skel, skel > 0.5)] | |
| if not polys: | |
| feedback.pushWarning( | |
| "No fractures detected. Try a lower threshold or a smaller minimum " | |
| "object size." | |
| ) | |
| gpd.GeoDataFrame(geometry=[], crs=crs).to_file(output_path, driver="GPKG") | |
| return output_path, 0 | |
| gdf = gpd.GeoDataFrame(geometry=polys, crs=crs) | |
| # 5. Affine transform to map coordinates | |
| aff = [transform.a, transform.b, transform.d, transform.e, transform.c, transform.f] | |
| gdf["geometry"] = gdf["geometry"].apply(lambda g: affine_transform(g, aff)) | |
| # 6. Buffer -> dissolve -> explode | |
| feedback.setProgress(55) | |
| gdf["geometry"] = gdf.buffer(res_x * buffer_px, join_style="mitre") | |
| gdf = gdf.dissolve().explode() | |
| # 7. Centerline | |
| feedback.setProgress(70) | |
| try: | |
| import pygeoops | |
| gdf["geometry"] = pygeoops.centerline( | |
| gdf["geometry"], | |
| densify_distance=res_x, | |
| min_branch_length=res_x * min_branch, | |
| simplifytolerance=res_x, | |
| ) | |
| except ImportError: | |
| feedback.pushInfo("pygeoops not available, skipping centerline.") | |
| # 8. Simplify | |
| feedback.setProgress(85) | |
| gdf = gdf.simplify(res_x * simplify_px, preserve_topology=False) | |
| gdf = gdf[~gdf.is_empty] | |
| gdf = gpd.GeoDataFrame(geometry=list(gdf), crs=crs) | |
| feedback.setProgress(95) | |
| gdf.to_file(output_path, driver="GPKG") | |
| feedback.pushInfo(f"Vectorize done: {len(gdf)} features.") | |
| feedback.setProgress(100) | |
| return output_path, len(gdf) | |