49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
from __future__ import annotations
|
|
import json
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import COLORMAPS, DEFAULT_CUSTOM_COLORMAP_STOPS, normalize_colormap_spec
|
|
|
|
|
|
@register_node(display_name="Color Map")
|
|
class ColorMap:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mode": (["preset", "custom"], {"default": "preset"}),
|
|
"preset": (list(COLORMAPS), {
|
|
"default": "viridis",
|
|
"show_when_widget_value": {"mode": ["preset"]},
|
|
}),
|
|
"stops": ("STRING", {
|
|
"default": json.dumps(list(DEFAULT_CUSTOM_COLORMAP_STOPS)),
|
|
"colormap_stops": True,
|
|
"show_when_widget_value": {"mode": ["custom"]},
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("COLORMAP",)
|
|
RETURN_NAMES = ("colormap",)
|
|
FUNCTION = "build"
|
|
|
|
DESCRIPTION = (
|
|
"Build a reusable colormap. Choose a preset, or create a custom gradient with min/max colours "
|
|
"and any number of intermediate stops."
|
|
)
|
|
|
|
def build(self, mode: str, preset: str, stops: str | None = None, stops_json: str | None = None) -> tuple:
|
|
if mode == "preset":
|
|
return ({"mode": "preset", "preset": normalize_colormap_spec(preset)},)
|
|
|
|
try:
|
|
raw_stops = stops if stops is not None else stops_json
|
|
stops_data = json.loads(raw_stops or "[]")
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError("Custom colormap stops must be valid JSON.") from exc
|
|
|
|
spec = normalize_colormap_spec({"mode": "custom", "stops": stops_data}, fallback=None)
|
|
if not (isinstance(spec, dict) and spec.get("mode") == "custom"):
|
|
raise ValueError("Custom colormap must include at least min and max colours.")
|
|
return (spec,)
|