91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import (
|
|
COLORMAPS,
|
|
DataField,
|
|
colormap_to_uint8,
|
|
normalize_for_colormap,
|
|
resolve_colormap_input,
|
|
)
|
|
|
|
|
|
@register_node(display_name="3D View")
|
|
class View3D:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"colormap": (["auto"] + list(COLORMAPS), {"hide_when_input_connected": "colormap_map"}),
|
|
"z_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10.0, "step": 0.05}),
|
|
"resolution": ("INT", {"default": 128, "min": 32, "max": 512, "step": 16}),
|
|
},
|
|
"optional": {
|
|
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "render"
|
|
|
|
OUTPUT_NODE = True
|
|
DESCRIPTION = (
|
|
"Interactive 3D surface view of a DATA_FIELD. "
|
|
"Drag to rotate, scroll to zoom. z_scale exaggerates height."
|
|
)
|
|
|
|
_broadcast_mesh_fn = None
|
|
_current_node_id: str = ""
|
|
|
|
def render(
|
|
self, field: DataField,
|
|
colormap: str, z_scale: float, resolution: int, colormap_map=None,
|
|
) -> tuple:
|
|
import base64
|
|
|
|
data = field.data
|
|
yres, xres = data.shape
|
|
|
|
step_y = max(1, yres // resolution)
|
|
step_x = max(1, xres // resolution)
|
|
z = data[::step_y, ::step_x].astype(np.float32)
|
|
ny, nx = z.shape
|
|
|
|
zmin, zmax = float(z.min()), float(z.max())
|
|
z_norm = normalize_for_colormap(
|
|
z,
|
|
offset=field.display_offset,
|
|
scale=field.display_scale,
|
|
data_min=float(field.data.min()),
|
|
data_max=float(field.data.max()),
|
|
)
|
|
|
|
resolved_colormap = resolve_colormap_input(
|
|
colormap,
|
|
colormap_input=colormap_map,
|
|
inherited=field.colormap,
|
|
default="gray",
|
|
)
|
|
colors_u8 = colormap_to_uint8(z_norm, resolved_colormap)
|
|
|
|
z_b64 = base64.b64encode(z.tobytes()).decode()
|
|
colors_b64 = base64.b64encode(colors_u8.tobytes()).decode()
|
|
|
|
mesh_data = {
|
|
"width": nx,
|
|
"height": ny,
|
|
"z_data": z_b64,
|
|
"colors": colors_b64,
|
|
"z_min": zmin,
|
|
"z_max": zmax,
|
|
"z_scale": float(z_scale * 0.1),
|
|
"x_range": [float(field.xoff), float(field.xoff + field.xreal)],
|
|
"y_range": [float(field.yoff), float(field.yoff + field.yreal)],
|
|
}
|
|
|
|
if View3D._broadcast_mesh_fn is not None:
|
|
View3D._broadcast_mesh_fn(View3D._current_node_id, mesh_data)
|
|
|
|
return ()
|