diff --git a/molecularnodes/blender/attribute.py b/molecularnodes/blender/attribute.py index e32fb05f..a0b0991e 100644 --- a/molecularnodes/blender/attribute.py +++ b/molecularnodes/blender/attribute.py @@ -1,11 +1,9 @@ from dataclasses import dataclass from enum import Enum from typing import Type -from functools import reduce import bpy import numpy as np -from numpy.ma.core import prod @dataclass @@ -121,6 +119,9 @@ def guess_atype_from_array(array: np.ndarray) -> AttributeType: return AttributeTypes.FLOAT.value class Attribute: + """ + Wrapper around a Blender attribute to provide a more convenient interface with numpy arrays + """ def __init__(self, attribute: bpy.types.Attribute): self.attribute = attribute self.n_attr = len(attribute.data) @@ -158,6 +159,17 @@ def dtype(self) -> Type: def n_values(self) -> int: return np.prod(self.shape, dtype=int) + def from_array(self, array: np.ndarray) -> None: + """ + Set the attribute data from a numpy array + """ + if array.shape != self.shape: + raise ValueError( + f"Array shape {array.shape} does not match attribute shape {self.shape}" + ) + + self.attribute.data.foreach_set(self.value_name, array.reshape(-1)) + def as_array(self) -> np.ndarray: """ diff --git a/molecularnodes/blender/mesh.py b/molecularnodes/blender/mesh.py index 5fa369a8..1e6a7fde 100644 --- a/molecularnodes/blender/mesh.py +++ b/molecularnodes/blender/mesh.py @@ -10,7 +10,7 @@ AttributeTypes, guess_atype_from_array, ) -from .object import ObjectTracker, create_object +from .object import ObjectTracker, create_object, evaluate_object def centre(array: np.ndarray): @@ -96,7 +96,7 @@ def named_attribute( obj: bpy.types.Object, name="position", evaluate=False ) -> np.ndarray: """ - Get the attribute data from the object. + Get the named attribute data from the object, optionally evaluating modifiers first. Parameters: object (bpy.types.Object): The Blender object. @@ -107,20 +107,17 @@ def named_attribute( """ if evaluate: obj = evaluate_object(obj) - attribute_names = obj.data.attributes.keys() verbose = False - if name not in attribute_names: + try: + attr = Attribute(obj.data.attributes[name]) + except KeyError: + message = f"The selected attribute '{name}' does not exist on the mesh." if verbose: - raise AttributeError( - f"The selected attribute '{name}' does not exist on the mesh. \ - Possible attributes are: {attribute_names=}" - ) - else: - raise AttributeError( - f"The selected attribute '{name}' does not exist on the mesh." - ) + message += f"Possible attributes are: {obj.data.attributes.keys()}" + + raise AttributeError(message) - return Attribute(obj.data.attributes[name]).as_array() + return attr.as_array() def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object: @@ -153,10 +150,6 @@ def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types. return obj -def evaluate_object(obj): - "Return an object which has the modifiers evaluated." - obj.update_tag() - return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) def evaluate_using_mesh(obj): diff --git a/molecularnodes/blender/object.py b/molecularnodes/blender/object.py index b4737254..568261e1 100644 --- a/molecularnodes/blender/object.py +++ b/molecularnodes/blender/object.py @@ -63,6 +63,11 @@ def latest(self): return self.new_objects()[-1] +def evaluate_object(obj): + "Return an object which has the modifiers evaluated." + obj.update_tag() + return obj.evaluated_get(bpy.context.evaluated_depsgraph_get()) + def create_object( vertices: np.ndarray = [], edges: np.ndarray = [],