Nuke_Workgroup_12.2/pythonpluginsUI/denoice/denoice.py

import os
import nuke
import nukescripts
import re
import subprocess
import tempfile

# martin@keller.io
# numpy install:
# https://jurajtomori.wordpress.com/2018/11/28/nuke-tip-installing-numpy-for-nuke-python-on-windows/

'''noice 6.2.1.0 [359aaa66] - the Arnold denoiser
Usage:  noice [options] ...
  -i,    --input %s           Input image, multiple images can be included for stability, only first input image will be denoised
  -ef,   --extraframes %d     Number of frames to pad before and after for temporal stability (0 - 2, default:0)
  -f,    --frames %d          Number of frames to denoise in the sequence (default:1)
  -fe,   --feature %s         Use this AOV as a denoising feature
  -pr,   --patchradius %d     Neighborhood patch radius, size of pixel neighborhood to compare (0 - 6, default:3)
  -sr,   --searchradius %d    Search radius, higher values mean a wider search for similar pixel neighborhoods (6 - 21, default:9)
  -v,    --variance %f        Variance threshold, higher values mean more aggressive denoising (0.0 - 1.0, default:0.25)
  -l,    --aov %s             Light AOV name to be co-denoised, multiple can be specified
  -o,    --output %s          Output image
  -t,    --threads %d         Number of threads (0: use all cores, negative numbers specify number of idle cores)
  -irgba,--ignore_rgba        Ignore RGBA layer: will not denoise it or copy it to output file
  -n,    --notices            Print copyright notices'''

class denoice(KellerNukePlugin):

    def configurePlugin(self):
        menubar = nuke.menu("Nuke")
        self.m = menubar.addMenu("&Render")
        self.m.addCommand('Denoise with Noice', 'import denoice; denoice.addPanel()')


    def unconfigurePlugin(self):
        self.m.removeItem("Denoise with Noice")



#######################################################################################
#
# GLOBALS
#
#######################################################################################
DISABLE_AOVS = ['albedo', 'default', 'variance', 'Matte', 'matte', 'Crypto', 'crypto', 'denoise_albedo', 'P', 'N', 'Z']
# Muster
MSTRSERVER = "192.168.0.251"
SUBMISSION_USER ="admin"
PASSWORD = ""
MRTOOL = "\\\\hades\\p\\env\\muster9\\Mrtool.exe"
USERNAME = "admin"
TEMPLATE = "2008"
MUSTERFOLDER = "34758"
POOL = "arnold_16thread"
PSIZE = 1
PRIO = 10


def find_node_layers(node):
    node_channels = node.channels()
    store_layers = []
    for channel in node_channels:
        cut = channel.split(".")
        last = cut[0]
        store_layers.append(last)
    node_layers = list(set(store_layers))
    node_layers.sort()

    # put on top if rgb in layers
    if 'rgba' in node_layers:
        node_layers.remove("rgba")
        node_layers.insert(0, "rgba")
    return node_layers


