84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""Distribution coercion — transform data to match a target distribution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from math import ceil
|
|
from scipy.stats import norm
|
|
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField
|
|
|
|
|
|
def _coerce_block(data: np.ndarray, distribution: str, n_levels: int) -> np.ndarray:
|
|
"""Coerce a flat or 2-D block to the target distribution, returning same shape."""
|
|
shape = data.shape
|
|
flat = data.ravel().astype(np.float64)
|
|
n_pixels = flat.size
|
|
if n_pixels == 0:
|
|
return data.copy()
|
|
|
|
indices = np.argsort(flat, kind="mergesort")
|
|
|
|
if distribution == "uniform":
|
|
target = np.linspace(float(flat.min()), float(flat.max()), n_pixels)
|
|
elif distribution == "gaussian":
|
|
eps = 0.5 / n_pixels
|
|
quantiles = np.linspace(eps, 1.0 - eps, n_pixels)
|
|
target = norm.ppf(quantiles) * float(flat.std()) + float(flat.mean())
|
|
elif distribution == "levels":
|
|
n_levels = max(2, int(n_levels))
|
|
level_values = np.linspace(float(flat.min()), float(flat.max()), n_levels)
|
|
target = np.repeat(level_values, ceil(n_pixels / n_levels))[:n_pixels]
|
|
else:
|
|
raise ValueError(f"Unknown distribution: {distribution}")
|
|
|
|
result = np.empty_like(flat)
|
|
result[indices] = target
|
|
return result.reshape(shape)
|
|
|
|
|
|
@register_node(display_name="Distribution Coercion")
|
|
class DistributionCoercion:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"distribution": (["uniform", "gaussian", "levels"], {"default": "uniform"}),
|
|
"n_levels": ("INT", {
|
|
"default": 4,
|
|
"min": 2,
|
|
"max": 1000,
|
|
"show_when_widget_value": {"distribution": ["levels"]},
|
|
}),
|
|
"processing": (["field", "rows"], {"default": "field"}),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('DATA_FIELD', 'result'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Transform pixel values so their distribution matches a target shape "
|
|
"(uniform, Gaussian, or discrete levels) using rank-based reassignment. "
|
|
"Equivalent to Gwyddion's coerce.c module."
|
|
)
|
|
|
|
KEYWORDS = ("coerce", "histogram matching", "equalize", "uniform", "gaussian", "quantize")
|
|
|
|
def process(self, field: DataField, distribution: str, n_levels: int,
|
|
processing: str) -> tuple:
|
|
data = np.asarray(field.data, dtype=np.float64)
|
|
|
|
if processing == "rows":
|
|
result = np.empty_like(data)
|
|
for i in range(data.shape[0]):
|
|
result[i] = _coerce_block(data[i], distribution, n_levels)
|
|
else:
|
|
result = _coerce_block(data, distribution, n_levels)
|
|
|
|
return (field.replace(data=result),)
|