Robowflex  v0.1
Making MoveIt Easy
utils.py
Go to the documentation of this file.
1 '''Helper functions for the Robowflex Visualization library.
2 '''
3 
4 import bpy
5 import mathutils
6 
7 import math
8 import os.path
9 import subprocess
10 import logging
11 import random
12 import yaml
13 
14 
15 def resolve_package(path):
16  '''Resolves `package://` URLs to their canonical form. The path does not need
17  to exist, but the package does. Can be used to write new files in packages.
18 
19  Returns "" on failure.
20  '''
21  if not path:
22  return ""
23 
24  package_name = ""
25  package_path1 = ""
26  PREFIX = "package://"
27  if PREFIX in path:
28  path = path[len(PREFIX):] # Remove "package://"
29  if "/" not in path:
30  package_name = path
31  path = ""
32  else:
33  package_name = path[:path.find("/")]
34  path = path[path.find("/"):]
35 
36  package_path1 = subprocess.check_output(
37  ["rospack", "find", package_name]).decode().strip()
38 
39  elif "~" in path:
40  path = os.path.expanduser(path)
41 
42  new_path = os.path.realpath(package_path1 + path)
43  return new_path
44 
45 
46 def resolve_path(path):
47  '''Resolves `package://` URLs and relative file paths to their canonical form.
48  Returns "" on failure.
49  '''
50  full_path = resolve_package(path)
51  if not os.path.exists(full_path):
52  logging.warn("File {} does not exist".format(full_path))
53  return ""
54  return full_path
55 
56 
58  '''Selects and returns all children, recursively.'''
59  items = [item]
60  item.select_set(True)
61  for child in item.children:
62  items += select_all_children(child)
63 
64  return items
65 
66 
68  '''Applies smooth shading to the provided object.
69  '''
70  if bpy.context.mode != 'OBJECT':
71  bpy.ops.object.mode_set(mode = 'OBJECT')
72 
73  deselect_all()
74  set_active(item)
75  item.select_set(True)
76 
77  if item.type == 'MESH':
78  bpy.ops.object.shade_smooth()
79 
80  deselect_all()
81 
82 
83 def apply_edge_split(item, angle = math.pi / 8):
84  '''Applies the edge-split modifier to the provided object.
85  '''
86  if bpy.context.mode != 'OBJECT':
87  bpy.ops.object.mode_set(mode = 'OBJECT')
88 
89  deselect_all()
90  set_active(item)
91  item.select_set(True)
92  try:
93  bpy.ops.object.modifier_add(type = 'EDGE_SPLIT')
94  bpy.context.object.modifiers["EdgeSplit"].split_angle = angle
95  except:
96  pass
97 
98  deselect_all()
99 
100 
101 def find_object_in_collection(coll_name, item_name):
102  '''Retrieves an object by name from a collection.
103  '''
104  collection = bpy.data.collections.get(coll_name)
105  for item in collection.objects:
106  if item.name == item_name:
107  return item
108 
109  return None
110 
111 
113  '''Moves selected objects to collection.
114  '''
115  collection = bpy.data.collections.get(name)
116  selected = bpy.context.selected_objects
117  for select in selected:
118  old = find_collection(select)
119 
120  collection.objects.link(select)
121  old.objects.unlink(select)
122 
123 
124 def find_collection(item):
125  '''Find the current collection of an object.
126  '''
127  collections = item.users_collection
128  if len(collections) > 0:
129  return collections[0]
130  return bpy.context.scene.collection
131 
132 
133 def get_collection(name):
134  '''Get a collection.
135  '''
136  collection = bpy.data.collections.get(name)
137  return collection
138 
139 
140 def remove_collection(name, remove_objects = True):
141  '''Removes a collection and all its contents.
142  '''
143  collection = bpy.data.collections.get(name)
144 
145  if collection:
146  if remove_objects:
147  obs = [o for o in collection.objects if o.users == 1]
148  while obs:
149  bpy.data.objects.remove(obs.pop())
150 
151  bpy.data.collections.remove(collection)
152 
153 
154 def make_collection(name):
155  '''Creates a new collection, deleting collections with the same name.
156  '''
157 
158  remove_collection(name)
159 
160  collection = bpy.data.collections.new(name)
161  bpy.context.scene.collection.children.link(collection)
162  return collection
163 
164 
166  bpy.ops.object.select_all(action = 'DESELECT')
167 
168 
169 # def parent_object(parent, child):
170 # deselect_all()
171 # parent.select_set(True)
172 # child.select_set(True)
173 # set_active(child)
174 
175 # bpy.ops.object.parent_set(type = 'OBJECT')
176 # deselect_all()
177 
178 
179 def create_object_parent(parent, child):
180  if bpy.context.mode != 'OBJECT':
181  bpy.ops.object.mode_set(mode = 'OBJECT')
182 
183  deselect_all()
184  set_active(child)
185  child.select_set(True)
186 
187  constraint = get_object_parent(parent, child)
188  if not constraint:
189  try:
190  bpy.ops.object.constraint_add(type = 'CHILD_OF')
191  constraint = bpy.context.object.constraints["Child Of"]
192 
193  constraint.name = "{}_to_{}".format(parent.name, child.name)
194  constraint.target = parent
195  constraint.influence = 0.
196  except:
197  pass
198 
199  deselect_all()
200  return constraint
201 
202 
203 def get_object_parent(parent, child):
204  if bpy.context.mode != 'OBJECT':
205  bpy.ops.object.mode_set(mode = 'OBJECT')
206 
207  deselect_all()
208  set_active(child)
209  child.select_set(True)
210 
211  try:
212  name = "{}_to_{}".format(parent.name, child.name)
213  constraint = bpy.context.object.constraints[name]
214  return constraint
215  except:
216  pass
217 
218  return None
219 
220 
221 def parent_object(parent, child, frame):
222  # Set current frame so inverse matrix can properly be computed
223  bpy.context.scene.frame_set(frame)
224  constraint = create_object_parent(parent, child)
225 
226  # Set object as parented in frame
227  constraint.influence = 1
228  constraint.keyframe_insert(data_path = "influence", frame = frame)
229 
230  # Compute inverse transform for proper parenting
231  deselect_all()
232  set_active(child)
233  child.select_set(True)
234 
235  context_py = bpy.context.copy()
236  context_py["constraint"] = constraint
237 
238  bpy.ops.constraint.childof_set_inverse(context_py,
239  constraint = constraint.name,
240  owner = "OBJECT")
241 
242  deselect_all()
243 
244  # Set object as unparented in prior frame
245  constraint.influence = 0
246  constraint.keyframe_insert(data_path = "influence", frame = frame - 1)
247 
248 
249 def unparent_object(parent, child, frame):
250  bpy.context.scene.frame_set(frame - 1)
251  constraint = get_object_parent(parent, child)
252  if not constraint:
253  return
254 
255  # Set object as parented in prior frame
256  constraint.influence = 1
257  constraint.keyframe_insert(data_path = "influence", frame = frame - 1)
258 
259  # Set object as unparented in prior frame
260  constraint.influence = 0
261  constraint.keyframe_insert(data_path = "influence", frame = frame)
262 
263 
264 def set_active(item):
265  bpy.context.view_layer.objects.active = item
266 
267 
268 def add_material(item, material):
269  '''Adds a material to an object.
270  '''
271  if item.data.materials:
272  # assign to 1st material slot
273  item.data.materials[0] = material
274 
275  else:
276  # no slots
277  item.data.materials.append(material)
278 
279 
280 def set_color(obj, element):
281  if 'color' in element:
282  # TODO: figure out a better way to make new materials?
283  mat = bpy.data.materials.new(name = str(random.randint(1, 100000)))
284  if len(element['color']) > 3:
285  mat.diffuse_color = element['color']
286  else:
287  mat.diffuse_color = element['color'] + (1., )
288 
289  add_material(obj, mat)
290 
291 
292 def pose_to_quat(pose):
293  '''Takes a pose dict and extracts the orientation quaternion. ROS quaternions or XYZW, but Blender's are WXYZ, so
294  reorder them.
295  '''
296 
297  if 'x' in pose['orientation']:
298  q = pose['orientation']
299  return mathutils.Quaternion([q['w'], q['x'], q['y'], q['z']])
300 
301  if isinstance(pose['orientation'][0], str):
302  # Means there's a NAN somewhere.
303  return mathutils.Quaternion((1.0, 0.0, 0.0, 0.0))
304  else:
305  return mathutils.Quaternion(pose['orientation'][3:] +
306  pose['orientation'][:3])
307 
308 
309 def pose_to_vec(pose):
310  '''Takes a pose dict and extracts the position vector.
311  '''
312  if 'x' in pose['position']:
313  v = pose['position']
314  return mathutils.Vector([v['x'], v['y'], v['z']])
315 
316  return mathutils.Vector(pose['position'])
317 
318 
319 def pose_add(obj, pose1, pose2):
320  '''Adds the second pose after the first. For something like link origins, the
321  joint pose should be passed first, then the link origin
322  '''
323 
324  # Get the objects for the pose vectors and quaternions.
325  p1_vec = pose_to_vec(pose1)
326  q1 = pose_to_quat(pose1)
327  p2_vec = pose_to_vec(pose2)
328  q2 = pose_to_quat(pose2)
329 
330  # Rotate pose2's vec by pose1's quaternion.
331  p2_vec.rotate(q1)
332  obj.location = p1_vec + p2_vec
333 
334  # Apply pose2's quaternion to the existing rotation.
335  q2.rotate(q1)
336  obj.rotation_mode = 'QUATERNION'
337  obj.rotation_quaternion = q2
338 
339 
340 def set_pose(obj, pose):
341  '''Sets the pose of a blender object by passing in a pose dict.
342  '''
343 
344  obj.location = pose_to_vec(pose)
345  obj.rotation_mode = 'QUATERNION'
346  obj.rotation_quaternion = pose_to_quat(pose)
347 
348 
349 def read_YAML_data(file_name):
350  '''Returns the yaml data structure of the data stored.
351  '''
352  full_name = resolve_path(file_name)
353  if not full_name:
354  logging.warn('Cannot open {}'.format(file_name))
355  return None
356  with open(full_name) as input_file:
357  return yaml.load(input_file.read(), Loader=yaml.SafeLoader)
358 
359 
360 def remove_doubles(item, threshold = 0.0001):
361  if bpy.context.mode != 'OBJECT':
362  bpy.ops.object.mode_set(mode = 'OBJECT')
363 
364  deselect_all()
365  set_active(item)
366  item.select_set(True)
367  if item.type == 'MESH':
368  bpy.ops.object.mode_set(mode = 'EDIT')
369  bpy.ops.mesh.select_all(action = 'SELECT')
370  bpy.ops.mesh.remove_doubles(threshold = threshold)
371  bpy.ops.object.mode_set(mode = 'OBJECT')
372 
373  deselect_all()
374 
375 
377  if bpy.context.mode != 'OBJECT':
378  bpy.ops.object.mode_set(mode = 'OBJECT')
379 
380  deselect_all()
381  set_active(item)
382  item.select_set(True)
383  if item.type == 'MESH':
384  bpy.ops.object.mode_set(mode = 'EDIT')
385  bpy.ops.mesh.select_all(action = 'SELECT')
386  bpy.ops.mesh.select_mode(type = 'FACE')
387  bpy.ops.mesh.select_interior_faces()
388  bpy.ops.mesh.delete(type = 'FACE')
389  bpy.ops.object.mode_set(mode = 'OBJECT')
390 
391  deselect_all()
392 
393 
394 def clear_alpha(obj):
395  for mat in obj.data.materials:
396  if not mat:
397  continue
398  for node in mat.node_tree.nodes:
399  if 'Alpha' in node.inputs:
400  node.inputs['Alpha'].default_value = 1.
401 
402 
404  xyz = [0, 0, 0]
405  rpy = [0, 0, 0]
406 
407  if xml.origin:
408  rpy = xml.origin.rpy
409  xyz = xml.origin.xyz
410 
411  rot = mathutils.Euler(rpy, 'XYZ').to_matrix()
412  rot.resize_4x4()
413  pos = mathutils.Matrix.Translation(xyz)
414  return pos @ rot
std::string format(const std::string &fmt, Args &&... args)
Recursion base case, return string form of formatted arguments.
Definition: log.h:60
def remove_inner_faces(item)
Definition: utils.py:376
def find_object_in_collection(coll_name, item_name)
Definition: utils.py:101
def create_object_parent(parent, child)
Definition: utils.py:179
def set_pose(obj, pose)
Definition: utils.py:340
def set_color(obj, element)
Definition: utils.py:280
def move_selected_to_collection(name)
Definition: utils.py:112
def get_object_parent(parent, child)
Definition: utils.py:203
def remove_doubles(item, threshold=0.0001)
Definition: utils.py:360
def unparent_object(parent, child, frame)
Definition: utils.py:249
def apply_smooth_shade(item)
Definition: utils.py:67
def pose_add(obj, pose1, pose2)
Definition: utils.py:319
def read_YAML_data(file_name)
Definition: utils.py:349
def remove_collection(name, remove_objects=True)
Definition: utils.py:140
def parent_object(parent, child, frame)
Definition: utils.py:221
def add_material(item, material)
Definition: utils.py:268
def apply_edge_split(item, angle=math.pi/8)
Definition: utils.py:83
def select_all_children(item)
Definition: utils.py:57