bake_roto.py

import nuke.rotopaint as rp, _curvelib as cl, nukescripts, nuke, math, time

class ShapePanel(nukescripts.PythonPanel):
	def __init__(self, node):
		nukescripts.PythonPanel.__init__(self, 'RotoPaint Elements')
		self.rp_node = node
		self.frame_range = nuke.String_Knob('frame_range', 'Frame Range', '{0}-{1}'.format(nuke.root().firstFrame(), nuke.root().lastFrame()))
		self.type_filter_knob = nuke.Enumeration_Knob('element', 'Element Type', ['Shapes', 'Strokes'])
		self.element_knob = nuke.Enumeration_Knob('curve', 'Curve', [])
		self.cv_knob = nuke.Enumeration_Knob('cv', 'CV', [])
		self.cv_knob.clearFlag( nuke.STARTLINE )
		for k in (self.frame_range, self.type_filter_knob, self.element_knob, self.cv_knob):
			self.addKnob(k)
		self.elements = {}
		self.get_elements()

	def parse_layer(self, layer, elements, parents):
		"""Return a list of all Shapes or Strokes in a Roto or Rotopaint node's heirarchy, with a list of all parent Layers for each item."""
		for e in layer:
			if isinstance(e, rp.Shape):
				elements['Shapes'].append([e, parents])
			elif isinstance(e, rp.Stroke):
				elements['Strokes'].append([e, parents])
			elif isinstance(e, rp.Layer):
				parents_copy = parents[:]
				parents_copy.insert(0, e)
				elements = self.parse_layer(e, elements, parents_copy)
		return elements

	def get_elements(self):
		"""	Get all Shape Stroke and Layer elements in our Roto/Paint node."""
		self.elements = { 'Shapes':[], 'Strokes':[] }
		self.curves_knob = self.rp_node['curves']
		self.root_layer = self.curves_knob.rootLayer
		self.elements = self.parse_layer(self.root_layer, self.elements, [self.root_layer])

	def knobChanged(self, knob):
		"""	Modify contents of element_knob and cv_knob: if Shapes are selected, element_knob shows shapes.
			If Strokes are selected, element_knob shows Strokes. 
			And Whatever element_knob is selected, the number of CVs is displayed in cv_knob"""
		if knob is self.type_filter_knob or knob.name() == 'showPanel':
			self.element_knob.setValues( [x[0].name for x in self.elements[self.type_filter_knob.value()]])
			# Update the element knob also
			#self.selected_element_and_parents
		if knob is self.element_knob or knob.name() == 'showPanel':
			# To get the name value of an enumKnob, you have to do .enumName(int(knob.getValue())): http://forums.thefoundry.co.uk/phpBB2/viewtopic.php?p=15856&sid=a17f123d8eee3f16ded415be6ac478f5
			selected_element_name = self.element_knob.enumName(int(self.element_knob.getValue()))
			# Find the shape or stroke object that matches the selected element name
			for e in self.elements[self.type_filter_knob.value()]:
				if e[0].name == selected_element_name:
					# Variable to access element and associated parent layer objects outside of class
					self.selected_element_and_parents = e
					parentnames = [p.name for p in e[1]]
					parentnames.reverse()
					# toElement() function does not include root in path
					parentnames.remove('Root')
					element_path = '/'.join(parentnames)
				else:
					continue
			# element_path_name is the path to the element, eg: 'Layer1/NestedLayer2/Bezier1'
			element_path_name = element_path+'/'+selected_element_name
			selected_element = self.curves_knob.toElement( element_path_name )
			# This list comprehension generates range of numbers matching CV points
			self.cv_knob.setValues([i for i, x in enumerate(selected_element)]) 

		if knob is self.cv_knob:
			self.selected_cv = int(self.cv_knob.enumName(int(self.cv_knob.getValue())))
			# Set frame_range knob to match first and last keyframes on selected CV if Shape (Strokes have no method to get point keyframes)
			if self.type_filter_knob.value() == 'Shapes':
				cv_keys = self.selected_element_and_parents[0][self.selected_cv].center.getControlPointKeyTimes()
				self.frame_range.setValue('{0}-{1}'.format(int(cv_keys[0]), int(cv_keys[-1])))
			else:
				self.frame_range.setValue('{0}-{1}'.format(nuke.root().firstFrame(), nuke.root().lastFrame()))
		#!TODO - Stroke support is broken, need to initialize cv on start so don't need to change it for it to work
		# Also if there is only 1 shape or stroke, the element will not get passed down


