78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
"""Grain marking — mark grains by height, slope, or curvature criteria."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from scipy.ndimage import label, sobel
|
|
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField
|
|
from backend.nodes.helpers import bool_to_mask
|
|
|
|
|
|
@register_node(display_name="Grain Mark")
|
|
class GrainMark:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"criterion": (["height", "slope", "curvature"], {"default": "height"}),
|
|
"threshold_low": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"threshold_high": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"min_size": ("INT", {"default": 10, "min": 1, "max": 100000, "step": 1}),
|
|
"inverted": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('IMAGE', 'mask'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Mark grains by thresholding height, slope magnitude, or curvature. "
|
|
"Thresholds are relative (0-1) to the data range. Small regions below "
|
|
"min_size pixels are removed. Use inverted to mark valleys instead of peaks. "
|
|
)
|
|
|
|
KEYWORDS = ("threshold", "segment", "peak", "particle")
|
|
|
|
def process(self, field: DataField, criterion: str, threshold_low: float,
|
|
threshold_high: float, min_size: int, inverted: bool) -> tuple:
|
|
data = np.asarray(field.data, dtype=np.float64)
|
|
|
|
if criterion == "height":
|
|
values = data
|
|
elif criterion == "slope":
|
|
gx = sobel(data, axis=1)
|
|
gy = sobel(data, axis=0)
|
|
values = np.sqrt(gx**2 + gy**2)
|
|
elif criterion == "curvature":
|
|
gxx = sobel(sobel(data, axis=1), axis=1)
|
|
gyy = sobel(sobel(data, axis=0), axis=0)
|
|
values = np.abs(gxx + gyy)
|
|
else:
|
|
raise ValueError(f"Unknown criterion: {criterion!r}")
|
|
|
|
# Normalize to [0, 1]
|
|
vmin, vmax = values.min(), values.max()
|
|
if vmax > vmin:
|
|
norm = (values - vmin) / (vmax - vmin)
|
|
else:
|
|
norm = np.zeros_like(values)
|
|
|
|
# Apply thresholds
|
|
binary = (norm >= threshold_low) & (norm <= threshold_high)
|
|
|
|
if inverted:
|
|
binary = ~binary
|
|
|
|
# Remove small regions
|
|
labeled, n_labels = label(binary.astype(np.int32))
|
|
for gid in range(1, n_labels + 1):
|
|
if (labeled == gid).sum() < min_size:
|
|
binary[labeled == gid] = False
|
|
|
|
return (bool_to_mask(binary),)
|