Node Wrangler: do not add reroutes to unavailable outputs
[blender-addons.git] / object_scatter / operator.py
blob51736f693ec12e1a2ac2daa77325f16ef9947724
1 # SPDX-FileCopyrightText: 2018-2023 Blender Foundation
3 # SPDX-License-Identifier: GPL-2.0-or-later
5 import bpy
6 import gpu
7 import blf
8 import math
9 import enum
10 import random
12 from itertools import islice
13 from mathutils.bvhtree import BVHTree
14 from mathutils import Vector, Matrix, Euler
15 from gpu_extras.batch import batch_for_shader
17 from bpy_extras.view3d_utils import (
18 region_2d_to_vector_3d,
19 region_2d_to_origin_3d
23 # Modal Operator
24 ################################################################
26 class ScatterObjects(bpy.types.Operator):
27 bl_idname = "object.scatter"
28 bl_label = "Scatter Objects"
29 bl_options = {'REGISTER', 'UNDO'}
31 @classmethod
32 def poll(cls, context):
33 return (
34 currently_in_3d_view(context)
35 and context.active_object is not None
36 and context.active_object.mode == 'OBJECT')
38 def invoke(self, context, event):
39 self.target_object = context.active_object
40 self.objects_to_scatter = get_selected_non_active_objects(context)
42 if self.target_object is None or len(self.objects_to_scatter) == 0:
43 self.report({'ERROR'}, "Select objects to scatter and a target object")
44 return {'CANCELLED'}
46 self.base_scale = get_max_object_side_length(self.objects_to_scatter)
48 self.targets = []
49 self.active_target = None
50 self.target_cache = {}
52 self.enable_draw_callback()
53 context.window_manager.modal_handler_add(self)
54 return {'RUNNING_MODAL'}
56 def modal(self, context, event):
57 context.area.tag_redraw()
59 if not event_is_in_region(event, context.region) and self.active_target is None:
60 return {'PASS_THROUGH'}
62 if event.type == 'ESC':
63 return self.finish('CANCELLED')
65 if event.type == 'RET' and event.value == 'PRESS':
66 self.create_scatter_object()
67 return self.finish('FINISHED')
69 event_used = self.handle_non_exit_event(event)
70 if event_used:
71 return {'RUNNING_MODAL'}
72 else:
73 return {'PASS_THROUGH'}
75 def handle_non_exit_event(self, event):
76 if self.active_target is None:
77 if event.type == 'LEFTMOUSE' and event.value == 'PRESS':
78 self.active_target = StrokeTarget()
79 self.active_target.start_build(self.target_object)
80 return True
81 else:
82 build_state = self.active_target.continue_build(event)
83 if build_state == BuildState.FINISHED:
84 self.targets.append(self.active_target)
85 self.active_target = None
86 self.remove_target_from_cache(self.active_target)
87 return True
89 return False
91 def enable_draw_callback(self):
92 self._draw_callback_view = bpy.types.SpaceView3D.draw_handler_add(self.draw_view, (), 'WINDOW', 'POST_VIEW')
93 self._draw_callback_px = bpy.types.SpaceView3D.draw_handler_add(self.draw_px, (), 'WINDOW', 'POST_PIXEL')
95 def disable_draw_callback(self):
96 bpy.types.SpaceView3D.draw_handler_remove(self._draw_callback_view, 'WINDOW')
97 bpy.types.SpaceView3D.draw_handler_remove(self._draw_callback_px, 'WINDOW')
99 def draw_view(self):
100 for target in self.iter_targets():
101 target.draw()
103 draw_matrices_batches(list(self.iter_matrix_batches()))
105 def draw_px(self):
106 draw_text((20, 20, 0), "Instances: " + str(len(self.get_all_matrices())))
108 def finish(self, return_value):
109 self.disable_draw_callback()
110 bpy.context.area.tag_redraw()
111 return {return_value}
113 def create_scatter_object(self):
114 matrix_chunks = make_random_chunks(
115 self.get_all_matrices(), len(self.objects_to_scatter))
117 collection = bpy.data.collections.new("Scatter")
118 bpy.context.collection.children.link(collection)
120 for obj, matrices in zip(self.objects_to_scatter, matrix_chunks):
121 make_duplicator(collection, obj, matrices)
123 def get_all_matrices(self):
124 settings = self.get_current_settings()
126 matrices = []
127 for target in self.iter_targets():
128 self.ensure_target_is_in_cache(target)
129 matrices.extend(self.target_cache[target].get_matrices(settings))
130 return matrices
132 def iter_matrix_batches(self):
133 settings = self.get_current_settings()
134 for target in self.iter_targets():
135 self.ensure_target_is_in_cache(target)
136 yield self.target_cache[target].get_batch(settings)
138 def iter_targets(self):
139 yield from self.targets
140 if self.active_target is not None:
141 yield self.active_target
143 def ensure_target_is_in_cache(self, target):
144 if target not in self.target_cache:
145 entry = TargetCacheEntry(target, self.base_scale)
146 self.target_cache[target] = entry
148 def remove_target_from_cache(self, target):
149 self.target_cache.pop(self.active_target, None)
151 def get_current_settings(self):
152 return bpy.context.scene.scatter_properties.to_settings()
154 class TargetCacheEntry:
155 def __init__(self, target, base_scale):
156 self.target = target
157 self.last_used_settings = None
158 self.base_scale = base_scale
159 self.settings_changed()
161 def get_matrices(self, settings):
162 self._handle_new_settings(settings)
163 if self.matrices is None:
164 self.matrices = self.target.get_matrices(settings)
165 return self.matrices
167 def get_batch(self, settings):
168 self._handle_new_settings(settings)
169 if self.gpu_batch is None:
170 self.gpu_batch = create_batch_for_matrices(self.get_matrices(settings), self.base_scale)
171 return self.gpu_batch
173 def _handle_new_settings(self, settings):
174 if settings != self.last_used_settings:
175 self.settings_changed()
176 self.last_used_settings = settings
178 def settings_changed(self):
179 self.matrices = None
180 self.gpu_batch = None
183 # Duplicator Creation
184 ######################################################
186 def make_duplicator(target_collection, source_object, matrices):
187 triangle_scale = 0.1
189 duplicator = triangle_object_from_matrices(source_object.name + " Duplicator", matrices, triangle_scale)
190 duplicator.instance_type = 'FACES'
191 duplicator.use_instance_faces_scale = True
192 duplicator.show_instancer_for_viewport = True
193 duplicator.show_instancer_for_render = False
194 duplicator.instance_faces_scale = 1 / triangle_scale
196 copy_obj = source_object.copy()
197 copy_obj.name = source_object.name + " - copy"
198 copy_obj.location = (0, 0, 0)
199 copy_obj.parent = duplicator
201 target_collection.objects.link(duplicator)
202 target_collection.objects.link(copy_obj)
204 def triangle_object_from_matrices(name, matrices, triangle_scale):
205 mesh = triangle_mesh_from_matrices(name, matrices, triangle_scale)
206 return bpy.data.objects.new(name, mesh)
208 def triangle_mesh_from_matrices(name, matrices, triangle_scale):
209 mesh = bpy.data.meshes.new(name)
210 vertices, polygons = mesh_data_from_matrices(matrices, triangle_scale)
211 mesh.from_pydata(vertices, [], polygons)
212 mesh.update()
213 mesh.validate()
214 return mesh
216 unit_triangle_vertices = (
217 Vector((-3**-0.25, -3**-0.75, 0)),
218 Vector((3**-0.25, -3**-0.75, 0)),
219 Vector((0, 2/3**0.75, 0)))
221 def mesh_data_from_matrices(matrices, triangle_scale):
222 vertices = []
223 polygons = []
224 triangle_vertices = [triangle_scale * v for v in unit_triangle_vertices]
226 for i, matrix in enumerate(matrices):
227 vertices.extend((matrix @ v for v in triangle_vertices))
228 polygons.append((i * 3 + 0, i * 3 + 1, i * 3 + 2))
230 return vertices, polygons
233 # Target Provider
234 #################################################
236 class BuildState(enum.Enum):
237 FINISHED = enum.auto()
238 ONGOING = enum.auto()
240 class TargetProvider:
241 def start_build(self, target_object):
242 pass
244 def continue_build(self, event):
245 return BuildState.FINISHED
247 def get_matrices(self, scatter_settings):
248 return []
250 def draw(self):
251 pass
253 class StrokeTarget(TargetProvider):
254 def start_build(self, target_object):
255 self.points = []
256 self.bvhtree = bvhtree_from_object(target_object)
257 self.batch = None
259 def continue_build(self, event):
260 if event.type == 'LEFTMOUSE' and event.value == 'RELEASE':
261 return BuildState.FINISHED
263 mouse_pos = (event.mouse_region_x, event.mouse_region_y)
264 location, *_ = shoot_region_2d_ray(self.bvhtree, mouse_pos)
265 if location is not None:
266 self.points.append(location)
267 self.batch = None
268 return BuildState.ONGOING
270 def draw(self):
271 if self.batch is None:
272 self.batch = create_line_strip_batch(self.points)
273 draw_line_strip_batch(self.batch, color=(1.0, 0.4, 0.1, 1.0), thickness=5)
275 def get_matrices(self, scatter_settings):
276 return scatter_around_stroke(self.points, self.bvhtree, scatter_settings)
278 def scatter_around_stroke(stroke_points, bvhtree, settings):
279 scattered_matrices = []
280 for point, local_seed in iter_points_on_stroke_with_seed(stroke_points, settings.density, settings.seed):
281 matrix = scatter_from_source_point(bvhtree, point, local_seed, settings)
282 scattered_matrices.append(matrix)
283 return scattered_matrices
285 def iter_points_on_stroke_with_seed(stroke_points, density, seed):
286 for i, (start, end) in enumerate(iter_pairwise(stroke_points)):
287 segment_seed = sub_seed(seed, i)
288 segment_vector = end - start
290 segment_length = segment_vector.length
291 amount = round_random(segment_length * density, segment_seed)
293 for j in range(amount):
294 t = random_uniform(sub_seed(segment_seed, j, 0))
295 origin = start + t * segment_vector
296 yield origin, sub_seed(segment_seed, j, 1)
298 def scatter_from_source_point(bvhtree, point, seed, settings):
299 # Project displaced point on surface
300 radius = random_uniform(sub_seed(seed, 0)) * settings.radius
301 offset = random_vector(sub_seed(seed, 2)) * radius
302 location, normal, *_ = bvhtree.find_nearest(point + offset)
303 assert location is not None
304 normal.normalize()
306 up_direction = normal if settings.use_normal_rotation else Vector((0, 0, 1))
308 # Scale
309 min_scale = settings.scale * (1 - settings.random_scale)
310 max_scale = settings.scale
311 scale = random_uniform(sub_seed(seed, 1), min_scale, max_scale)
313 # Location
314 location += normal * settings.normal_offset * scale
316 # Rotation
317 z_rotation = Euler((0, 0, random_uniform(sub_seed(seed, 3), 0, 2 * math.pi))).to_matrix()
318 up_rotation = up_direction.to_track_quat('Z', 'X').to_matrix()
319 local_rotation = random_euler(sub_seed(seed, 3), settings.rotation).to_matrix()
320 rotation = local_rotation @ up_rotation @ z_rotation
322 return Matrix.Translation(location) @ rotation.to_4x4() @ scale_matrix(scale)
325 # Drawing
326 #################################################
328 box_vertices = (
329 (-1, -1, 1), ( 1, -1, 1), ( 1, 1, 1), (-1, 1, 1),
330 (-1, -1, -1), ( 1, -1, -1), ( 1, 1, -1), (-1, 1, -1))
332 box_indices = (
333 (0, 1, 2), (2, 3, 0), (1, 5, 6), (6, 2, 1),
334 (7, 6, 5), (5, 4, 7), (4, 0, 3), (3, 7, 4),
335 (4, 5, 1), (1, 0, 4), (3, 2, 6), (6, 7, 3))
337 box_vertices = tuple(Vector(vertex) * 0.5 for vertex in box_vertices)
339 def draw_matrices_batches(batches):
340 shader = get_uniform_color_shader()
341 shader.bind()
342 shader.uniform_float("color", (0.4, 0.4, 1.0, 0.3))
344 gpu.state.blend_set('ALPHA')
345 gpu.state.depth_mask_set(False)
347 for batch in batches:
348 batch.draw(shader)
350 gpu.state.blend_set('NONE')
351 gpu.state.depth_mask_set(True)
353 def create_batch_for_matrices(matrices, base_scale):
354 coords = []
355 indices = []
357 scaled_box_vertices = [base_scale * vertex for vertex in box_vertices]
359 for matrix in matrices:
360 offset = len(coords)
361 coords.extend((matrix @ vertex for vertex in scaled_box_vertices))
362 indices.extend(tuple(index + offset for index in element) for element in box_indices)
364 batch = batch_for_shader(get_uniform_color_shader(),
365 'TRIS', {"pos" : coords}, indices = indices)
366 return batch
369 def draw_line_strip_batch(batch, color, thickness=1):
370 shader = get_uniform_color_shader()
371 gpu.state.line_width_set(thickness)
372 shader.bind()
373 shader.uniform_float("color", color)
374 batch.draw(shader)
376 def create_line_strip_batch(coords):
377 return batch_for_shader(get_uniform_color_shader(), 'LINE_STRIP', {"pos" : coords})
380 def draw_text(location, text, size=15, color=(1, 1, 1, 1)):
381 font_id = 0
382 ui_scale = bpy.context.preferences.system.ui_scale
383 blf.position(font_id, *location)
384 blf.size(font_id, round(size * ui_scale))
385 blf.draw(font_id, text)
388 # Utilities
389 ########################################################
392 Pythons random functions are designed to be used in cases
393 when a seed is set once and then many random numbers are
394 generated. To improve the user experience I want to have
395 full control over how random seeds propagate through the
396 functions. This is why I use custom random functions.
398 One benefit is that changing the object density does not
399 generate new random positions for all objects.
402 def round_random(value, seed):
403 probability = value % 1
404 if probability < random_uniform(seed):
405 return math.floor(value)
406 else:
407 return math.ceil(value)
409 def random_vector(x, min=-1, max=1):
410 return Vector((
411 random_uniform(sub_seed(x, 0), min, max),
412 random_uniform(sub_seed(x, 1), min, max),
413 random_uniform(sub_seed(x, 2), min, max)))
415 def random_euler(x, factor):
416 return Euler(tuple(random_vector(x) * factor))
418 def random_uniform(x, min=0, max=1):
419 return random_int(x) / 2147483648 * (max - min) + min
421 def random_int(x):
422 x = (x<<13) ^ x
423 return (x * (x * x * 15731 + 789221) + 1376312589) & 0x7fffffff
425 def sub_seed(seed, index, index2=0):
426 return random_int(seed * 3243 + index * 5643 + index2 * 54243)
429 def currently_in_3d_view(context):
430 return context.space_data.type == 'VIEW_3D'
432 def get_selected_non_active_objects(context):
433 return set(context.selected_objects) - {context.active_object}
435 def make_random_chunks(sequence, chunk_amount):
436 sequence = list(sequence)
437 random.shuffle(sequence)
438 return make_chunks(sequence, chunk_amount)
440 def make_chunks(sequence, chunk_amount):
441 length = math.ceil(len(sequence) / chunk_amount)
442 return [sequence[i:i+length] for i in range(0, len(sequence), length)]
444 def iter_pairwise(sequence):
445 return zip(sequence, islice(sequence, 1, None))
447 def bvhtree_from_object(object):
448 import bmesh
449 bm = bmesh.new()
451 depsgraph = bpy.context.evaluated_depsgraph_get()
452 object_eval = object.evaluated_get(depsgraph)
453 mesh = object_eval.to_mesh()
454 bm.from_mesh(mesh)
455 bm.transform(object.matrix_world)
457 bvhtree = BVHTree.FromBMesh(bm)
458 object_eval.to_mesh_clear()
459 return bvhtree
461 def shoot_region_2d_ray(bvhtree, position_2d):
462 region = bpy.context.region
463 region_3d = bpy.context.space_data.region_3d
465 origin = region_2d_to_origin_3d(region, region_3d, position_2d)
466 direction = region_2d_to_vector_3d(region, region_3d, position_2d)
468 location, normal, index, distance = bvhtree.ray_cast(origin, direction)
469 return location, normal, index, distance
471 def scale_matrix(factor):
472 m = Matrix.Identity(4)
473 m[0][0] = factor
474 m[1][1] = factor
475 m[2][2] = factor
476 return m
478 def event_is_in_region(event, region):
479 return (region.x <= event.mouse_x <= region.x + region.width
480 and region.y <= event.mouse_y <= region.y + region.height)
482 def get_max_object_side_length(objects):
483 return max(
484 max(obj.dimensions[0] for obj in objects),
485 max(obj.dimensions[1] for obj in objects),
486 max(obj.dimensions[2] for obj in objects)
489 def get_uniform_color_shader():
490 return gpu.shader.from_builtin('UNIFORM_COLOR')
493 # Registration
494 ###############################################
496 def register():
497 bpy.utils.register_class(ScatterObjects)
499 def unregister():
500 bpy.utils.unregister_class(ScatterObjects)