def TransformPoint(point, frame, item):
	"""Transform a point with a layer or object's transform properties, given a point and a frame. Compatible with > Nuke 6.2 only."""
	extramatrix = item.getTransform().evaluate(frame).getMatrix()
	vector = nuke.math.Vector4(point[0], point[1], 1, 1)
	x = (vector[0] * extramatrix[0]) + (vector[1] * extramatrix[1]) + extramatrix[2] + extramatrix[3]
	y = (vector[0] * extramatrix[4]) + (vector[1] * extramatrix[5]) + extramatrix[6] + extramatrix[7]
	z = (vector[0] * extramatrix[8]) + (vector[1] * extramatrix[9]) + extramatrix[10] + extramatrix[11]
	w = (vector[0] * extramatrix[12]) + (vector[1] * extramatrix[13]) + extramatrix[14] + extramatrix[15]
	new_point = nuke.math.Vector4(x, y, z, w) / w
	return new_point


def CVPointToTrack(rp_node, element_type, element, parents, cv, framerange):
	"""	Takes a point of an element and "un-parents" it from its parent transforms,
		then generates tracking data for that point.
	"""
	# Add knob to receive cv point track data, create another tab if it doesn't already exist
	if rp_node.knob('user_tracks') == None:
		rp_node.addKnob(nuke.Tab_Knob('user_tracks', 'Tracks'))
	track_knob = nuke.XY_Knob('{0}_cv{1}'.format(element.name, cv))
	rp_node.addKnob(track_knob)
	track_knob.setAnimated()
	track_knob_anim = [track_knob.animation(0), track_knob.animation(1)]

	# Create a list of nuke.AnimationKey objects for each unparented cv point value
	cv_anim = [[],[]]
	for frame in xrange(framerange[0], framerange[1]+1):
		#??? Would these be better/faster to do with getAnimationCurve?
		if element_type == 'Shapes':
			# Get parented position of point on shape
			cv_pos = [element[cv].center.getPosition(frame)[0], element[cv].center.getPosition(frame)[1]]
		elif element_type == 'Strokes':
			cv_pos = [element[cv].getPosition(frame)[0], element[cv].getPosition(frame)[1]]
		all_parents = parents[:]
		# Include self in parents, in case there are transforms on the object in question
		all_parents.insert(0, element)
		#print [x.name for x in all_parents]
		# Unparent the cv point from the transforms of all parents
		unparented_cv_pos = cv_pos
		for parent in all_parents:
			unparented_cv_pos = TransformPoint(unparented_cv_pos, frame, parent)
			#print parent.name, unparented_cv_pos
		for i in [0,1]:
			cv_anim[i].append(nuke.AnimationKey(frame, unparented_cv_pos[i]))
	#print [cv_anim]
	for i in [0,1]:
		track_knob_anim[i].addKey(cv_anim[i])

def init():
	rp_node = nuke.selectedNode()
	p = ShapePanel(rp_node)
	p.showModalDialog()
	element_type = p.type_filter_knob.value()
	element = p.selected_element_and_parents[0]
	parents = p.selected_element_and_parents[1]

	cv = p.selected_cv
	framerange = map(int, p.frame_range.value().split('-'))
	print "starting cgpoitntotrack function with", element_type, element.name, [x.name for x in parents], cv, framerange
	CVPointToTrack(rp_node, element_type, element, parents, cv, framerange)

if __name__=="__main__":
	init()