refactor nodes into standalone file
This commit is contained in:
37
backend/nodes/fix_zero.py
Normal file
37
backend/nodes/fix_zero.py
Normal file
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
import numpy as np
|
||||
from backend.node_registry import register_node
|
||||
from backend.data_types import DataField
|
||||
|
||||
|
||||
@register_node(display_name="Fix Zero")
|
||||
class FixZero:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"field": ("DATA_FIELD",),
|
||||
"method": (["min", "mean", "median"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("DATA_FIELD",)
|
||||
RETURN_NAMES = ("zeroed",)
|
||||
FUNCTION = "process"
|
||||
|
||||
DESCRIPTION = (
|
||||
"Shift data so that the minimum (or mean/median) is zero. "
|
||||
"Equivalent to fix_zero in Gwyddion's level.c."
|
||||
)
|
||||
|
||||
def process(self, field: DataField, method: str) -> tuple:
|
||||
data = field.data.copy()
|
||||
if method == "min":
|
||||
data -= data.min()
|
||||
elif method == "mean":
|
||||
data -= data.mean()
|
||||
elif method == "median":
|
||||
data -= np.median(data)
|
||||
else:
|
||||
raise ValueError(f"Unknown method: {method}")
|
||||
return (field.replace(data=data),)
|
||||
Reference in New Issue
Block a user