299 lines
11 KiB
Python
299 lines
11 KiB
Python
from __future__ import annotations
|
|
import base64
|
|
import io
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.execution_context import emit_mesh, emit_warning
|
|
from backend.data_types import (
|
|
COLORMAPS,
|
|
DataField,
|
|
ImageData,
|
|
MeshModel,
|
|
_annotation_context_from_field,
|
|
colormap_to_uint8,
|
|
normalize_for_colormap,
|
|
resolve_colormap_input,
|
|
)
|
|
|
|
|
|
def _darken_colors(colors: np.ndarray, factor: float) -> np.ndarray:
|
|
return np.clip(np.rint(colors.astype(np.float32) * factor), 0, 255).astype(np.uint8)
|
|
|
|
|
|
def _grid_triangle_indices(nx: int, ny: int, *, reverse: bool = False) -> list[list[int]]:
|
|
faces: list[list[int]] = []
|
|
for iy in range(ny - 1):
|
|
for ix in range(nx - 1):
|
|
a = iy * nx + ix
|
|
b = a + 1
|
|
c = a + nx
|
|
d = c + 1
|
|
if reverse:
|
|
faces.append([a, b, c])
|
|
faces.append([b, d, c])
|
|
else:
|
|
faces.append([a, c, b])
|
|
faces.append([b, c, d])
|
|
return faces
|
|
|
|
|
|
def _surface_extent_scale(xreal: float, yreal: float, nx: int, ny: int) -> tuple[float, float]:
|
|
def _resolve_span(value: float, fallback_points: int) -> float:
|
|
try:
|
|
span = abs(float(value))
|
|
except (TypeError, ValueError):
|
|
span = 0.0
|
|
if not np.isfinite(span) or span <= 0.0:
|
|
span = float(max(fallback_points - 1, 1))
|
|
return span
|
|
|
|
x_span = _resolve_span(xreal, nx)
|
|
y_span = _resolve_span(yreal, ny)
|
|
max_span = max(x_span, y_span)
|
|
if not np.isfinite(max_span) or max_span <= 0.0:
|
|
max_span = 1.0
|
|
return (x_span / max_span, y_span / max_span)
|
|
|
|
|
|
def _build_mesh_model(
|
|
z: np.ndarray,
|
|
colors_u8: np.ndarray,
|
|
z_scale: float,
|
|
make_solid: bool,
|
|
lateral_extent: tuple[float, float] = (1.0, 1.0),
|
|
) -> MeshModel:
|
|
ny, nx = z.shape
|
|
zmin = float(z.min())
|
|
zmax = float(z.max())
|
|
z_range = zmax - zmin if zmax != zmin else 1.0
|
|
x_extent, y_extent = lateral_extent
|
|
|
|
top_vertices = np.empty((nx * ny, 3), dtype=np.float32)
|
|
top_colors = colors_u8.reshape(-1, 3).astype(np.uint8)
|
|
for iy in range(ny):
|
|
py = (iy / max(ny - 1, 1) - 0.5) * y_extent
|
|
for ix in range(nx):
|
|
idx = iy * nx + ix
|
|
px = (ix / max(nx - 1, 1) - 0.5) * x_extent
|
|
pz = ((float(z[iy, ix]) - zmin) / z_range - 0.5) * z_scale
|
|
top_vertices[idx] = (px, pz, py)
|
|
|
|
faces = _grid_triangle_indices(nx, ny)
|
|
if not make_solid:
|
|
return MeshModel(vertices=top_vertices, faces=np.asarray(faces, dtype=np.int32), colors=top_colors)
|
|
|
|
base_y = float(top_vertices[:, 1].min())
|
|
bottom_vertices = top_vertices.copy()
|
|
bottom_vertices[:, 1] = base_y
|
|
bottom_colors = _darken_colors(top_colors, 0.35)
|
|
|
|
vertices = np.vstack([top_vertices, bottom_vertices]).astype(np.float32)
|
|
colors = np.vstack([top_colors, bottom_colors]).astype(np.uint8)
|
|
|
|
bottom_offset = len(top_vertices)
|
|
faces.extend([[a + bottom_offset, b + bottom_offset, c + bottom_offset] for a, b, c in _grid_triangle_indices(nx, ny, reverse=True)])
|
|
|
|
def _add_wall(a: int, b: int):
|
|
faces.append([a, a + bottom_offset, b])
|
|
faces.append([b, a + bottom_offset, b + bottom_offset])
|
|
|
|
for ix in range(nx - 1):
|
|
_add_wall(ix, ix + 1)
|
|
top_row = (ny - 1) * nx
|
|
_add_wall(top_row + ix + 1, top_row + ix)
|
|
for iy in range(ny - 1):
|
|
_add_wall((iy + 1) * nx, iy * nx)
|
|
_add_wall(iy * nx + (nx - 1), (iy + 1) * nx + (nx - 1))
|
|
|
|
return MeshModel(vertices=vertices, faces=np.asarray(faces, dtype=np.int32), colors=colors)
|
|
|
|
|
|
@register_node(display_name="3D View")
|
|
class View3D:
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD", {"label": "mesh"}),
|
|
"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}),
|
|
"make_solid": ("BOOLEAN", {"default": False}),
|
|
"camera_azimuth": ("FLOAT", {"default": 0.0, "hidden": True}),
|
|
"camera_polar": ("FLOAT", {"default": 1.1, "hidden": True}),
|
|
"camera_distance": ("FLOAT", {"default": 1.8, "hidden": True}),
|
|
"camera_target_x": ("FLOAT", {"default": 0.0, "hidden": True}),
|
|
"camera_target_y": ("FLOAT", {"default": 0.0, "hidden": True}),
|
|
"camera_target_z": ("FLOAT", {"default": 0.0, "hidden": True}),
|
|
"viewport_snapshot": ("STRING", {"default": "", "hidden": True}),
|
|
},
|
|
"optional": {
|
|
"map_field": ("DATA_FIELD", {"label": "map"}),
|
|
"colormap_map": ("COLORMAP", {"label": "colormap"}),
|
|
},
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('MESH_MODEL', 'mesh'),
|
|
('IMAGE', 'viewport'),
|
|
)
|
|
FUNCTION = "render"
|
|
|
|
OUTPUT_NODE = True
|
|
DESCRIPTION = (
|
|
"Interactive 3D surface view of a DATA_FIELD. "
|
|
"Use the mesh input for geometry and optionally a second map input for coloring. "
|
|
"Drag to rotate, middle-drag to pan, and right-drag or scroll to zoom. z_scale exaggerates height."
|
|
)
|
|
|
|
KEYWORDS = ("surface", "mesh", "render", "perspective", "height map", "visualize", "opengl")
|
|
|
|
def render(
|
|
self, field: DataField,
|
|
colormap: str, z_scale: float, resolution: int, make_solid: bool = False,
|
|
camera_azimuth: float = 0.0, camera_polar: float = 1.1, camera_distance: float = 1.8,
|
|
camera_target_x: float = 0.0, camera_target_y: float = 0.0, camera_target_z: float = 0.0,
|
|
viewport_snapshot: str = "",
|
|
map_field: DataField | None = None, colormap_map=None,
|
|
) -> tuple:
|
|
from scipy.ndimage import map_coordinates
|
|
|
|
data = field.data
|
|
yres, xres = data.shape
|
|
|
|
phys_ratio = field.xreal / field.yreal if field.yreal else 1.0
|
|
pixel_ratio = xres / yres if yres else 1.0
|
|
if abs(phys_ratio / pixel_ratio - 1.0) > 0.02:
|
|
emit_warning(
|
|
f"Non-square pixels ({xres}\u00d7{yres} px). "
|
|
f"The 3D surface shows the physical scan area, not the pixel grid."
|
|
)
|
|
|
|
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())
|
|
color_field = map_field if map_field is not None else field
|
|
color_data = color_field.data
|
|
|
|
if color_field is field and color_data.shape == z.shape:
|
|
color_samples = z
|
|
elif color_field is field:
|
|
color_samples = color_data[::step_y, ::step_x].astype(np.float32)
|
|
else:
|
|
x_phys = np.linspace(field.xoff, field.xoff + field.xreal, nx, dtype=np.float64)
|
|
y_phys = np.linspace(field.yoff, field.yoff + field.yreal, ny, dtype=np.float64)
|
|
grid_y, grid_x = np.meshgrid(y_phys, x_phys, indexing="ij")
|
|
|
|
map_x = np.clip(
|
|
(grid_x - color_field.xoff) / max(color_field.xreal, 1e-12) * max(color_field.xres - 1, 0),
|
|
0.0,
|
|
max(color_field.xres - 1, 0),
|
|
)
|
|
map_y = np.clip(
|
|
(grid_y - color_field.yoff) / max(color_field.yreal, 1e-12) * max(color_field.yres - 1, 0),
|
|
0.0,
|
|
max(color_field.yres - 1, 0),
|
|
)
|
|
color_samples = map_coordinates(
|
|
color_data.astype(np.float64),
|
|
[map_y, map_x],
|
|
order=1,
|
|
mode="nearest",
|
|
).astype(np.float32)
|
|
|
|
z_norm = normalize_for_colormap(
|
|
color_samples,
|
|
offset=color_field.display_offset,
|
|
scale=color_field.display_scale,
|
|
data_min=float(color_field.data.min()),
|
|
data_max=float(color_field.data.max()),
|
|
)
|
|
|
|
resolved_colormap = resolve_colormap_input(
|
|
colormap,
|
|
colormap_input=colormap_map,
|
|
inherited=color_field.colormap,
|
|
default="gray",
|
|
)
|
|
colors_u8 = colormap_to_uint8(z_norm, resolved_colormap)
|
|
surface_extent = _surface_extent_scale(field.xreal, field.yreal, nx, ny)
|
|
mesh_model = _build_mesh_model(
|
|
z,
|
|
colors_u8,
|
|
float(z_scale * 0.1),
|
|
bool(make_solid),
|
|
lateral_extent=surface_extent,
|
|
)
|
|
|
|
z_b64 = base64.b64encode(z.tobytes()).decode()
|
|
colors_b64 = base64.b64encode(colors_u8.tobytes()).decode()
|
|
positions_b64 = base64.b64encode(np.asarray(mesh_model.vertices, dtype=np.float32).tobytes()).decode()
|
|
indices_b64 = base64.b64encode(np.asarray(mesh_model.faces, dtype=np.uint32).tobytes()).decode()
|
|
mesh_colors_b64 = None
|
|
if mesh_model.colors is not None:
|
|
mesh_colors_b64 = base64.b64encode(np.asarray(mesh_model.colors, dtype=np.uint8).tobytes()).decode()
|
|
|
|
mesh_data = {
|
|
"width": nx,
|
|
"height": ny,
|
|
"z_data": z_b64,
|
|
"colors": colors_b64,
|
|
"positions": positions_b64,
|
|
"indices": indices_b64,
|
|
"vertex_colors": mesh_colors_b64,
|
|
"z_min": zmin,
|
|
"z_max": zmax,
|
|
"z_scale": float(z_scale * 0.1),
|
|
"make_solid": bool(make_solid),
|
|
"camera_azimuth": float(camera_azimuth),
|
|
"camera_polar": float(camera_polar),
|
|
"camera_distance": float(camera_distance),
|
|
"camera_target_x": float(camera_target_x),
|
|
"camera_target_y": float(camera_target_y),
|
|
"camera_target_z": float(camera_target_z),
|
|
"x_range": [float(field.xoff), float(field.xoff + field.xreal)],
|
|
"y_range": [float(field.yoff), float(field.yoff + field.yreal)],
|
|
"surface_extent_x": float(surface_extent[0]),
|
|
"surface_extent_y": float(surface_extent[1]),
|
|
}
|
|
|
|
emit_mesh(mesh_data)
|
|
|
|
annotation_context = _annotation_context_from_field(color_field, resolved_colormap)
|
|
annotation_context["xreal"] = float(field.xreal)
|
|
annotation_context["si_unit_xy"] = str(field.si_unit_xy)
|
|
viewport_image = ImageData(
|
|
self._decode_viewport_snapshot(viewport_snapshot),
|
|
metadata={
|
|
"annotation_context": annotation_context,
|
|
"viewport_camera": {
|
|
"azimuth": float(camera_azimuth),
|
|
"polar": float(camera_polar),
|
|
"distance": float(camera_distance),
|
|
"target_x": float(camera_target_x),
|
|
"target_y": float(camera_target_y),
|
|
"target_z": float(camera_target_z),
|
|
},
|
|
},
|
|
)
|
|
return (mesh_model, viewport_image)
|
|
|
|
def _decode_viewport_snapshot(self, snapshot: str) -> np.ndarray:
|
|
text = str(snapshot or "").strip()
|
|
if not text.startswith("data:image/"):
|
|
return np.zeros((1, 1, 3), dtype=np.uint8)
|
|
|
|
try:
|
|
header, payload = text.split(",", 1)
|
|
raw = base64.b64decode(payload)
|
|
from PIL import Image
|
|
image = Image.open(io.BytesIO(raw)).convert("RGB")
|
|
return np.asarray(image, dtype=np.uint8)
|
|
except Exception:
|
|
return np.zeros((1, 1, 3), dtype=np.uint8)
|