class denoicePanel(nukescripts.PythonPanel):
    def __init__(self):
        nukescripts.panels.PythonPanel.__init__(self, "KellerDenoicePanel", scrollable=True)


        self.node = nuke.selectedNode()
        self.nodeLayers = find_node_layers(self.node)
        self.layerKnobs = []

        # just disable layers
        disable = [lay for lay in self.nodeLayers if any(x in lay for x in DISABLE_AOVS)]
        num = len(self.nodeLayers)
        for layer in self.nodeLayers:
            self.layer = nuke.Boolean_Knob(layer, layer, True)
            self.layer.setFlag(nuke.STARTLINE)
            self.addKnob(self.layer)
            self.layerKnobs.append(self.layer)
            if layer in disable:
                self.layer.setValue(False)

        # interface selection
        self.div1 = nuke.Text_Knob("div1", " ", "")
        self.addKnob(self.div1)
        self.all = nuke.PyScript_Knob("select_all", " Select All ")
        self.addKnob(self.all)
        self.invert = nuke.PyScript_Knob("invert_selection", " Invert Selection ")
        self.addKnob(self.invert)
        self.clear = nuke.PyScript_Knob("clear_selection", " Clear Selection ")
        self.addKnob(self.clear)
        self.div2 = nuke.Text_Knob("div2", " ", "")
        self.addKnob(self.div2)

        # settings
        # variance
        self.variance = nuke.Double_Knob("variance", "Variance", arraysize=1)
        self.variance.setRange(0,1)
        self.variance.setTooltip(
            'Default: 0.25\nThe strength of the filter is determined by the variance parameter, the higher the variance the more forceful the denoising will be. '
            'For variance maybe 0.25 (the default) / 0.5 / 0.75 are good low/min/max values. How aggressive the Arnold denoiser is in removing noise can be controlled by '
            'setting a variance threshold with the command line argument -variance (-v). The default value is 0.25, higher values will make the denoising more aggressive by '
            'considering similar neighborhoods that have bigger color disparities.')
        self.addKnob(self.variance)
        self.variance.setValue(0.25)

        # patchradius
        self.patchradius = nuke.Int_Knob("patchradius", "Patchradius", arraysize=1)
        self.patchradius.setRange(0, 10)
        self.patchradius.setTooltip(
            'Default: 3\nThis increases the softness of the denoising (while always preserving the features). The default is 3 (on the big side), but a low value would be maybe 0 or 1, middle 3, and high 5. '
            'For every pixel, the Arnold denoiser will consider its neighborhood patch and look for other pixels with similar neighborhood patches. The radius of this neighborhood can be controlled with the -patchradius '
            ' (or -pr) command-line argument. The default value is set to 3, which gives a 7x7 square neighborhood.')
        self.addKnob(self.patchradius)
        self.patchradius.setValue(3)

        # searchradius
        self.searchradius = nuke.Int_Knob("searchradius", "Searchradius", arraysize=1)
        self.searchradius.setRange(0, 50)
        self.searchradius.setTooltip(
            'Default: 9\nThis is the area over which similar neighborhoods are found. The higher the better, but it will increase '
            'the cost of denoising. For every pixel noice will search a square area with a radius set with the '
            'command line argument -searchradius (-sr). The bigger this area the bigger the denoising stability and '
            'the higher the chance that similar neighborhoods to be considered will be found. The default value is 9, '
            'which gives a 19x19 square neighborhood. Setting it to 21 (a search window of 42 x 42) will look over a '
            'pixel area equivalent to loading 5 frames.')
        self.addKnob(self.searchradius)
        self.searchradius.setValue(9)

        # extraFrames
        self.temporalFrames = nuke.Int_Knob("temporalFrames", "ExtraFrames", arraysize=1)
        self.temporalFrames.setRange(0, 2)
        self.temporalFrames.setTooltip('This argument (-ef n or --extraframes n) specifies how many additional source frames before and after the current one should be used, for improved stability in animation sequences.\n(0-2. default:0)')
        self.addKnob(self.temporalFrames)
        self.temporalFrames.setValue(0)

        # threads
        self.threads = nuke.Int_Knob("threads", "Threads", arraysize=1)
        self.threads.setTooltip('Number of threads (0: use all cores, negative numbers specify number of idle cores)')
        self.addKnob(self.threads)
        self.threads.setValue(-2)

        # div
        self.div3 = nuke.Text_Knob("div3", " ", "")
        self.addKnob(self.div3)

        # filepatn in
        self.inputpath = nuke.File_Knob("input", "Input")
        self.addKnob(self.inputpath)
        self.inputpath.setTooltip('Needs to be in #### syntax. Not in %04d')
        self.inputpath.setValue(nuke.selectedNode().knobs()['file'].value().replace('%04d','####'))

        # filepath out
        self.outputpath = nuke.File_Knob("output", "Output")
        file = nuke.selectedNode().knobs()['file'].value().replace('%04d', '####')
        file = file.replace('####', 'denoise.####')
        self.outputpath.setValue(file)
        self.addKnob(self.outputpath)

        # set path
        self.setpath = nuke.PyScript_Knob("setpath", " Set Path ")
        self.layer.setFlag(nuke.STARTLINE)
        self.addKnob(self.setpath)

        # startframe
        self.startframe = nuke.Int_Knob("startframe", "Startframe")
        self.addKnob(self.startframe)

        # endframe
        self.endframe = nuke.Int_Knob("endframe", "Endframe")
        self.addKnob(self.endframe)
        self.endframe.clearFlag(nuke.STARTLINE)

        # set to input
        self.setToInput = nuke.PyScript_Knob("setToInput", " Set to Input ")
        self.addKnob(self.setToInput)
        selected = nuke.selectedNode()
        fr = selected.frameRange()
        start = fr.first()
        end = fr.last()
        self.startframe.setValue(start)
        self.endframe.setValue(end)
        self.setToInput.clearFlag(nuke.STARTLINE)

        # set framerange
        self.setFrameRange = nuke.PyScript_Knob("setFrameRange", " Set Framerange ")
        self.addKnob(self.setFrameRange)
        self.setFrameRange.clearFlag(nuke.STARTLINE)

        # filepath
        self.arnoldBinPath = nuke.File_Knob("arnoldBinPath", " Arnold Bin Path ")
        self.addKnob(self.arnoldBinPath)
        self.arnoldBinPath.setValue('$env(PROJECT_ROOT_3D)/000_env/arnold/_noice/bin/')
        self.arnoldBinPath.setTooltip('Please be sure to select the bin folder with a "/" at the end.')

        # div
        self.div4 = nuke.Text_Knob("div4", " ", "")
        self.addKnob(self.div4)

        # render mode
        self.renderChoice = nuke.Enumeration_Knob('renderChoice', 'Render On', ['muster', 'local_machine'])
        self.renderChoice.setTooltip('local_machine: creates a .bat file with all frames, saves it to the temp directory and runs it on the local machine.\nmuster: sends the job to muster.')
        self.addKnob(self.renderChoice)

        # muster pool
        self.musterPool = nuke.String_Knob("musterPool", "Muster Pool")
        self.musterPool.setValue(POOL)
        self.addKnob(self.musterPool)

        # muster pool
        self.musterFolder = nuke.String_Knob("musterFolder", "Muster Folder ID")
        self.musterFolder.setValue(MUSTERFOLDER)
        self.addKnob(self.musterFolder)

        # div
        self.div5 = nuke.Text_Knob("div5", " ", "")
        self.addKnob(self.div5)

        # window
        #height = len(self.layerKnobs) * 20
        height = 1200
        width = 800
        #if height > 1200:
        #    width = 700
        #    height = 1200
        self.setMinimumSize(width, height)
        #self.setMaximumSize(800, 1050)


    def knobChanged(self, knob):
        if knob.name() == "select_all":
            for layer in self.layerKnobs:
                layer.setValue(True)

        if knob.name() == "invert_selection":
            selected = []
            for layer in self.layerKnobs:
                if layer.value() is True:
                    selected.append(layer)
                    layer.setValue(False)
            inverse = list(set(self.layerKnobs).difference(set(selected)))
            for check in inverse:
                check.setValue(True)

        if knob.name() == "clear_selection":
            for layer in self.layerKnobs:
                layer.setValue(False)

        if knob.name() == "setToInput":
            selected = nuke.selectedNode()
            fr = selected.frameRange()
            start = fr.first()
            end = fr.last()
            self.startframe.setValue(start)
            self.endframe.setValue(end)

        if knob.name() == "setFrameRange":
            self.startframe.setValue(int(nuke.root()['first_frame'].value()))
            self.endframe.setValue(int(nuke.root()['last_frame'].value()))

        if knob.name() == "setpath":
            file = nuke.selectedNode().knobs()['file'].value().replace('%04d','####')
            file = file.replace('####','denoise.####')
            self.outputpath.setValue(file)

        if knob.name() == "renderChoice":
            if self.renderChoice.value() == "local_machine":
                self.musterPool.setVisible(False)
                self.musterFolder.setVisible(False)
            if self.renderChoice.value() == "muster":
                self.musterPool.setVisible(True)
                self.musterFolder.setVisible(True)


        if knob.name() == "OK":
            if self.renderChoice.value() == "local_machine":
                print('Denoising locally......')
                self.executeNoiceLocally()
            if self.renderChoice.value() == "muster":
                print('Denoising on Muster......')
                self.sendToMuster()


    def executeNoiceLocally(self):
        # saving selected aovs
        light_aov_names = []
        for layer in self.layerKnobs:
            if layer.value() is True:
                light_aov_names.append(layer.name())

        binpath = self.arnoldBinPath.evaluate().replace('\\','/')

        filename_input_sequence_number = re.sub(r'#+', lambda m: r'{{:0{}d}}'.format(len(m.group(0))), self.inputpath.value().replace('%04d','####'))
        filename_output_sequence_number = re.sub(r'#+', lambda m: r'{{:0{}d}}'.format(len(m.group(0))), self.outputpath.value().replace('%04d','####'))

        command = ''
        for i in range(self.startframe.value(), self.endframe.value() + 1):
            light_aov_string = ""

            for k in light_aov_names:
                print('Denoice: Info. AOV: %s') % k
                light_aov_string += "-aov " + str(k) + " "

            command += '"' + binpath + '/noice.exe' + '"' + ' ' \
                  + '-patchradius ' + str(self.patchradius.value()) + ' ' \
                  + '-searchradius ' + str(self.searchradius.value()) + ' ' \
                  + '-variance ' + str(self.variance.value()) + ' ' \
                  + light_aov_string \
                  + '-i ' + filename_input_sequence_number.format(i) + ' ' \
                  + '-t ' + str(self.threads.value()) + ' ' \
                  + '-ef ' + str(self.temporalFrames.value()) + ' ' \
                  + '-output ' + filename_output_sequence_number.format(i)

            print('_________________________________________________________________________________________________')
            print('')
            print('--------------------------------> Denoise Command for Frame %s <---------------------------') %i
            print('_________________________________________________________________________________________________')
            print('')
            print('')
            print('')

            if i is not self.endframe.value():
                command += "\n"
        print(command)

        temp = tempfile.NamedTemporaryFile(prefix='Noice_', suffix='.bat', delete=False)
        temp.write(command)
        print('Created temp file is:', temp.name)
        temp.close()

        filepath = 'cmd /c start "Denoice " "%s"' %temp.name
        p = subprocess.call(filepath, shell=True, stdout=subprocess.PIPE)


    def fixfilename(self,infile):

        infile = infile.replace("p:/", "//hades/p/")
        infile = infile.replace("P:/", "//hades/p/")
        infile = infile.replace("o:/", "//calculon/o/")
        infile = infile.replace("O:/", "//calculon/o/")
        infile = infile.replace("/", "\\")

        return infile


    def sendToMuster(self):
        # jobname
        jobname = nuke.selectedNode().knobs()['file'].value().split("/")[-1].replace('.%04d.exr',"__denoise")
        # aovs
        # saving selected aovs
        light_aov_names = []
        light_aov_string = ""
        for layer in self.layerKnobs:
            if layer.value() is True:
                light_aov_names.append(layer.name())
        for k in light_aov_names:
            print('Denoice: Info. AOV: %s') % k
            light_aov_string += "-l " + str(k) + " "
        light_aov_string = light_aov_string.rstrip()
        light_aov_string = '"' + light_aov_string + '"'

        a_bin_path = self.fixfilename(self.arnoldBinPath.evaluate())
        a_input_path = self.fixfilename(self.inputpath.evaluate().replace('%04d','####'))
        a_output_path = self.fixfilename(self.outputpath.evaluate().replace('%04d','####'))

        # building Command
        cmd = '%s -s %s -u %s -b -e %s -n %s' % (MRTOOL, MSTRSERVER, USERNAME, TEMPLATE, jobname)
        cmd += ' -pk %i -pr %i -pool "%s" -parent %i' % (int(PSIZE), int(PRIO), self.musterPool.value(), int(self.musterFolder.value()))
        cmd += ' -attr "NOICEBIN" %s 0' % a_bin_path
        cmd += ' -attr "PATCHRADIUS" %i 0' % int(self.patchradius.value())
        cmd += ' -attr "SEARCHRADIUS" %i 0' % int(self.searchradius.value())
        cmd += ' -attr "VARIANCE" %s 0' % self.variance.value()
        cmd += ' -sf %i' % int(self.startframe.value())
        cmd += ' -ef %i' % int(self.endframe.value())
        cmd += ' -bf %i' % 1
        cmd += ' -f %s' % a_input_path

        cmd += ' -attr "LIGHTAOVSTRING" %s 0' % light_aov_string
        cmd += ' -attr "TEMPORALFRAMES" %i 0' % int(self.temporalFrames.value())
        cmd += ' -attr "DENOISETHREADS" %i 0' % int(self.threads.value())
        cmd += ' -attr "OUTPUTFILE" %s 0' % a_output_path

        # send
        print (cmd)
        print (os.system(cmd))
        return


def addPanel():
    try:
        selected = nuke.selectedNode()
    except:
        nuke.message("Denoice: Error. Please select a Read node with EXR file")
        return
    if selected.Class() == "Read":
        if os.path.splitext(selected["file"].value())[1] == ".exr":
            launch_panel = denoicePanel()
            launch_panel.showModalDialog()
        else:
            nuke.message("Denoice: Error. Please select a Read node with EXR file")
    else:
        nuke.message("Denoice: Error. Please select a Read node")