adding more nodes
This commit is contained in:
57
backend/nodes/laplace_interpolation.py
Normal file
57
backend/nodes/laplace_interpolation.py
Normal file
@@ -0,0 +1,57 @@
|
||||
"""Laplace interpolation — fill masked regions by solving the Laplace equation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import DataField
|
||||
from backend.nodes.helpers import mask_to_bool
|
||||
|
||||
|
||||
@register_node(display_name="Laplace Interpolation")
|
||||
class LaplaceInterpolation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"field": ("DATA_FIELD",),
|
||||
"mask": ("IMAGE",),
|
||||
"iterations": ("INT", {"default": 500, "min": 10, "max": 10000, "step": 10}),
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUTS = (
|
||||
('DATA_FIELD', 'filled'),
|
||||
)
|
||||
FUNCTION = "process"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Fill masked (missing) regions by solving the Laplace equation with "
|
||||
"Dirichlet boundary conditions from surrounding pixels. "
|
||||
"Produces a smooth, harmonic interpolation without overshooting. "
|
||||
)
|
||||
|
||||
def process(self, field: DataField, mask: np.ndarray, iterations: int) -> tuple:
|
||||
data = np.asarray(field.data, dtype=np.float64).copy()
|
||||
hole = mask_to_bool(mask)
|
||||
|
||||
if not hole.any():
|
||||
return (field.replace(data=data),)
|
||||
|
||||
# Initialize masked pixels to mean of unmasked neighbours or global mean
|
||||
valid_mean = data[~hole].mean() if (~hole).any() else 0.0
|
||||
data[hole] = valid_mean
|
||||
|
||||
# Iterative Jacobi relaxation: replace each masked pixel with
|
||||
# the mean of its 4-connected neighbours
|
||||
padded = np.pad(data, 1, mode='edge')
|
||||
hole_padded = np.pad(hole, 1, mode='constant', constant_values=False)
|
||||
|
||||
for _ in range(iterations):
|
||||
avg = (padded[:-2, 1:-1] + padded[2:, 1:-1] +
|
||||
padded[1:-1, :-2] + padded[1:-1, 2:]) / 4.0
|
||||
padded[1:-1, 1:-1][hole] = avg[hole]
|
||||
|
||||
data = padded[1:-1, 1:-1].copy()
|
||||
return (field.replace(data=data),)
|
||||
Reference in New Issue
Block a user