1 files changed, 269 insertions, 0 deletions
diff --git a/share/nuke/ocionuke/cdl.py b/share/nuke/ocionuke/cdl.py
new file mode 100644
index 0000000..9ed76ae
--- /dev/null
+++ b/share/nuke/ocionuke/cdl.py
@@ -0,0 +1,269 @@
+"""Various utilities relating to the OCIOCDLTransform node
+"""
+
+import nuke
+import nukescripts
+import PyOpenColorIO as OCIO
+import xml.etree.ElementTree as ET
+
+
+def _node_to_cdltransform(node):
+    """From an OCIOCDLTransform node, returns a PyOpenColorIO
+    CDLTransform object, which could be used to write XML
+    """
+
+    # Color_Knob.value returns single float if control is not
+    # expanded, so use value(index=...) to always get three values
+    slope = [node['slope'].value(x) for x in range(3)]
+    offset = [node['offset'].value(x) for x in range(3)]
+    power = [node['power'].value(x) for x in range(3)]
+    sat = node['saturation'].value()
+    cccid = node['cccid'].value()
+
+    cdl = OCIO.CDLTransform()
+    cdl.setSlope(slope)
+    cdl.setOffset(offset)
+    cdl.setPower(power)
+    cdl.setSat(sat)
+    cdl.setID(cccid)
+
+    return cdl
+
+
+def _cdltransform_to_node(cdl, node):
+    """From an XML string, populates the parameters on an
+    OCIOCDLTransform node
+    """
+
+    # Treat "node" as a dictionary of knobs, as the "node" argument could be
+    # a the return value of PythonPanel.knob(), as in SelectCCCIDPanel
+
+    node['slope'].setValue(cdl.getSlope())
+    node['offset'].setValue(cdl.getOffset())
+    node['power'].setValue(cdl.getPower())
+    node['saturation'].setValue(cdl.getSat())
+    node['cccid'].setValue(cdl.getID())
+
+
+def _xml_to_cdltransforms(xml):
+    """Given some XML as a string, returns a list of CDLTransform
+    objects for each ColorCorrection (returns a one-item list for a
+    .cc file)
+    """
+
+    tree = ET.fromstring(xml)
+
+    # Strip away xmlns
+    for elem in tree.getiterator():
+        if elem.tag.startswith("{"):
+             elem.tag = elem.tag.partition("}")[2]
+
+    filetype = tree.tag
+
+    if filetype == "ColorCorrection":
+        ccxml = ET.tostring(tree)
+        cdl = OCIO.CDLTransform()
+        cdl.setXML(ccxml)
+        return [cdl]
+
+    elif filetype == "ColorCorrectionCollection":
+        allcdl = []
+        for cc in tree.getchildren():
+            if cc.tag != "ColorCorrection": continue
+            ccxml = ET.tostring(cc)
+            cdl = OCIO.CDLTransform()
+            cdl.setXML(ccxml)
+            allcdl.append(cdl)
+        return allcdl
+
+    else:
+        raise RuntimeError(
+            "The supplied file did not have the correct root element, expected"
+            " 'ColorCorrection' or 'ColorCorrectionCollection', got %r" % (filetype))
+
+
+def _cdltransforms_to_xml(allcc):
+    """Given a list of CDLTransform objects, returns an XML string
+    """
+
+    root = ET.Element("ColorCorrectionCollection")
+    root.attrib['xmlns'] = 'urn:ASC:CDL:v1.2'
+
+    for cc in allcc:
+        cur = ET.fromstring(cc.getXML())
+
+        # Strip away xmlns
+        for elem in cur.getiterator():
+            if elem.tag.startswith("{"):
+                elem.tag = elem.tag.partition("}")[2]
+        root.append(cur)
+
+    return ET.tostring(root)
+
+
+class SelectCCCIDPanel(nukescripts.PythonPanel):
+    """Allows the user to select from a list of CDLTransform
+    objects
+    """
+
+    def __init__(self, allcdl):
+        super(SelectCCCIDPanel, self).__init__()
+        self.available = {}
+        for cur in allcdl:
+            self.available[cur.getID()] = cur
+
+        self.addKnob(nuke.Enumeration_Knob("cccid", "cccid", self.available.keys()))
+        self.addKnob(nuke.Text_Knob("divider"))
+        self.addKnob(nuke.Color_Knob("slope"))
+        self.addKnob(nuke.Color_Knob("offset"))
+        self.addKnob(nuke.Color_Knob("power"))
+        self.addKnob(nuke.Double_Knob("saturation"))
+
+    def selected(self):
+        return self.available[self.knobs()['cccid'].value()]
+
+    def knobChanged(self, knob):
+        """When the user selects a cccid, a grade-preview knobs are set.
+
+        This method is triggered when any knob is changed, which has the
+        useful side-effect of preventing changing the preview values, while
+        keeping them selectable for copy-and-paste.
+        """
+        _cdltransform_to_node(self.selected(), self.knobs())
+
+
+def export_as_cc(node = None, filename = None):
+    """Export a OCIOCDLTransform node as a ColorCorrection XML file
+    (.cc)
+
+    If node is None, "nuke.thisNode()" will be used. If filename is
+    not specified, the user will be prompted.
+    """
+
+    if node is None:
+        node = nuke.thisNode()
+
+    cdl = _node_to_cdltransform(node)
+
+    if filename is None:
+        ccfilename = nuke.getFilename("Color Correction filename", pattern = "*.cc")
+        if ccfilename is None:
+            # User clicked cancel
+            return
+
+    xml = cdl.getXML()
+    print "Writing to %s - contents:\n%s" % (ccfilename, xml)
+    open(ccfilename, "w").write(xml)
+
+
+def import_cc_from_xml(node = None, filename = None):
+    """Import a ColorCorrection XML (.cc) into a OCIOCDLTransform node.
+
+    If node is None, "nuke.thisNode()" will be used. If filename is
+    not specified, the user will be prompted.
+    """
+
+    if node is None:
+        node = nuke.thisNode()
+
+    if filename is None:
+        ccfilename = nuke.getFilename("Color Correction filename", pattern = "*.cc *.ccc")
+        if ccfilename is None:
+            # User clicked cancel
+            return
+
+    xml = open(ccfilename).read()
+
+    allcc = _xml_to_cdltransforms(xml)
+
+    if len(allcc) == 1:
+        _cdltransform_to_node(allcc[0], node)
+    elif len(allcc) > 1:
+        do_selectcccid = nuke.ask(
+            "Selected a ColorCorrectionCollection, do you wish to select a ColorCorrection from this file?")
+        if do_selectcccid:
+            sel = SelectCCCIDPanel(allcc)
+            okayed = sel.showModalDialog()
+            if okayed:
+                cc = sel.selected()
+                _cdltransform_to_node(cc, node)
+        else:
+            return
+    else:
+        nuke.message("The supplied file (%r) contained no ColorCorrection's" % ccfilename)
+        return
+
+
+def export_multiple_to_ccc(filename = None):
+    """Exported all selected OCIOCDLTransform nodes to a
+    ColorCorrectionCollection XML file (.ccc)
+    """
+
+    if filename is None:
+        filename = nuke.getFilename("Color Correction XML file", pattern = "*.cc *.ccc")
+        if filename is None:
+            # User clicked cancel
+            return
+
+    allcc = []
+    for node in nuke.selectedNodes("OCIOCDLTransform"):
+        allcc.append(_node_to_cdltransform(node))
+
+    xml = _cdltransforms_to_xml(allcc)
+    print "Writing %r, contents:\n%s" % (filename, xml)
+    open(filename, "w").write(xml)
+
+
+def import_multiple_from_ccc(filename = None):
+    """Import a ColorCorrectionCollection file (.ccc) into multiple
+    OCIOCDLTransform nodes. Also creates a single node for a .cc file
+    """
+
+    if filename is None:
+        filename = nuke.getFilename("Color Correction XML file", pattern = "*.cc *.ccc")
+        if filename is None:
+            # User clicked cancel
+            return
+
+    xml = open(filename).read()
+    allcc = _xml_to_cdltransforms(xml)
+
+    def _make_node(cdl):
+        newnode = nuke.nodes.OCIOCDLTransform(inputs = nuke.selectedNodes()[:1])
+        _cdltransform_to_node(cdl, newnode)
+        newnode['label'].setValue("id: [value cccid]")
+
+    if len(allcc) > 0:
+        for cc in allcc:
+            _make_node(cc)
+    else:
+        nuke.message("The supplied file (%r) contained no ColorCorrection's" % filename)
+
+
+def select_cccid_for_filetransform(node = None, fileknob = 'file', cccidknob = 'cccid'):
+    """Select cccid button for the OCIOFileTransform node, also used
+    in OCIOCDLTransform. Presents user with list of cccid's within the
+    specified .ccc file, and sets the cccid knob to the selected ID.
+    """
+
+    if node is None:
+        node = nuke.thisNode()
+
+    filename = node[fileknob].value()
+
+    try:
+        xml = open(filename).read()
+    except IOError, e:
+        nuke.message("Error opening src file: %s" % e)
+        raise
+
+    allcc = _xml_to_cdltransforms(xml)
+
+    if len(allcc) == 0:
+        nuke.message("The file (%r) contains no ColorCorrection's")
+        return
+
+    sel = SelectCCCIDPanel(allcc)
+    okayed = sel.showModalDialog()
+    if okayed:
+        node[cccidknob].setValue(sel.selected().getID())