/share/nuke/ocionuke/cdl.py

  1# SPDX-License-Identifier: BSD-3-Clause
  2# Copyright Contributors to the OpenColorIO Project.
  3
  4"""Various utilities relating to the OCIOCDLTransform node
  5"""
  6
  7import nuke
  8import PyOpenColorIO as OCIO
  9import xml.etree.ElementTree as ET
 10
 11
 12def _node_to_cdltransform(node):
 13    """From an OCIOCDLTransform node, returns a PyOpenColorIO
 14    CDLTransform object, which could be used to write XML
 15    """
 16
 17    # Color_Knob.value returns single float if control is not
 18    # expanded, so use value(index=...) to always get three values
 19    slope = [node['slope'].value(x) for x in range(3)]
 20    offset = [node['offset'].value(x) for x in range(3)]
 21    power = [node['power'].value(x) for x in range(3)]
 22    sat = node['saturation'].value()
 23    cccid = node['cccid'].value()
 24
 25    cdl = OCIO.CDLTransform()
 26    cdl.setSlope(slope)
 27    cdl.setOffset(offset)
 28    cdl.setPower(power)
 29    cdl.setSat(sat)
 30    cdl.setID(cccid)
 31
 32    return cdl
 33
 34
 35def _cdltransform_to_node(cdl, node):
 36    """From an XML string, populates the parameters on an
 37    OCIOCDLTransform node
 38    """
 39
 40    # Treat "node" as a dictionary of knobs, as the "node" argument could be
 41    # a the return value of PythonPanel.knob(), as in SelectCCCIDPanel
 42
 43    node['slope'].setValue(cdl.getSlope())
 44    node['offset'].setValue(cdl.getOffset())
 45    node['power'].setValue(cdl.getPower())
 46    node['saturation'].setValue(cdl.getSat())
 47    node['cccid'].setValue(cdl.getID())
 48
 49
 50def _xml_colorcorrection_to_cdltransform(xml, cccposition):
 51    """From an XML string, return a CDLTransform object, if the cccid
 52    is absent from the XML the given cccposition will be used instead 
 53    """
 54    ccxml = ET.tostring(xml)
 55    cdl = OCIO.CDLTransform()
 56    cdl.setXML(ccxml)
 57    if "cccid" not in xml.getchildren():
 58	cdl.setID(str(cccposition))
 59
 60    return cdl
 61
 62def _xml_to_cdltransforms(xml):
 63    """Given some XML as a string, returns a list of CDLTransform
 64    objects for each ColorCorrection (returns a one-item list for a
 65    .cc file)
 66    """
 67
 68    tree = ET.fromstring(xml)
 69
 70    # Strip away xmlns
 71    for elem in tree.getiterator():
 72        if elem.tag.startswith("{"):
 73            elem.tag = elem.tag.partition("}")[2]
 74
 75    filetype = tree.tag
 76
 77    cccposition = 0
 78    if filetype == "ColorCorrection":
 79        return [_xml_colorcorrection_to_cdltransform(tree, cccposition)]
 80
 81    elif filetype in ["ColorCorrectionCollection", "ColorDecisionList"]:
 82        allcdl = []
 83        for cc in tree.getchildren():
 84            if cc.tag == "ColorDecision":
 85                cc = cc.getchildren()[0] # TODO: something better here
 86            if cc.tag != "ColorCorrection": continue
 87            allcdl.append(_xml_colorcorrection_to_cdltransform(cc, cccposition))
 88            cccposition += 1
 89        return allcdl
 90
 91    else:
 92        raise RuntimeError(
 93            "The supplied file did not have the correct root element, expected"
 94            " 'ColorCorrection' or 'ColorCorrectionCollection' or 'ColorDecisionList', got %r" % (filetype))
 95
 96
 97def _cdltransforms_to_xml(allcc):
 98    """Given a list of CDLTransform objects, returns an XML string
 99    """
