Source code for nukescripts.snap3d

# Copyright (c) 2010 The Foundry Visionmongers Ltd.  All Rights Reserved.

import math
import nuke_internal as nuke
import _nukemath

#
# Predefined snapping functions.
#

# Lazy functions to call on "thisNode"
[docs]def translateThisNodeToPoints(): return translateToPoints(nuke.thisNode())
[docs]def translateRotateThisNodeToPoints(): return translateRotateToPoints(nuke.thisNode())
[docs]def translateRotateScaleThisNodeToPoints(): return translateRotateScaleToPoints(nuke.thisNode())
# Lazy functions to determine the vertex selection # These are the external user functions # Much hard work is done in obtaining the selection in these functions
[docs]def translateToPoints(nodeToSnap): ''' Translate the specified node to the average position of the current vertex selection in the active viewer. The nodeToSnap must contain a 'translate' knob and the transform order must be 'SRT'. @type nodeToSnap: nuke.Node @param nodeToSnap: Node to translate ''' return translateSelectionToPoints(nodeToSnap, getSelection())
[docs]def translateRotateToPoints(nodeToSnap): ''' Translate the specified node to the average position of the current vertex selection in the active viewer and rotate to the orientation of the (mean squares) best fit plane for the selection. The nodeToSnap must contain 'translate' and 'rotate' knobs, the transform order must be 'SRT' and the rotation order must be 'ZXY'. @type nodeToSnap: nuke.Node @param nodeToSnap: Node to translate and rotate ''' return translateRotateSelectionToPoints(nodeToSnap, getSelection())
[docs]def translateRotateScaleToPoints(nodeToSnap): ''' Translate the specified node to the average position of the current vertex selection in the active viewer, rotate to the orientation of the (mean squares) best fit plane for the selection and scale to the extents of the selection. The nodeToSnap must contain 'translate', 'rotate' and 'scale' knobs, the transform order must be 'SRT' and the rotation order must be 'ZXY'. @type nodeToSnap: nuke.Node @param nodeToSnap: Node to translate, rotate and scale ''' return translateRotateScaleSelectionToPoints(nodeToSnap, getSelection())
# Verification wrappers
[docs]def translateSelectionToPoints(nodeToSnap, vertexSelection): try: verifyNodeToSnap(nodeToSnap, ["translate", "xform_order", "rot_order"]) verifyVertexSelection(vertexSelection, 1) translateToPointsVerified(nodeToSnap, vertexSelection) except ValueError as e: nuke.message(str(e))
[docs]def translateRotateSelectionToPoints(nodeToSnap, vertexSelection): try: verifyNodeToSnap(nodeToSnap, ["translate", "rotate", "xform_order", "rot_order"]) verifyVertexSelection(vertexSelection, 1) # This can use the normal direction translateRotateToPointsVerified(nodeToSnap, vertexSelection) except ValueError as e: nuke.message(str(e))
[docs]def translateRotateScaleSelectionToPoints(nodeToSnap, vertexSelection): try: verifyNodeToSnap(nodeToSnap, ["translate", "rotate", "scaling", "xform_order", "rot_order"]) verifyVertexSelection(vertexSelection, 3) translateRotateScaleToPointsVerified(nodeToSnap, vertexSelection) except ValueError as e: nuke.message(str(e))
[docs]def verifyVertexSelection(vertexSelection, minLen): if len(vertexSelection) < minLen: raise ValueError('Selection must be at least %d points' % minLen)
# Verification functions
[docs]def verifyNodeToSnap(nodeToSnap, knobList): # Check the knobs nodeKnobs = nodeToSnap.knobs() for knob in knobList: if knob not in nodeKnobs: raise ValueError('Snap requires "%s" knob' % knob) # Verify the transform order verifyNodeOrder(nodeToSnap, "xform_order", "SRT") # Verify the rotation order as necessary if "rotate" in knobList: verifyNodeOrder(nodeToSnap, "rot_order", "ZXY")
[docs]def verifyNodeOrder(node, knobName, orderName): orderKnob = node.knob(knobName) # Is there a better way than this? order = orderKnob.enumName( int(orderKnob.getValue()) ) if orderName != order: raise ValueError('Snap requires "%s" %s' % (orderName, knobName))
# Info for each vertex
[docs]class VertexInfo: def __init__(self, objnum, index, value, position): self.objnum = objnum self.index = index self.value = value self.position = position # This is updated on applying a transform
# Selection container
[docs]class VertexSelection: def __init__(self): self.vertexInfoSet = set() self.length = 0 def __len__(self): return self.length # Convenience function to allow direct iteration def __iter__(self): # Since a set is not iterable use a generator function # Don't change the set while iterating! for info in self.vertexInfoSet: yield info def add(self, vertexInfo): self.vertexInfoSet.add(vertexInfo) self.length += 1 def points(self): # Generate an iterable list of the positions points = [] for info in self.vertexInfoSet: points += [info.position] return points def indices(self): # Generate a searchable dictionary of the positions indices = {} i = 0 for info in self.vertexInfoSet: indices[info.index] = i i += 1 return indices def translate(self, vector): for info in self.vertexInfoSet: info.position += vector def inverseRotate(self, vector, order): #nuke.tprint(vector) #nuke.tprint(order) # Create a rotation matrix m = _nukemath.Matrix3() m.makeIdentity() for axis in order: if axis == "X": # Apply X rotation m.rotateX(vector[0]) #nuke.tprint("rotateX: %f" % vector[0]) elif axis == "Y": # Apply Y rotation m.rotateY(vector[1]) #nuke.tprint("rotateY: %f" % vector[1]) elif axis == "Z": # Apply Z rotation m.rotateZ(vector[2]) #nuke.tprint("rotateZ: %f" % vector[2]) else: raise ValueError("Invalid rotation axis: %s" % axis) # Now determine the inverse/transpose matrix #nuke.tprint(m) transpose(m) #nuke.tprint(m) # Apply the matrix to the vertices for info in self.vertexInfoSet: info.position = m * info.position def scale(self, vector): for info in self.vertexInfoSet: info.position[0] *= vector[0] info.position[1] *= vector[1] info.position[2] *= vector[2]
# Helper function since Matrix3 has no transpose operation
[docs]def transpose(m): t = m[0+3*1] m[0+3*1] = m[1+3*0] m[1+3*0] = t t = m[0+3*2] m[0+3*2] = m[2+3*0] m[2+3*0] = t t = m[1+3*2] m[1+3*2] = m[2+3*1] m[2+3*1] = t return
# Core snapping functions
[docs]def translateToPointsVerified(nodeToSnap, vertexSelection): # Find the average position centre = calcAveragePosition(vertexSelection) # Move the nodeToSnap to the average position nodeToSnap['translate'].setValue(centre) # Subtract this translation from the vertexSelection inverseTranslation = -centre vertexSelection.translate(inverseTranslation)
[docs]def scaleToPointsVerified(nodeToScale, vertexSelection): # Scale the nodeToScale to fit the bounding box of the selected points scale = calcBounds(vertexSelection) nodeToScale['scaling'].setValue(scale) # Apply the inverse scale to the points inverseScale = _nukemath.Vector3(1/scale[0], 1/scale[1], 1/scale[2]) vertexSelection.scale(inverseScale)
[docs]def rotateToPointsVerified(nodeToSnap, vertexSelection): # Get the normal of the points norm = averageNormal(vertexSelection) # Calculate the rotation vector rotationVec = calcRotationVector(vertexSelection, norm) # Convert to degrees rotationDegrees = _nukemath.Vector3( math.degrees(rotationVec.x), math.degrees(rotationVec.y), math.degrees(rotationVec.z) ) # Set the node transform nodeToSnap['rotate'].setValue(rotationDegrees) # Apply the reverse rotation to the points vertexSelection.inverseRotate(rotationVec, "YXZ")
[docs]def translateRotateToPointsVerified(nodeToSnap, vertexSelection): # Note that the vertexSelection positions are updated as we go translateToPointsVerified(nodeToSnap, vertexSelection) rotateToPointsVerified(nodeToSnap, vertexSelection)
[docs]def translateRotateScaleToPointsVerified(nodeToSnap, vertexSelection): # Note that the vertexSelection positions are updated as we go translateRotateToPointsVerified(nodeToSnap, vertexSelection) scaleToPointsVerified(nodeToSnap, vertexSelection)
# Get the rotation vector
[docs]def calcRotationVector(vertexSelection, norm): # Collate a point set from the vertex selection points = vertexSelection.points() # Find a best fit plane with three or more points if len(vertexSelection) >= 3: planeTri = nuke.geo.bestFitPlane(*points) rotationVec = planeRotation(planeTri, norm) elif len(vertexSelection) == 2: # Choose the axes dependent on the line direction u = points[1] - points[0] u.normalize() w = norm v = w.cross(u) v.normalize() # Update w w = v.cross(u) # Fabricate a tri (tuple) planeTri = (_nukemath.Vector3(0.0, 0.0, 0.0), u, v) rotationVec = planeRotation(planeTri, norm) elif len(vertexSelection) == 1: # Choose the axes dependent on the normal direction w = norm w.normalize() if abs(w.x) < abs(w.y): v = _nukemath.Vector3(0.0, 1.0, 0.0) u = w.cross(v) u.normalize() # Update v v = w.cross(u) else: u = _nukemath.Vector3(1.0, 0.0, 0.0) v = w.cross(u) v.normalize() # Update v u = w.cross(v) # Fabricate a tri (tuple) planeTri = (_nukemath.Vector3(0.0, 0.0, 0.0), u, v) rotationVec = planeRotation(planeTri, norm) # In fact this only handles ZXY (see planeRotation) rotationVec.z = 0 return rotationVec
# # Geometric functions #
[docs]def calcAveragePosition(vertexSelection): ''' Calculate the average position of all points. @param points: An iterable sequence of _nukemath.Vector3 objects. @return: A _nukemath.Vector3 containing the average of all the points. ''' pos = _nukemath.Vector3(0, 0, 0) count = 0 for info in vertexSelection: point = info.position pos += point count += 1 if count == 0: return pos /= count return pos
[docs]def calcBounds(vertexSelection): # Get the bounding box of all the selected points # Avoid zero size to allow inverse scaling (1/scale) high = _nukemath.Vector3(1e-20, 1e-20, 1e-20) low = _nukemath.Vector3(-1e-20, -1e-20, -1e-20) for info in vertexSelection: pos = info.position for i in range(len(pos)): if pos[i] < low[i]: low[i] = pos[i] elif pos[i] > high[i]: high[i] = pos[i] bounds = high - low return bounds
[docs]def planeRotation(tri, norm=None): ''' Calculate the rotations around the X, Y and Z axes that will align a plane perpendicular to the Z axis with the given triangle. @param tri: A list or tuple of 3 _nukemath.Vector3 objects. The 3 points must describe the plane (i.e. they must not be collinear). @return: A _nukemath.Vector3 object where the x coordinate is the angle of rotation around the x axis and so on. @raise ValueError: if the three points are collinear. ''' # Get the vectors along two edges of the triangle described by the three points. a = tri[1] - tri[0] b = tri[2] - tri[0] # Calculate the basis vectors for an orthonormal basis, where: # - u is parallel to a # - v is perpendicular to u and lies in the plane defined by a and b # - w is perpendicular to both u and v u = _nukemath.Vector3(a) u.normalize() w = a.cross(b) w.normalize() v = w.cross(u) # If a normal was passed in, check to make sure that the one we're generating # is aligned close to the same way if norm != None: if w.dot(norm) < 0.0: w.x = -w.x w.y = -w.y w.z = -w.z # Don't mirror! v.x = -v.x v.y = -v.y v.z = -v.z # Now find the rotation angles necessary to align a card to the uv plane. m = ( (u[0], v[0], w[0]), (u[1], v[1], w[1]), (u[2], v[2], w[2]) ) if abs(m[1][2]) == 1.0: ry = 0.0 rx = (math.pi / 2.0) * -m[1][2] rz = math.atan2(-m[0][1], m[0][0]) else: cosx = math.sqrt(m[0][2] ** 2 + m[2][2] ** 2) if cosx == 0: cosx = 1.0 rx = math.atan2(-m[1][2], cosx) rz = math.atan2(m[1][0] / cosx, m[1][1] / cosx) ry = math.atan2(m[0][2] / cosx, m[2][2] / cosx) return _nukemath.Vector3(rx, ry, rz)
# # Helper functions #
[docs]def averageNormal(vertexSelection): ''' averageNormal(selectionThreshold -> _nukemath.Vector3 Return a _nukemath.Vector3 which is the average of the normals of all selected points ''' if not nuke.activeViewer(): return None # Put the indices in a dictionary for fast searching fastIndices = vertexSelection.indices() found = False norm = _nukemath.Vector3(0.0, 0.0, 0.0) for theNode in allNodesWithGeoSelectKnob(): geoSelectKnob = theNode['geo_select'] sel = geoSelectKnob.getSelection() objs = geoSelectKnob.getGeometry() for o in range(len(sel)): objPrimitives = objs[o].primitives() objTransform = objs[o].transform() # Use a dictionary for fast searching visitedPrimitives = {} for prim in objPrimitives: # This will be slow! if prim not in visitedPrimitives: for pt in prim.points(): # This will be slow! if pt in fastIndices: found = True n = prim.normal() n = _nukemath.Vector3(n[0], n[1], n[2]) n = objTransform.vtransform(n) norm += n visitedPrimitives[prim] = pt break if found == False: return None norm.normalize() return norm
[docs]def anySelectedPoint(selectionThreshold=0.5): ''' anySelectedPoint(selectionThreshold) -> _nukemath.Vector3 Return a selected point from the active viewer or the first viewer with a selection. The selectionThreshold parameter is used when working with a soft selection. Only points with a selection level >= the selection threshold will be returned by this function. ''' if not nuke.activeViewer(): return None for n in allNodesWithGeoSelectKnob(): geoSelectKnob = n['geo_select'] sel = geoSelectKnob.getSelection() objs = geoSelectKnob.getGeometry() for o in range(len(sel)): objSelection = sel[o] for p in range(len(objSelection)): if objSelection[p] >= selectionThreshold: pos = objs[o].points()[p] tPos = objs[o].transform() * _nukemath.Vector4(pos.x, pos.y, pos.z, 1.0) return _nukemath.Vector3(tPos.x, tPos.y, tPos.z) return None
[docs]def selectedPoints(selectionThreshold=0.5): ''' selectedPoints(selectionThreshold) -> iterator Return an iterator which yields the position of every point currently selected in the Viewer in turn. The selectionThreshold parameter is used when working with a soft selection. Only points with a selection level >= the selection threshold will be returned by this function. ''' if not nuke.activeViewer(): return for info in selectedVertexInfos(selectionThreshold): yield info.position
[docs]def getSelection(selectionThreshold=0.5): # Build a VertexSelection from VertexInfos vertexSelection = VertexSelection() for info in selectedVertexInfos(selectionThreshold): vertexSelection.add(info) return vertexSelection
[docs]def selectedVertexInfos(selectionThreshold=0.5): ''' selectedVertexInfos(selectionThreshold) -> iterator Return an iterator which yields a tuple of the index and position of each point currently selected in the Viewer in turn. The selectionThreshold parameter is used when working with a soft selection. Only points with a selection level >= the selection threshold will be returned by this function. ''' if not nuke.activeViewer(): return for n in allNodesWithGeoSelectKnob(): geoSelectKnob = n['geo_select'] sel = geoSelectKnob.getSelection() objs = geoSelectKnob.getGeometry() for o in range(len(sel)): objSelection = sel[o] objPoints = objs[o].points() objTransform = objs[o].transform() for p in range(len(objSelection)): value = objSelection[p] if value >= selectionThreshold: pos = objPoints[p] tPos = objTransform * _nukemath.Vector4(pos.x, pos.y, pos.z, 1.0) yield VertexInfo(o, p, value, _nukemath.Vector3(tPos.x, tPos.y, tPos.z))
[docs]def anySelectedVertexInfo(selectionThreshold=0.5): ''' anySelectedVertexInfo(selectionThreshold) -> VertexInfo Returns a single VertexInfo for a selected point. If more than one point is selected, one of them will be chosen arbitrarily. The selectionThreshold parameter is used when working with a soft selection. Only points with a selection level >= the selection threshold will be returned by this function. ''' if not nuke.activeViewer(): return None for n in allNodesWithGeoSelectKnob(): geoSelectKnob = n['geo_select'] sel = geoSelectKnob.getSelection() objs = geoSelectKnob.getGeometry() for o in range(len(sel)): objSelection = sel[o] objPoints = objs[o].points() objTransform = objs[o].transform() for p in range(len(objSelection)): value = objSelection[p] if value >= selectionThreshold: pos = objPoints[p] tPos = objTransform * _nukemath.Vector4(pos.x, pos.y, pos.z, 1.0) return VertexInfo(o, p, value, _nukemath.Vector3(tPos.x, tPos.y, tPos.z)) return None
[docs]def allNodes(node = nuke.root()): ''' allNodes() -> iterator Return an iterator which yields all nodes in the current script. This includes nodes inside groups. They will be returned in top-down, depth-first order. ''' yield node if hasattr(node, "nodes"): for child in node.nodes(): for n in allNodes(child): yield n
[docs]def allNodesWithGeoSelectKnob(): if nuke.activeViewer(): preferred_nodes = [n for n in nuke.activeViewer().getGeometryNodes() if 'geo_select' in n.knobs()] else: preferred_nodes = [] nodes = preferred_nodes + [n for n in allNodes() if 'geo_select' in n.knobs() and n not in preferred_nodes] return nodes
[docs]def cameraProjectionMatrix(cameraNode): 'Calculate the projection matrix for the camera based on its knob values.' # Matrix to transform points into camera-relative coords. camTransform = cameraNode['transform'].value().inverse() # Matrix to take the camera projection knobs into account roll = float(cameraNode['winroll'].getValue()) scale_x, scale_y = [float(v) for v in cameraNode['win_scale'].getValue()] translate_x, translate_y = [float(v) for v in cameraNode['win_translate'].getValue()] m = _nukemath.Matrix4() m.makeIdentity() m.rotateZ(math.radians(roll)) m.scale(1.0 / scale_x, 1.0 / scale_y, 1.0) m.translate(-translate_x, -translate_y, 0.0) # Projection matrix based on the focal length, aperture and clipping planes of the camera focal_length = float(cameraNode['focal'].getValue()) h_aperture = float(cameraNode['haperture'].getValue()) near = float(cameraNode['near'].getValue()) far = float(cameraNode['far'].getValue()) projection_mode = int(cameraNode['projection_mode'].getValue()) p = _nukemath.Matrix4() p.projection(focal_length / h_aperture, near, far, projection_mode == 0) # Matrix to translate the projected points into normalised pixel coords format = nuke.root()['format'].value() imageAspect = float(format.height()) / float(format.width()) t = _nukemath.Matrix4() t.makeIdentity() t.translate( 1.0, 1.0 - (1.0 - imageAspect / float(format.pixelAspect())), 0.0 ) # Matrix to scale normalised pixel coords into actual pixel coords. x_scale = float(format.width()) / 2.0 y_scale = x_scale * format.pixelAspect() s = _nukemath.Matrix4() s.makeIdentity() s.scale(x_scale, y_scale, 1.0) # The projection matrix transforms points into camera coords, modifies based # on the camera knob values, projects points into clip coords, translates the # clip coords so that they lie in the range 0,0 - 2,2 instead of -1,-1 - 1,1, # then scales the clip coords to proper pixel coords. return s * t * p * m * camTransform
[docs]def projectPoints(camera=None, points=None): ''' projectPoint(camera, points) -> list of nuke.math.Vector2 Project the given 3D point through the camera to get 2D pixel coordinates. @param camera: The Camera node or name of the Camera node to use for projecting the point. @param points: A list or tuple of either nuke.math.Vector3 or of list/tuples of three float values representing the 3D points. @raise ValueError: If camera or point is invalid. ''' camNode = None if isinstance(camera, nuke.Node): camNode = camera elif isinstance(camera, str): camNode = nuke.toNode(camera) else: raise ValueError("Argument camera must be a node or the name of a node.") camMatrix = cameraProjectionMatrix(camNode) if camMatrix == None: raise RuntimeError("snap3d.cameraProjectionMatrix() returned None for camera.") if not ( isinstance(points, list) or isinstance(points, tuple) ): raise ValueError("Argument points must be a list or tuple.") for point in points: # Would be nice to not do this for every item but since lists/tuples can # containg anything... if isinstance(point, nuke.math.Vector3): pt = point elif isinstance(point, list) or isinstance(point, tuple): pt = nuke.math.Vector3(point[0], point[1], point[2]) else: raise ValueError("All items in points must be nuke.math.Vector3 or list/tuple of 3 floats.") tPos = camMatrix * nuke.math.Vector4(pt.x, pt.y, pt.z, 1.0) yield nuke.math.Vector2(tPos.x / tPos.w, tPos.y / tPos.w)
[docs]def projectPoint(camera=None, point=None): ''' projectPoint(camera, point) -> nuke.math.Vector2 Project the given 3D point through the camera to get 2D pixel coordinates. @param camera: The Camera node or name of the Camera node to use for projecting the point. @param point: A nuke.math.Vector3 or of list/tuple of three float values representing the 3D point. @raise ValueError: If camera or point is invalid. ''' return next(projectPoints( camera, (point,) ))
[docs]def projectSelectedPoints(cameraName='Camera1'): ''' projectSelectedPoints(cameraName='Camera1') -> iterator yielding nuke.math.Vector2 Using the specified camera, project all of the selected points into 2D pixel coordinates and return their locations. @param cameraName: Optional name of the Camera node to use for projecting the points. If omitted, will look for a node called Camera1. ''' camNode = nuke.toNode(cameraName) camMatrix = cameraProjectionMatrix(camNode) for pt in selectedPoints(): tPos = camMatrix * nuke.math.Vector4(pt.x, pt.y, pt.z, 1.0) yield nuke.math.Vector2(tPos.x / tPos.w, tPos.y / tPos.w)
# # Managing the list of snapping functions. # # The list of snapping functions. Treat this as a read-only variable; if you # want to add a new snapping function call addSnapFunc (below) snapFuncs = []
[docs]def addSnapFunc(label, func): ''' addSnapFunc(label, func) -> None Add a new snapping function to the list. The label parameter is the text that will appear in (eg.) an Enumeration_Knob for the function. It cannot be the same as any existing snap function label (if it is, the function will abort without changing anything). The func parameter is the snapping function. It must be a callable object taking a single parameter: the node to perform the snapping on. ''' if label in dict(snapFuncs): return snapFuncs.append( (label, func) )
[docs]def callSnapFunc(nodeToSnap=None): ''' callSnapFunc(nodeToSnap) -> None Call the snapping function on a node. The nodeToSnap parameter is optional. If it's not specified, or is None, we use the result of nuke.thisNode() instead. The node must have an Enumeration_Knob called "snapFunc" which selects the snapping function to call. ''' if nodeToSnap is None: nodeToSnap = nuke.thisNode() # Make sure that the nodeToSnap has a snapFunc knob if "snapFunc" not in nodeToSnap.knobs(): # TODO: warn the user that we can't snap this node. return snapFunc = dict(snapFuncs)[nodeToSnap['snapFunc'].value()] snapFunc(nodeToSnap)