159 lines
5.8 KiB
Python
159 lines
5.8 KiB
Python
from __future__ import annotations
|
|
from typing import Any
|
|
import numpy as np
|
|
|
|
from ..bpy_ns_types import NfBoolProperty, NfEnumProperty, notify_tree_changed
|
|
from ....registry import register_class
|
|
from .base import SliceNode, SlicingCollection, get_all_children, SliceNodeJoiner, SlicingGroup, SlicingObject
|
|
|
|
@register_class
|
|
class SLICENODE_GetMesh(SliceNode):
|
|
bl_label = "Get Mesh"
|
|
|
|
include_children: NfBoolProperty(name="Include Children", default=True) # pyright: ignore[reportInvalidTypeForm]
|
|
|
|
def init(self, context):
|
|
self.add_input("SLICESOCK_Object")
|
|
self.add_output("SLICESOCK_Mesh")
|
|
|
|
def draw_buttons(self, context, layout):
|
|
assert layout
|
|
layout.prop(self, "include_children")
|
|
|
|
def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingCollection | None:
|
|
if socket_name != "Mesh": return None
|
|
|
|
obj: Any = self.resolve("Object")
|
|
if not obj: return None
|
|
|
|
family = [obj] + (get_all_children(obj) if self.include_children else [])
|
|
|
|
sc: SlicingCollection = SlicingCollection(family, obj.name)
|
|
|
|
return sc
|
|
|
|
@register_class
|
|
class SLICENODE_JoinMesh(SliceNodeJoiner):
|
|
bl_label = "Join Mesh Parts"
|
|
socket_type = "SLICESOCK_Mesh"
|
|
|
|
def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingCollection:
|
|
from ....infra.blender.blender_slice_nodes_generic import resolve_input
|
|
merged_coll = resolve_input(self.inputs[0]) or SlicingCollection([],'')
|
|
|
|
for input in self.inputs[1:]:
|
|
coll = resolve_input(input)
|
|
if coll:
|
|
merged_coll.objects += coll.objects
|
|
|
|
return merged_coll
|
|
|
|
@register_class
|
|
class SLICENODE_JoinMeshGroups(SliceNodeJoiner):
|
|
bl_label = "Join Mesh Groups"
|
|
socket_type = "SLICESOCK_Mesh"
|
|
|
|
def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingGroup:
|
|
from ....infra.blender.blender_slice_nodes_generic import resolve_input
|
|
group: SlicingGroup = SlicingGroup([])
|
|
|
|
for input in self.inputs:
|
|
coll = resolve_input(input)
|
|
if coll:
|
|
group.collections[coll.objects[0].name] = coll
|
|
|
|
return group
|
|
|
|
@register_class
|
|
class SLICENODE_SeparateMesh(SliceNode):
|
|
bl_label = "Separate Mesh"
|
|
|
|
comparison: NfEnumProperty(
|
|
name="Comparison",
|
|
items=[
|
|
('EQUAL', "Equal", "=="),
|
|
('NOT_EQUAL', "Not Equal", "!="),
|
|
('GREATER_THAN', "Greater Than", ">"),
|
|
('LESS_THAN', "Less Than", "<"),
|
|
('GREATER_EQUAL', "Greater Equal", ">="),
|
|
('LESS_EQUAL', "Less Equal", "<="),
|
|
],
|
|
default='EQUAL',
|
|
update=notify_tree_changed
|
|
) # pyright: ignore[reportInvalidTypeForm]
|
|
|
|
def init(self, context):
|
|
self.add_input("SLICESOCK_Mesh")
|
|
self.add_input("SLICESOCK_String", "Attribute")
|
|
self.add_input("SLICESOCK_Float", "Value")
|
|
self.add_output("SLICESOCK_Mesh", "Selection")
|
|
self.add_output("SLICESOCK_Mesh", "Inverted")
|
|
|
|
def draw_buttons(self, context, layout):
|
|
assert layout
|
|
layout.prop(self, "comparison", text="")
|
|
|
|
def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingCollection | None:
|
|
coll: SlicingCollection | None = self.resolve("Mesh")
|
|
attr_name: str = self.resolve("Attribute")
|
|
value: float = self.resolve("Value")
|
|
|
|
if not coll: return None
|
|
|
|
res_selection = SlicingCollection([], coll.name + "_selection")
|
|
res_inverted = SlicingCollection([], coll.name + "_inverted")
|
|
|
|
for so in coll.objects:
|
|
if attr_name not in so.attributes:
|
|
# If attribute missing, entire object goes to inverted
|
|
res_inverted.objects.append(so)
|
|
continue
|
|
|
|
attr = so.attributes[attr_name]
|
|
vals = attr["values"]
|
|
|
|
mask = np.zeros(len(so.mesh), dtype=bool)
|
|
if self.comparison == 'EQUAL': mask = vals == value
|
|
elif self.comparison == 'NOT_EQUAL': mask = vals != value
|
|
elif self.comparison == 'GREATER_THAN': mask = vals > value
|
|
elif self.comparison == 'LESS_THAN': mask = vals < value
|
|
elif self.comparison == 'GREATER_EQUAL': mask = vals >= value
|
|
elif self.comparison == 'LESS_EQUAL': mask = vals <= value
|
|
|
|
if np.all(mask):
|
|
res_selection.objects.append(so)
|
|
elif not np.any(mask):
|
|
res_inverted.objects.append(so)
|
|
else:
|
|
# Split the object
|
|
so_sel = SlicingObject()
|
|
so_sel.name = so.name + "_sel"
|
|
so_sel.parent = so.parent
|
|
so_sel.object_type = so.object_type
|
|
so_sel.extruder = so.extruder
|
|
so_sel.mesh = so.mesh[mask]
|
|
for k, v in so.attributes.items():
|
|
so_sel.attributes[k] = {
|
|
"domain": v["domain"],
|
|
"type": v["type"],
|
|
"values": v["values"][mask]
|
|
}
|
|
res_selection.objects.append(so_sel)
|
|
|
|
so_inv = SlicingObject()
|
|
so_inv.name = so.name + "_inv"
|
|
so_inv.parent = so.parent
|
|
so_inv.object_type = so.object_type
|
|
so_inv.extruder = so.extruder
|
|
so_inv.mesh = so.mesh[~mask]
|
|
for k, v in so.attributes.items():
|
|
so_inv.attributes[k] = {
|
|
"domain": v["domain"],
|
|
"type": v["type"],
|
|
"values": v["values"][~mask]
|
|
}
|
|
res_inverted.objects.append(so_inv)
|
|
|
|
if socket_name == "Selection":
|
|
return res_selection
|
|
return res_inverted
|