link_nodes.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893
  1. from .node_container_common import *
  2. from bpy.types import Bone, NodeTree
  3. from .base_definitions import MantisNode, GraphError, FLOAT_EPSILON
  4. from .link_socket_templates import *
  5. def TellClasses():
  6. return [
  7. # special
  8. LinkInherit,
  9. # copy
  10. LinkCopyLocation,
  11. LinkCopyRotation,
  12. LinkCopyScale,
  13. LinkCopyTransforms,
  14. LinkTransformation,
  15. # limit
  16. LinkLimitLocation,
  17. LinkLimitRotation,
  18. LinkLimitScale,
  19. LinkLimitDistance,
  20. # tracking
  21. LinkStretchTo,
  22. LinkDampedTrack,
  23. LinkLockedTrack,
  24. LinkTrackTo,
  25. #misc
  26. LinkInheritConstraint,
  27. LinkArmature,
  28. # IK
  29. LinkInverseKinematics,
  30. LinkSplineIK,
  31. # stuff that snaps or limits a bone
  32. LinkFloor,
  33. LinkShrinkWrap,
  34. # Drivers
  35. LinkDrivenParameter,
  36. ]
  37. # set the name if it is available, otherwise just use the constraint's nice name
  38. set_constraint_name = lambda nc : nc.evaluate_input("Name") if nc.evaluate_input("Name") else nc.__class__.__name__
  39. class MantisLinkNode(MantisNode):
  40. def __init__(self, signature : tuple,
  41. base_tree : NodeTree,
  42. socket_templates : list[SockTemplate]=[]):
  43. super().__init__(signature, base_tree, socket_templates)
  44. self.node_type = 'LINK'
  45. self.prepared = True; self.bObject=[]
  46. def evaluate_input(self, input_name, index=0):
  47. # should catch 'Target', 'Pole Target' and ArmatureConstraint targets, too
  48. if ('Target' in input_name) and input_name not in ["Target Space", "Use Target Z"]:
  49. socket = self.inputs.get(input_name)
  50. if socket.is_linked:
  51. return socket.links[0].from_node
  52. return None
  53. else:
  54. return super().evaluate_input(input_name)
  55. def gen_property_socket_map(self) -> dict:
  56. props_sockets = super().gen_property_socket_map()
  57. if (os := self.inputs.get("Owner Space")) and os.is_connected and os.links[0].from_node.node_type == 'XFORM':
  58. del props_sockets['owner_space']
  59. if (ts := self.inputs.get("Target Space")) and ts.is_connected and ts.links[0].from_node.node_type == 'XFORM':
  60. del props_sockets['target_space']
  61. return props_sockets
  62. def set_custom_space(self):
  63. for c in self.bObject:
  64. if (os := self.inputs.get("Owner Space")) and os.is_connected and os.links[0].from_node.node_type == 'XFORM':
  65. c.owner_space='CUSTOM'
  66. xf = self.inputs["Owner Space"].links[0].from_node.bGetObject(mode="OBJECT")
  67. if isinstance(xf, Bone):
  68. c.space_object=self.inputs["Owner Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  69. else:
  70. c.space_object=xf
  71. if ts := self.inputs.get("Target_Space") and ts.is_connected and ts.links[0].from_node.node_type == 'XFORM':
  72. c.target_space='CUSTOM'
  73. xf = self.inputs["Target_Space Space"].links[0].from_node.bGetObject(mode="OBJECT")
  74. if isinstance(xf, Bone):
  75. c.space_object=self.inputs["Target_Space Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  76. else:
  77. c.space_object=xf
  78. def GetxForm(nc, output_name="Output Relationship"):
  79. break_condition= lambda node : node.node_type=='XFORM'
  80. xforms = trace_line_up_branching(nc, output_name, break_condition)
  81. return_me=[]
  82. for xf in xforms:
  83. if xf.node_type != 'XFORM':
  84. continue
  85. if xf in return_me:
  86. continue
  87. return_me.append(xf)
  88. return return_me
  89. def reset_execution(self):
  90. super().reset_execution()
  91. self.prepared = True; self.bObject = []
  92. def bFinalize(self, bContext=None):
  93. finish_drivers(self)
  94. #*#-------------------------------#++#-------------------------------#*#
  95. # L I N K N O D E S
  96. #*#-------------------------------#++#-------------------------------#*#
  97. class LinkInherit(MantisLinkNode):
  98. '''A node representing inheritance'''
  99. def __init__(self, signature, base_tree):
  100. super().__init__(signature, base_tree, LinkInheritSockets)
  101. self.init_parameters()
  102. self.set_traverse([('Parent', 'Inheritance')])
  103. self.executed = True
  104. def GetxForm(self):
  105. # I think this is only run in display update.
  106. trace = trace_single_line_up(self, "Inheritance")
  107. for node in trace[0]:
  108. if (node.node_type == 'XFORM'):
  109. return node
  110. raise GraphError("%s is not connected to a downstream xForm" % self)
  111. class LinkCopyLocation(MantisLinkNode):
  112. '''A node representing Copy Location'''
  113. def __init__(self, signature : tuple,
  114. base_tree : NodeTree,):
  115. super().__init__(signature, base_tree, LinkCopyLocationSockets)
  116. additional_parameters = { "Name":None }
  117. self.init_parameters(additional_parameters=additional_parameters)
  118. self.set_traverse([("Input Relationship", "Output Relationship")])
  119. def bRelationshipPass(self, context):
  120. prepare_parameters(self)
  121. for xf in self.GetxForm():
  122. c = xf.bGetObject().constraints.new('COPY_LOCATION')
  123. self.get_target_and_subtarget(c)
  124. print(wrapGreen("Creating ")+wrapWhite("Copy Location")+
  125. wrapGreen(" Constraint for bone: ") +
  126. wrapOrange(xf.bGetObject().name))
  127. if constraint_name := self.evaluate_input("Name"):
  128. c.name = constraint_name
  129. self.bObject.append(c)
  130. self.set_custom_space()
  131. props_sockets = self.gen_property_socket_map()
  132. evaluate_sockets(self, c, props_sockets)
  133. self.executed = True
  134. class LinkCopyRotation(MantisLinkNode):
  135. '''A node representing Copy Rotation'''
  136. def __init__(self, signature, base_tree):
  137. super().__init__(signature, base_tree, LinkCopyRotationSockets)
  138. additional_parameters = { "Name":None }
  139. self.init_parameters(additional_parameters=additional_parameters)
  140. self.set_traverse([("Input Relationship", "Output Relationship")])
  141. def bRelationshipPass(self, context):
  142. prepare_parameters(self)
  143. for xf in self.GetxForm():
  144. c = xf.bGetObject().constraints.new('COPY_ROTATION')
  145. self.get_target_and_subtarget(c)
  146. print(wrapGreen("Creating ")+wrapWhite("Copy Rotation")+
  147. wrapGreen(" Constraint for bone: ") +
  148. wrapOrange(xf.bGetObject().name))
  149. rotation_order = self.evaluate_input("RotationOrder")
  150. if ((rotation_order == 'QUATERNION') or (rotation_order == 'AXIS_ANGLE')):
  151. c.euler_order = 'AUTO'
  152. else:
  153. try:
  154. c.euler_order = rotation_order
  155. except TypeError: # it's a driver or incorrect
  156. c.euler_order = 'AUTO'
  157. if constraint_name := self.evaluate_input("Name"):
  158. c.name = constraint_name
  159. self.bObject.append(c)
  160. self.set_custom_space()
  161. props_sockets = self.gen_property_socket_map()
  162. evaluate_sockets(self, c, props_sockets)
  163. self.executed = True
  164. class LinkCopyScale(MantisLinkNode):
  165. '''A node representing Copy Scale'''
  166. def __init__(self, signature, base_tree):
  167. super().__init__(signature, base_tree, LinkCopyScaleSockets)
  168. additional_parameters = { "Name":None }
  169. self.init_parameters(additional_parameters=additional_parameters)
  170. self.set_traverse([("Input Relationship", "Output Relationship")])
  171. def bRelationshipPass(self, context):
  172. prepare_parameters(self)
  173. for xf in self.GetxForm():
  174. c = xf.bGetObject().constraints.new('COPY_SCALE')
  175. self.get_target_and_subtarget(c)
  176. print(wrapGreen("Creating ")+wrapWhite("Copy Scale")+
  177. wrapGreen(" Constraint for bone: ") +
  178. wrapOrange(xf.bGetObject().name))
  179. if constraint_name := self.evaluate_input("Name"):
  180. c.name = constraint_name
  181. self.bObject.append(c)
  182. if self.inputs["Owner Space"].is_connected and self.inputs["Owner Space"].links[0].from_node.node_type == 'XFORM':
  183. c.owner_space='CUSTOM'
  184. xf = self.inputs["Owner Space"].links[0].from_node.bGetObject(mode="OBJECT")
  185. if isinstance(xf, Bone):
  186. c.space_object=self.inputs["Owner Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  187. else:
  188. c.space_object=xf
  189. if self.inputs["Target Space"].is_connected and self.inputs["Target Space"].links[0].from_node.node_type == 'XFORM':
  190. c.target_space='CUSTOM'
  191. xf = self.inputs["Target Space"].links[0].from_node.bGetObject(mode="OBJECT")
  192. if isinstance(xf, Bone):
  193. c.space_object=self.inputs["Owner Space"].links[0].from_node.bGetParentArmature(); c.space_subtarget=xf.name
  194. else:
  195. c.space_object=xf
  196. props_sockets = self.gen_property_socket_map()
  197. evaluate_sockets(self, c, props_sockets)
  198. self.executed = True
  199. class LinkCopyTransforms(MantisLinkNode):
  200. '''A node representing Copy Transfoms'''
  201. def __init__(self, signature, base_tree):
  202. super().__init__(signature, base_tree, LinkCopyTransformsSockets)
  203. additional_parameters = { "Name":None }
  204. self.init_parameters(additional_parameters=additional_parameters)
  205. self.set_traverse([("Input Relationship", "Output Relationship")])
  206. def bRelationshipPass(self, context):
  207. prepare_parameters(self)
  208. for xf in self.GetxForm():
  209. c = xf.bGetObject().constraints.new('COPY_TRANSFORMS')
  210. self.get_target_and_subtarget(c)
  211. print(wrapGreen("Creating ")+wrapWhite("Copy Transforms")+
  212. wrapGreen(" Constraint for bone: ") +
  213. wrapOrange(xf.bGetObject().name))
  214. if constraint_name := self.evaluate_input("Name"):
  215. c.name = constraint_name
  216. self.bObject.append(c)
  217. self.set_custom_space()
  218. props_sockets = self.gen_property_socket_map()
  219. evaluate_sockets(self, c, props_sockets)
  220. self.executed = True
  221. class LinkTransformation(MantisLinkNode):
  222. '''A node representing Copy Transfoms'''
  223. def __init__(self, signature, base_tree):
  224. super().__init__(signature, base_tree, LinkTransformationSockets)
  225. self.init_parameters(additional_parameters={"Name":None })
  226. self.set_traverse([("Input Relationship", "Output Relationship")])
  227. def ui_modify_socket(self, ui_socket, socket_name=None):
  228. from_suffix, to_suffix = '', ''
  229. if self.evaluate_input("Map From") == 'ROTATION': from_suffix='_rot'
  230. elif self.evaluate_input("Map From") == 'SCALE': from_suffix='_scale'
  231. if self.evaluate_input("Map To") == 'ROTATION': to_suffix='_rot'
  232. elif self.evaluate_input("Map To") == 'SCALE': to_suffix='_scale'
  233. if ('To' in ui_socket.name or 'From' in ui_socket.name) and (from_suffix or to_suffix):
  234. for s_temp in self.socket_templates:
  235. if s_temp.name == ui_socket.name: break
  236. if 'from' in s_temp.blender_property:
  237. socket_name=s_temp.blender_property+from_suffix
  238. else:
  239. socket_name=s_temp.blender_property+to_suffix
  240. return self.update_socket_value(socket_name, ui_socket.default_value)
  241. return super().ui_modify_socket(ui_socket, socket_name)
  242. def bRelationshipPass(self, context):
  243. prepare_parameters(self)
  244. for xf in self.GetxForm():
  245. c = xf.bGetObject().constraints.new('TRANSFORM')
  246. self.get_target_and_subtarget(c)
  247. print(wrapGreen("Creating ")+wrapWhite("Transformation")+
  248. wrapGreen(" Constraint for bone: ") +
  249. wrapOrange(xf.bGetObject().name))
  250. if constraint_name := self.evaluate_input("Name"):
  251. c.name = constraint_name
  252. self.bObject.append(c)
  253. self.set_custom_space()
  254. props_sockets = self.gen_property_socket_map()
  255. # we have to fix the blender-property for scale/rotation
  256. # because Blender stores these separately.
  257. # I do not care that this code is ugly.
  258. from_suffix, to_replace = '', ''
  259. if self.evaluate_input("Map From") == 'ROTATION':
  260. from_suffix='_rot'
  261. elif self.evaluate_input("Map From") == 'SCALE':
  262. from_suffix='_scale'
  263. if self.evaluate_input("Map To") == 'ROTATION':
  264. to_replace='_rot'
  265. elif self.evaluate_input("Map To") == 'SCALE':
  266. to_replace='_scale'
  267. if from_suffix:
  268. for axis in ['x', 'y', 'z']:
  269. stub='from_min_'+axis
  270. props_sockets[stub+from_suffix]=props_sockets[stub]
  271. del props_sockets[stub]
  272. stub='from_max_'+axis
  273. props_sockets[stub+from_suffix]=props_sockets[stub]
  274. del props_sockets[stub]
  275. if to_replace:
  276. for axis in ['x', 'y', 'z']:
  277. stub='to_min_'+axis
  278. props_sockets[stub+to_replace]=props_sockets[stub]
  279. del props_sockets[stub]
  280. stub='to_max_'+axis
  281. props_sockets[stub+to_replace]=props_sockets[stub]
  282. del props_sockets[stub]
  283. evaluate_sockets(self, c, props_sockets)
  284. self.executed = True
  285. class LinkLimitLocation(MantisLinkNode):
  286. def __init__(self, signature, base_tree):
  287. super().__init__(signature, base_tree, LinkLimitLocationScaleSockets)
  288. self.init_parameters(additional_parameters={ "Name":None })
  289. self.set_traverse([("Input Relationship", "Output Relationship")])
  290. def bRelationshipPass(self, context):
  291. prepare_parameters(self)
  292. for xf in self.GetxForm():
  293. c = xf.bGetObject().constraints.new('LIMIT_LOCATION')
  294. print(wrapGreen("Creating ")+wrapWhite("Limit Location")+
  295. wrapGreen(" Constraint for bone: ") +
  296. wrapOrange(xf.bGetObject().name))
  297. if constraint_name := self.evaluate_input("Name"):
  298. c.name = constraint_name
  299. self.bObject.append(c)
  300. self.set_custom_space()
  301. props_sockets = self.gen_property_socket_map()
  302. evaluate_sockets(self, c, props_sockets)
  303. self.executed = True
  304. class LinkLimitRotation(MantisLinkNode):
  305. def __init__(self, signature, base_tree):
  306. super().__init__(signature, base_tree, LinkLimitRotationSockets)
  307. self.init_parameters(additional_parameters={ "Name":None })
  308. self.set_traverse([("Input Relationship", "Output Relationship")])
  309. def bRelationshipPass(self, context):
  310. prepare_parameters(self)
  311. for xf in self.GetxForm():
  312. c = xf.bGetObject().constraints.new('LIMIT_ROTATION')
  313. print(wrapGreen("Creating ")+wrapWhite("Limit Rotation")+
  314. wrapGreen(" Constraint for bone: ") +
  315. wrapOrange(xf.bGetObject().name))
  316. if constraint_name := self.evaluate_input("Name"):
  317. c.name = constraint_name
  318. self.bObject.append(c)
  319. self.set_custom_space()
  320. props_sockets = self.gen_property_socket_map()
  321. evaluate_sockets(self, c, props_sockets)
  322. self.executed = True
  323. class LinkLimitScale(MantisLinkNode):
  324. def __init__(self, signature, base_tree):
  325. super().__init__(signature, base_tree, LinkLimitLocationScaleSockets)
  326. self.init_parameters(additional_parameters={ "Name":None })
  327. self.set_traverse([("Input Relationship", "Output Relationship")])
  328. def bRelationshipPass(self, context):
  329. prepare_parameters(self)
  330. for xf in self.GetxForm():
  331. c = xf.bGetObject().constraints.new('LIMIT_SCALE')
  332. print(wrapGreen("Creating ")+wrapWhite("Limit Scale")+
  333. wrapGreen(" Constraint for bone: ") +
  334. wrapOrange(xf.bGetObject().name))
  335. if constraint_name := self.evaluate_input("Name"):
  336. c.name = constraint_name
  337. self.bObject.append(c)
  338. self.set_custom_space()
  339. props_sockets = self.gen_property_socket_map()
  340. evaluate_sockets(self, c, props_sockets)
  341. self.executed = True
  342. class LinkLimitDistance(MantisLinkNode):
  343. def __init__(self, signature, base_tree):
  344. super().__init__(signature, base_tree, LinkLimitDistanceSockets)
  345. self.init_parameters(additional_parameters={ "Name":None })
  346. self.set_traverse([("Input Relationship", "Output Relationship")])
  347. def bRelationshipPass(self, context):
  348. prepare_parameters(self)
  349. for xf in self.GetxForm():
  350. print(wrapGreen("Creating ")+wrapWhite("Limit Distance")+
  351. wrapGreen(" Constraint for bone: ") +
  352. wrapOrange(xf.bGetObject().name))
  353. c = xf.bGetObject().constraints.new('LIMIT_DISTANCE')
  354. self.get_target_and_subtarget(c)
  355. if constraint_name := self.evaluate_input("Name"):
  356. c.name = constraint_name
  357. self.bObject.append(c)
  358. self.set_custom_space()
  359. props_sockets = self.gen_property_socket_map()
  360. evaluate_sockets(self, c, props_sockets)
  361. self.executed = True
  362. # Tracking
  363. class LinkStretchTo(MantisLinkNode):
  364. def __init__(self, signature, base_tree):
  365. super().__init__(signature, base_tree, LinkStretchToSockets)
  366. self.init_parameters(additional_parameters={ "Name":None })
  367. self.set_traverse([("Input Relationship", "Output Relationship")])
  368. def bRelationshipPass(self, context):
  369. prepare_parameters(self)
  370. for xf in self.GetxForm():
  371. print(wrapGreen("Creating ")+wrapWhite("Stretch-To")+
  372. wrapGreen(" Constraint for bone: ") +
  373. wrapOrange(xf.bGetObject().name))
  374. c = xf.bGetObject().constraints.new('STRETCH_TO')
  375. self.get_target_and_subtarget(c)
  376. if constraint_name := self.evaluate_input("Name"):
  377. c.name = constraint_name
  378. self.bObject.append(c)
  379. props_sockets = self.gen_property_socket_map()
  380. evaluate_sockets(self, c, props_sockets)
  381. if (self.evaluate_input("Original Length") == 0):
  382. # this is meant to be set automatically.
  383. c.rest_length = xf.bGetObject().bone.length
  384. self.executed = True
  385. class LinkDampedTrack(MantisLinkNode):
  386. def __init__(self, signature, base_tree):
  387. super().__init__(signature, base_tree, LinkDampedTrackSockets)
  388. self.init_parameters(additional_parameters={ "Name":None })
  389. self.set_traverse([("Input Relationship", "Output Relationship")])
  390. def bRelationshipPass(self, context):
  391. prepare_parameters(self)
  392. for xf in self.GetxForm():
  393. print(wrapGreen("Creating ")+wrapWhite("Damped Track")+
  394. wrapGreen(" Constraint for bone: ") +
  395. wrapOrange(xf.bGetObject().name))
  396. c = xf.bGetObject().constraints.new('DAMPED_TRACK')
  397. self.get_target_and_subtarget(c)
  398. if constraint_name := self.evaluate_input("Name"):
  399. c.name = constraint_name
  400. self.bObject.append(c)
  401. props_sockets = self.gen_property_socket_map()
  402. evaluate_sockets(self, c, props_sockets)
  403. self.executed = True
  404. class LinkLockedTrack(MantisLinkNode):
  405. def __init__(self, signature, base_tree):
  406. super().__init__(signature, base_tree,LinkLockedTrackSockets)
  407. self.init_parameters(additional_parameters={"Name":None })
  408. self.set_traverse([("Input Relationship", "Output Relationship")])
  409. def bRelationshipPass(self, context):
  410. prepare_parameters(self)
  411. for xf in self.GetxForm():
  412. print(wrapGreen("Creating ")+wrapWhite("Locked Track")+
  413. wrapGreen(" Constraint for bone: ") +
  414. wrapOrange(xf.bGetObject().name))
  415. c = xf.bGetObject().constraints.new('LOCKED_TRACK')
  416. self.get_target_and_subtarget(c)
  417. if constraint_name := self.evaluate_input("Name"):
  418. c.name = constraint_name
  419. self.bObject.append(c)
  420. props_sockets = self.gen_property_socket_map()
  421. evaluate_sockets(self, c, props_sockets)
  422. self.executed = True
  423. class LinkTrackTo(MantisLinkNode):
  424. def __init__(self, signature, base_tree):
  425. super().__init__(signature, base_tree, LinkTrackToSockets)
  426. self.init_parameters(additional_parameters={"Name":None })
  427. self.set_traverse([("Input Relationship", "Output Relationship")])
  428. def bRelationshipPass(self, context):
  429. prepare_parameters(self)
  430. for xf in self.GetxForm():
  431. print(wrapGreen("Creating ")+wrapWhite("Track-To")+
  432. wrapGreen(" Constraint for bone: ") +
  433. wrapOrange(xf.bGetObject().name))
  434. c = xf.bGetObject().constraints.new('TRACK_TO')
  435. self.get_target_and_subtarget(c)
  436. if constraint_name := self.evaluate_input("Name"):
  437. c.name = constraint_name
  438. self.bObject.append(c)
  439. props_sockets = self.gen_property_socket_map()
  440. evaluate_sockets(self, c, props_sockets)
  441. self.executed = True
  442. class LinkInheritConstraint(MantisLinkNode):
  443. def __init__(self, signature, base_tree):
  444. super().__init__(signature, base_tree, LinkInheritConstraintSockets)
  445. self.init_parameters(additional_parameters={"Name":None })
  446. self.set_traverse([("Input Relationship", "Output Relationship")])
  447. def bRelationshipPass(self, context):
  448. prepare_parameters(self)
  449. for xf in self.GetxForm():
  450. print(wrapGreen("Creating ")+wrapWhite("Child-Of")+
  451. wrapGreen(" Constraint for bone: ") +
  452. wrapOrange(xf.bGetObject().name))
  453. c = xf.bGetObject().constraints.new('CHILD_OF')
  454. self.get_target_and_subtarget(c)
  455. if constraint_name := self.evaluate_input("Name"):
  456. c.name = constraint_name
  457. self.bObject.append(c)
  458. props_sockets = self.gen_property_socket_map()
  459. evaluate_sockets(self, c, props_sockets)
  460. c.set_inverse_pending
  461. self.executed = True
  462. class LinkInverseKinematics(MantisLinkNode):
  463. def __init__(self, signature, base_tree):
  464. super().__init__(signature, base_tree, LinkInverseKinematicsSockets)
  465. self.init_parameters(additional_parameters={"Name":None })
  466. self.set_traverse([("Input Relationship", "Output Relationship")])
  467. def get_base_ik_bone(self, ik_bone):
  468. chain_length : int = (self.evaluate_input("Chain Length"))
  469. if not isinstance(chain_length, (int, float)):
  470. raise GraphError(f"Chain Length must be an integer number in {self}::Chain Length")
  471. if chain_length == 0:
  472. chain_length = int("inf")
  473. base_ik_bone = ik_bone; i=1
  474. while (i<chain_length) and (base_ik_bone.parent):
  475. base_ik_bone=base_ik_bone.parent; i+=1
  476. return base_ik_bone
  477. # We need to do the calculation in a "full circle", meaning the pole_angle
  478. # can go over pi or less than -pi - but the actuall constraint value must
  479. # be clamped in that range.
  480. # so we simply wrap the value.
  481. # not very efficient but it's OK
  482. def set_pole_angle(self, constraint, angle: float) -> None:
  483. from math import pi
  484. from .utilities import wrap
  485. constraint.pole_angle = wrap(-pi, pi, angle)
  486. def calc_pole_angle_pre(self, c, ik_bone):
  487. """
  488. This function gets us most of the way to a correct IK pole angle. Unfortunately,
  489. due to the unpredictable nature of the iterative IK calculation, I can't figure
  490. out an exact solution. So we do a bisect search in calc_pole_angle_post().
  491. """
  492. # TODO: instead of these checks, convert all to armature local space. But this is tedious.
  493. if not c.target:
  494. raise GraphError(f"IK Constraint {self} must have target.")
  495. elif c.target.type != "ARMATURE":
  496. raise NotImplementedError(f"Currently, IK Constraint Target for {self} must be a bone within the same armature.")
  497. if c.pole_target.type != "ARMATURE":
  498. raise NotImplementedError(f"Currently, IK Constraint Pole Target for {self} must be a bone within the same armature.")
  499. ik_handle = c.target.pose.bones[c.subtarget]
  500. if ik_handle.id_data != ik_bone.id_data:
  501. raise NotImplementedError(f"Currently, IK Constraint Target for {self} must be a bone within the same armature.")
  502. ik_pole = c.pole_target.pose.bones[c.pole_subtarget]
  503. if ik_pole.id_data != ik_bone.id_data:
  504. raise NotImplementedError(f"Currently,IK Constraint Pole Target for {self} must be a bone within the same armature.")
  505. base_ik_bone = self.get_base_ik_bone(ik_bone)
  506. start_effector = base_ik_bone.bone.head_local
  507. end_effector = ik_handle.bone.head_local
  508. pole_location = ik_pole.bone.head_local
  509. # this is the X-Axis of the bone's rest-pose, added to its bone
  510. knee_location = base_ik_bone.bone.matrix_local.col[0].xyz+start_effector
  511. ik_axis = (end_effector-start_effector).normalized()
  512. from .utilities import project_point_to_plane
  513. pole_planar_projection = project_point_to_plane(pole_location, start_effector, ik_axis)
  514. # this planar projection is necessary because the IK axis is different than the base_bone's y axis
  515. planar_projection = project_point_to_plane(knee_location, start_effector, ik_axis)
  516. knee_direction =(planar_projection - start_effector).normalized()
  517. pole_direction =(pole_planar_projection - start_effector).normalized()
  518. return knee_direction.angle(pole_direction)
  519. def calc_pole_angle_post(self, c, ik_bone, context):
  520. """
  521. This function should give us a completely accurate result for IK.
  522. """
  523. from time import time
  524. start_time=time()
  525. def signed_angle(vector_u, vector_v, normal):
  526. # it seems that this fails if the vectors are exactly aligned under certain circumstances.
  527. angle = vector_u.angle(vector_v, 0.0) # So we use a fallback of 0
  528. # Normal specifies orientation
  529. if angle != 0 and vector_u.cross(vector_v).angle(normal) < 1:
  530. angle = -angle
  531. return angle
  532. # we have already checked for valid data.
  533. ik_handle = c.target.pose.bones[c.subtarget]
  534. base_ik_bone = self.get_base_ik_bone(ik_bone)
  535. start_effector = base_ik_bone.bone.head_local
  536. angle = c.pole_angle
  537. dg = context.view_layer.depsgraph
  538. dg.update()
  539. ik_axis = (ik_handle.bone.head_local-start_effector).normalized()
  540. center_point = start_effector +(ik_axis*base_ik_bone.bone.length)
  541. knee_direction = base_ik_bone.bone.tail_local - center_point
  542. current_knee_direction = base_ik_bone.tail-center_point
  543. error=signed_angle(current_knee_direction, knee_direction, ik_axis)
  544. if error == 0:
  545. prGreen("No Fine-tuning needed."); return
  546. # Flip it if needed
  547. dot_before=current_knee_direction.dot(knee_direction)
  548. if dot_before < 0 and angle!=0: # then it is not aligned and we should check the inverse
  549. angle = -angle; c.pole_angle=angle
  550. dg.update()
  551. current_knee_direction = base_ik_bone.tail-center_point
  552. dot_after=current_knee_direction.dot(knee_direction)
  553. if dot_after < dot_before: # they are somehow less aligned
  554. prPurple("Mantis has gone down an unexpected code path. Please report this as a bug.")
  555. angle = -angle; self.set_pole_angle(c, angle)
  556. dg.update()
  557. # now we can do a bisect search to find the best value.
  558. error_threshhold = FLOAT_EPSILON
  559. max_iterations=600
  560. error=signed_angle(current_knee_direction, knee_direction, ik_axis)
  561. if error == 0:
  562. prGreen("No Fine-tuning needed."); return
  563. angle+=error
  564. alt_angle = angle+(error*-2) # should be very near the center when flipped here
  565. # we still need to bisect search because the relationship of pole_angle <==> error is somewhat unpredictable
  566. upper_bounds = alt_angle if alt_angle > angle else angle
  567. lower_bounds = alt_angle if alt_angle < angle else angle
  568. i, error_identical = 0, 0
  569. while ( True ):
  570. if (i>=max_iterations):
  571. prOrange(f"IK Pole Angle Set reached max iterations of {i-error_identical} in {time()-start_time} seconds")
  572. break
  573. if (abs(error)<error_threshhold) or (upper_bounds<=lower_bounds) or (error_identical > 3):
  574. prPurple(f"IK Pole Angle Set converged after {i-error_identical} iterations with error={error} in {time()-start_time} seconds")
  575. break
  576. # get the center-point betweeen the bounds
  577. try_angle = lower_bounds + (upper_bounds-lower_bounds)/2
  578. self.set_pole_angle(c, try_angle); dg.update()
  579. prev_error = error
  580. error = signed_angle((base_ik_bone.tail-center_point), knee_direction, ik_axis)
  581. error_identical+= int(error == prev_error)
  582. if error>0: upper_bounds=try_angle
  583. if error<0: lower_bounds=try_angle
  584. i+=1
  585. def bRelationshipPass(self, context):
  586. prepare_parameters(self)
  587. for xf in self.GetxForm():
  588. print(wrapGreen("Creating ")+wrapOrange("Inverse Kinematics")+
  589. wrapGreen(" Constraint for bone: ") +
  590. wrapOrange(xf.bGetObject().name))
  591. ik_bone = xf.bGetObject()
  592. c = xf.bGetObject().constraints.new('IK')
  593. self.get_target_and_subtarget(c)
  594. self.get_target_and_subtarget(c, input_name = 'Pole Target')
  595. if constraint_name := self.evaluate_input("Name"):
  596. c.name = constraint_name
  597. self.bObject.append(c)
  598. c.chain_count = 1 # so that, if there are errors, this doesn't print
  599. # a whole bunch of circular dependency crap from having infinite chain length
  600. if (c.pole_target):
  601. self.set_pole_angle(c, self.calc_pole_angle_pre(c, ik_bone))
  602. props_sockets = self.gen_property_socket_map()
  603. evaluate_sockets(self, c, props_sockets)
  604. c.use_location = self.evaluate_input("Position") > 0
  605. c.use_rotation = self.evaluate_input("Rotation") > 0
  606. self.executed = True
  607. def bFinalize(self, bContext = None):
  608. # adding a test here
  609. if bContext:
  610. for i, constraint in enumerate(self.bObject):
  611. ik_bone = self.GetxForm()[i].bGetObject(mode='POSE')
  612. if constraint.pole_target:
  613. prWhite(f"Fine-tuning IK Pole Angle for {self}")
  614. # make sure to enable it first
  615. enabled_before = constraint.mute
  616. constraint.mute = False
  617. self.calc_pole_angle_post(constraint, ik_bone, bContext)
  618. constraint.mute = enabled_before
  619. super().bFinalize(bContext)
  620. def ik_report_error(pb, context, do_print=False):
  621. dg = context.view_layer.depsgraph
  622. dg.update()
  623. loc1, rot_quaternion1, scl1 = pb.matrix.decompose()
  624. loc2, rot_quaternion2, scl2 = pb.bone.matrix_local.decompose()
  625. location_error=(loc1-loc2).length
  626. rotation_error = rot_quaternion1.rotation_difference(rot_quaternion2).angle
  627. scale_error = (scl1-scl2).length
  628. if location_error < FLOAT_EPSILON: location_error = 0
  629. if abs(rotation_error) < FLOAT_EPSILON: rotation_error = 0
  630. if scale_error < FLOAT_EPSILON: scale_error = 0
  631. if do_print:
  632. print (f"IK Location Error: {location_error}")
  633. print (f"IK Rotation Error: {rotation_error}")
  634. print (f"IK Scale Error : {scale_error}")
  635. return (location_error, rotation_error, scale_error)
  636. # This is kinda a weird design decision?
  637. class LinkDrivenParameter(MantisLinkNode):
  638. '''A node representing an armature object'''
  639. def __init__(self, signature, base_tree):
  640. super().__init__(signature, base_tree, LinkDrivenParameterSockets)
  641. self.init_parameters(additional_parameters={ "Name":None })
  642. self.set_traverse([("Input Relationship", "Output Relationship")])
  643. def bRelationshipPass(self, bContext = None,):
  644. prepare_parameters(self)
  645. prGreen("Executing Driven Parameter node")
  646. prop = self.evaluate_input("Parameter")
  647. index = self.evaluate_input("Index")
  648. value = self.evaluate_input("Value")
  649. for xf in self.GetxForm():
  650. ob = xf.bGetObject(mode="POSE")
  651. # IMPORTANT: this node only works on pose bone attributes.
  652. self.bObject.append(ob)
  653. length=1
  654. if hasattr(ob, prop):
  655. try:
  656. length = len(getattr(ob, prop))
  657. except TypeError:
  658. pass
  659. except AttributeError:
  660. pass
  661. else:
  662. raise AttributeError(f"Cannot Set value {prop} on object because it does not exist.")
  663. def_value = 0.0
  664. if length>1:
  665. def_value=[0.0]*length
  666. self.parameters["Value"] = tuple( 0.0 if i != index else value for i in range(length))
  667. props_sockets = {
  668. prop: ("Value", def_value)
  669. }
  670. evaluate_sockets(self, ob, props_sockets)
  671. self.executed = True
  672. def bFinalize(self, bContext = None):
  673. driver = self.evaluate_input("Value")
  674. try:
  675. for i, val in enumerate(self.parameters["Value"]):
  676. from .drivers import MantisDriver
  677. if isinstance(val, MantisDriver):
  678. driver["ind"] = i
  679. val = driver
  680. except AttributeError:
  681. self.parameters["Value"] = driver
  682. except TypeError:
  683. self.parameters["Value"] = driver
  684. super().bFinalize(bContext)
  685. class LinkArmature(MantisLinkNode):
  686. '''A node representing an armature object'''
  687. def __init__(self, signature, base_tree,):
  688. super().__init__(signature, base_tree, LinkArmatureSockets)
  689. self.init_parameters(additional_parameters={"Name":None })
  690. self.set_traverse([("Input Relationship", "Output Relationship")])
  691. setup_custom_props(self) # <-- this takes care of the runtime-added sockets
  692. def bRelationshipPass(self, bContext = None,):
  693. prepare_parameters(self)
  694. for xf in self.GetxForm():
  695. print(wrapGreen("Creating ")+wrapOrange("Armature")+
  696. wrapGreen(" Constraint for bone: ") +
  697. wrapOrange(xf.bGetObject().name))
  698. c = xf.bGetObject().constraints.new('ARMATURE')
  699. if constraint_name := self.evaluate_input("Name"):
  700. c.name = constraint_name
  701. self.bObject.append(c)
  702. # get number of targets
  703. num_targets = len( list(self.inputs.values())[6:] )//2
  704. props_sockets = self.gen_property_socket_map()
  705. targets_weights = {}
  706. for i in range(num_targets):
  707. target = c.targets.new()
  708. target_input_name = list(self.inputs.keys())[i*2+6 ]
  709. weight_input_name = list(self.inputs.keys())[i*2+6+1]
  710. self.get_target_and_subtarget(target, target_input_name)
  711. weight_value=self.evaluate_input(weight_input_name)
  712. if not isinstance(weight_value, float):
  713. weight_value=0
  714. targets_weights[i]=weight_value
  715. props_sockets["targets[%d].weight" % i] = (weight_input_name, 0)
  716. # targets_weights.append({"weight":(weight_input_name, 0)})
  717. evaluate_sockets(self, c, props_sockets)
  718. for target, value in targets_weights.items():
  719. c.targets[target].weight=value
  720. self.executed = True
  721. class LinkSplineIK(MantisLinkNode):
  722. '''A node representing an armature object'''
  723. def __init__(self, signature, base_tree):
  724. super().__init__(signature, base_tree, LinkSplineIKSockets)
  725. self.init_parameters(additional_parameters={"Name":None })
  726. self.set_traverse([("Input Relationship", "Output Relationship")])
  727. def bRelationshipPass(self, bContext = None,):
  728. prepare_parameters(self)
  729. if not self.inputs['Target'].is_linked:
  730. raise GraphError(f"ERROR: {self} is not connected to a target curve.")
  731. for xf in self.GetxForm():
  732. print(wrapGreen("Creating ")+wrapOrange("Spline-IK")+
  733. wrapGreen(" Constraint for bone: ") +
  734. wrapOrange(xf.bGetObject().name))
  735. c = xf.bGetObject().constraints.new('SPLINE_IK')
  736. # set the spline - we need to get the right one
  737. spline_index = self.evaluate_input("Spline Index")
  738. from .utilities import get_extracted_spline_object
  739. proto_curve = self.inputs['Target'].links[0].from_node.bGetObject()
  740. curve = get_extracted_spline_object(proto_curve, spline_index, self.mContext)
  741. # link it to the view layer
  742. if (curve.name not in bContext.view_layer.active_layer_collection.collection.objects):
  743. bContext.view_layer.active_layer_collection.collection.objects.link(curve)
  744. c.target=curve
  745. if constraint_name := self.evaluate_input("Name"):
  746. c.name = constraint_name
  747. self.bObject.append(c)
  748. props_sockets = self.gen_property_socket_map()
  749. evaluate_sockets(self, c, props_sockets)
  750. self.executed = True
  751. class LinkFloor(MantisLinkNode):
  752. '''A node representing an armature object'''
  753. def __init__(self, signature, base_tree,):
  754. super().__init__(signature, base_tree, LinkFloorSockets)
  755. self.init_parameters(additional_parameters={"Name":None })
  756. self.set_traverse([("Input Relationship", "Output Relationship")])
  757. def bRelationshipPass(self, bContext = None,):
  758. prepare_parameters(self)
  759. for xf in self.GetxForm():
  760. print(wrapGreen("Creating ")+wrapOrange("Floor")+
  761. wrapGreen(" Constraint for bone: ") +
  762. wrapOrange(xf.bGetObject().name))
  763. c = xf.bGetObject().constraints.new('FLOOR')
  764. self.get_target_and_subtarget(c)
  765. if constraint_name := self.evaluate_input("Name"):
  766. c.name = constraint_name
  767. self.bObject.append(c)
  768. self.set_custom_space()
  769. props_sockets = self.gen_property_socket_map()
  770. evaluate_sockets(self, c, props_sockets)
  771. self.executed = True
  772. class LinkShrinkWrap(MantisLinkNode):
  773. '''A node representing a shrinkwrap relationship.'''
  774. def __init__(self, signature, base_tree,):
  775. super().__init__(signature, base_tree, LinkShrinkWrapSockets)
  776. self.init_parameters(additional_parameters={"Name":None })
  777. self.set_traverse([("Input Relationship", "Output Relationship")])
  778. def bRelationshipPass(self, bContext = None,):
  779. prepare_parameters(self)
  780. for xf in self.GetxForm():
  781. print(wrapGreen("Creating ")+wrapOrange("Shrinkwrap")+
  782. wrapGreen(" Constraint for bone: ") +
  783. wrapOrange(xf.bGetObject().name))
  784. c = xf.bGetObject().constraints.new('SHRINKWRAP')
  785. self.get_target_and_subtarget(c)
  786. if constraint_name := self.evaluate_input("Name"):
  787. c.name = constraint_name
  788. props_sockets = self.gen_property_socket_map()
  789. print (props_sockets.keys())
  790. self.bObject.append(c)
  791. # self.set_custom_space() # this needs to be overridden for me TODO
  792. evaluate_sockets(self, c, props_sockets)
  793. self.executed = True