Hugging Face's logo

Spaces:
ayoubft
/
fractex2D_tuto
Running

App Files Community
fractex2D_tuto
File size: 4,452 Bytes
fa7174b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
"""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)