47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
"""Median background subtraction — extract and subtract background using local median."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from scipy.ndimage import median_filter
|
|
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField
|
|
|
|
|
|
@register_node(display_name="Median Background")
|
|
class MedianBackground:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"radius": ("INT", {"default": 20, "min": 2, "max": 500, "step": 1}),
|
|
"output": (["subtracted", "background"], {"default": "subtracted"}),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('DATA_FIELD', 'result'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Extract background using a local median filter and subtract it. "
|
|
"The radius controls the filter window size — larger values capture "
|
|
"broader background variations. More robust than polynomial leveling "
|
|
"for surfaces with sparse tall features. "
|
|
)
|
|
|
|
KEYWORDS = ("rolling ball", "flatten", "level", "subtract", "baseline")
|
|
|
|
def process(self, field: DataField, radius: int, output: str) -> tuple:
|
|
data = np.asarray(field.data, dtype=np.float64)
|
|
size = 2 * radius + 1
|
|
background = median_filter(data, size=size)
|
|
|
|
if output == "background":
|
|
return (field.replace(data=background),)
|
|
else:
|
|
return (field.replace(data=data - background),)
|