Custom Knobs

For some models, we will want to change some of the model parameters from Nuke in order to control the output image generated by the model. For example, if we are using an object detection model, we might like to add an enumeration knob in Nuke to control which object the model is detecting in the input image. Similarly, if our model changes the colour of a person’s hair, we might like to add an enumeration knob in Nuke to control which hair colour is transferred to the person in the image.

These custom knobs can be defined in the CatFileCreator node when creating the .cat file, and can be used to control parameters defined in the model’s init() function. When this .cat file is loaded in the Inference node, all of the custom knobs will be visible and can be used to control the model output from Nuke.

In the following sections, we’ll detail how to define your model with parameters that can be controlled from Nuke, how to create custom knobs using the CatFileCreator, and how these knobs can be used in the Inference node.

TorchScript

If we consider our simple Cifar10ResNet example, let’s change the Cifar10ResNet class definition slightly so that it has an attribute called userLabel:

class Cifar10ResNet(torch.nn.Module):
    """
    This class is a wrapper for our custom ResNet model.
    """
    def __init__(self, userLabel = 0):

        super(Cifar10ResNet, self).__init__()
        self.model = ResNet(pretrained=True)
        self.model.eval()
        self.userLabel = userLabel

This is the parameter that we are going to control using a knob in Nuke. Any variables that you want to control using a knob must be declared like this, as an attribute of your model class.

Next we’ll define the model forward function as follows

def forward(self, input):
      """
      :param input: A torch.Tensor of size 1 x 3 x H x W representing the input image
      :return: A torch.Tensor of size 1 x 1 x H x W of zeros or ones
      """

      modelOutput = self.model.forward(input)
      modelLabel = int(torch.argmax(modelOutput[0]))

      if modelLabel == self.userLabel:
         output = torch.ones(1, 1, input.shape[2], input.shape[3])
      else:
         output = torch.zeros(1, 1, input.shape[2], input.shape[3])
      return output

With this new forward() function, our model will return a tensor of ones if the object defined by self.userLabel is detected, and will return a tensor of zeros otherwise. So the value stored in userLabel will define which object is detected by the model. We can use the usual steps to convert this model from Pytorch to a TorchScript .pt file.

CatFileCreator

Now that we have created a model with an attribute, we can use the CatFileCreator to define a knob to control this attribute. First, add the main CatFileCreator knob values as before:

Next, click and drag an enumeration knob to the top of the CatFileCreator node, and move to the ‘User’ tab to fill in the enumeration values as follows:

The most important knob here is the Name knob. The knob value you enter here MUST correspond to the name of the attribute that you want to control, as defined in your model’s init() function. In our case, this attribute is called userLabel. For our Cifar10ResNet model, label values of 0, 1, 2, 3, .., 9 correspond to objects plane, car, bird, cat, deer, etc, so these are the values we add to the Menu Items knob.

Once our enumeration knob values have been added, we’ll see the knob appear in the main CatFileCreator tab. Clicking the Create .cat file and Inference knob will create the .cat file and Inference node.

Note

Class attributes with the Optional[ ] type cannot be controlled by a custom knob in Nuke.

Only certain Nuke knobs can be added as custom knobs to the Inference node. Please see the Appendix for the list of these knobs. More details on adding knobs using the CatFileCreator can be found in the Nuke documentation.

Inference

If we open the Inference node in the properties panel, we can see that our custom Detect enumeration knob has been added. Now, if we connect the Inference node to an image of a dog, and select ‘dog’ in the Detect enumeration knob, the alpha channel returned from the Inference node will be all ones, indicating that there is a dog in the image. Changing the selected option in the enumeration knob to ‘cat’ will return an alpha channel of zeros, indicating that there is no cat in the image.

Note

When connecting a model attribute to a check box knob in Nuke, make sure to define the boolean attribute in your model __init__ using integer values of 0 or 1, rather than True or False:

class Cifar10ResNet(torch.nn.Module):
   def __init__(self, bool = 0):