100
101    root = ET.Element("ColorCorrectionCollection")
102    root.attrib['xmlns'] = 'urn:ASC:CDL:v1.2'
103
104    for cc in allcc:
105        cur = ET.fromstring(cc.getXML())
106
107        # Strip away xmlns
108        for elem in cur.getiterator():
109            if elem.tag.startswith("{"):
110                elem.tag = elem.tag.partition("}")[2]
111        root.append(cur)
112
113    return ET.tostring(root)
114
115
116def SelectCCCIDPanel(*args, **kwargs):
117    # Wrap class definition in a function, so nukescripts.PythonPanel
118    # is only accessed when ``SelectCCCIDPanel()`` is called,
119    # https://github.com/AcademySoftwareFoundation/OpenColorIO/issues/277
120
121    import nukescripts
122
123    class _SelectCCCIDPanel(nukescripts.PythonPanel):
124        """Allows the user to select from a list of CDLTransform
125        objects
126        """
127
128        def __init__(self, allcdl):
129            super(_SelectCCCIDPanel, self).__init__()
130            self.available = {}
131            for cur in allcdl:
132                self.available[cur.getID()] = cur
133
134            self.addKnob(nuke.Enumeration_Knob("cccid", "cccid", self.available.keys()))
135            self.addKnob(nuke.Text_Knob("divider"))
136            self.addKnob(nuke.Color_Knob("slope"))
137            self.addKnob(nuke.Color_Knob("offset"))
138            self.addKnob(nuke.Color_Knob("power"))
139            self.addKnob(nuke.Double_Knob("saturation"))
140
141        def selected(self):
142            return self.available[self.knobs()['cccid'].value()]
143
144        def knobChanged(self, knob):
145            """When the user selects a cccid, a grade-preview knobs are set.
146
147            This method is triggered when any knob is changed, which has the
148            useful side-effect of preventing changing the preview values, while
149            keeping them selectable for copy-and-paste.
150            """
151            _cdltransform_to_node(self.selected(), self.knobs())
152
153    return _SelectCCCIDPanel(*args, **kwargs)
154
155
156def export_as_cc(node = None, filename = None):
157    """Export a OCIOCDLTransform node as a ColorCorrection XML file
158    (.cc)
159
160    If node is None, "nuke.thisNode()" will be used. If filename is
161    not specified, the user will be prompted.
162    """
163
164    if node is None:
165        node = nuke.thisNode()
166
167    cdl = _node_to_cdltransform(node)
168
169    if filename is None:
170        ccfilename = nuke.getFilename("Color Correction filename", pattern = "*.cc")
171        if ccfilename is None:
172            # User clicked cancel
173            return
174
175    xml = cdl.getXML()
176    print "Writing to %s - contents:\n%s" % (ccfilename, xml)
177    open(ccfilename, "w").write(xml)
178
179
180def import_cc_from_xml(node = None, filename = None):
181    """Import a ColorCorrection XML (.cc) into a OCIOCDLTransform node.
182
183    If node is None, "nuke.thisNode()" will be used. If filename is
184    not specified, the user will be prompted.
185    """
186
187    if node is None:
188        node = nuke.thisNode()
189
190    if filename is None:
191        ccfilename = nuke.getFilename("Color Correction filename", pattern = "*.cc *.ccc *.cdl")
192        if ccfilename is None:
193            # User clicked cancel
194            return
195
196    xml = open(ccfilename).read()
197
198    allcc = _xml_to_cdltransforms(xml)
199
200    if len(allcc) == 1:
201        _cdltransform_to_node(allcc[0], node)
202    elif len(allcc) > 1:
203        do_selectcccid = nuke.ask(
204            "Selected a ColorCorrectionCollection, do you wish to select a ColorCorrection from this file?")
205        if do_selectcccid:
206            sel = SelectCCCIDPanel(allcc)
207            okayed = sel.showModalDialog()
208            if okayed:
209                cc = sel.selected()
210                _cdltransform_to_node(cc, node)
211        else:
212            return
213    else:
214        nuke.message("The supplied file (%r) contained no ColorCorrection's" % ccfilename)
215        return
216
217
218def export_multiple_to_ccc(filename = None):
219    """Exported all selected OCIOCDLTransform nodes to a
220    ColorCorrectionCollection XML file (.ccc)
221    """
222
223    if filename is None:
224        filename = nuke.getFilename("Color Correction XML file", pattern = "*.cc *.ccc")
225        if filename is None:
226            # User clicked cancel
227            return
228
229    allcc = []
230    for node in nuke.selectedNodes("OCIOCDLTransform"):
231        allcc.append(_node_to_cdltransform(node))
232
233    xml = _cdltransforms_to_xml(allcc)
234    print "Writing %r, contents:\n%s" % (filename, xml)
235    open(filename, "w").write(xml)
236
237
238def import_multiple_from_ccc(filename = None):
239    """Import a ColorCorrectionCollection file (.ccc) into multiple
240    OCIOCDLTransform nodes. Also creates a single node for a .cc file
241    """
242
243    if filename is None:
244        filename = nuke.getFilename("Color Correction XML file", pattern = "*.cc *.ccc *.cdl")
245        if filename is None:
246            # User clicked cancel
247            return
248
249    xml = open(filename).read()
250    allcc = _xml_to_cdltransforms(xml)
251
252    def _make_node(cdl):
253        newnode = nuke.nodes.OCIOCDLTransform(inputs = nuke.selectedNodes()[:1])
254        _cdltransform_to_node(cdl, newnode)
255        newnode['label'].setValue("id: [value cccid]")
256
257    if len(allcc) > 0:
258        for cc in allcc:
259            _make_node(cc)
260    else:
261        nuke.message("The supplied file (%r) contained no ColorCorrection's" % filename)
262
263
264def select_cccid_for_filetransform(node = None, fileknob = 'file', cccidknob = 'cccid'):
265    """Select cccid button for the OCIOFileTransform node, also used
266    in OCIOCDLTransform. Presents user with list of cccid's within the
267    specified .ccc file, and sets the cccid knob to the selected ID.
268    """
269
270    if node is None:
271        node = nuke.thisNode()
272
273    filename = node[fileknob].value()
274
275    try:
276        xml = open(filename).read()
277    except IOError, e:
278        nuke.message("Error opening src file: %s" % e)
279        raise
280
281    allcc = _xml_to_cdltransforms(xml)
282
283    if len(allcc) == 0:
284        nuke.message("The file (%r) contains no ColorCorrection's")
285        return
286
287    sel = SelectCCCIDPanel(allcc)
288    okayed = sel.showModalDialog()
289    if okayed:
290        node[cccidknob].setValue(sel.selected().getID())