Machine Learning Image Upscaling in Gaffer (Patreon)

Following up on my recent article dipping a toe into Gaffer's new machine learning tools, let's learn how we can implement the open source ONNX Super Resolution model.

ONNX Super Resolution upscales an input image by 3x. The model is hardcoded to accept an image at 224x224 resolution and output one at 672x672. While chatting about this topic with the Gaffer devs, John, Chief Gaffer, suggested we could use the model to upscale any size image by first dividing it into 224x224 tiles.

This ended up working great!

Let's walk though how it works. First, download the .onnx file here: https://github.com/onnx/models/tree/main/validated/vision/super_resolution/sub_pixel_cnn_2016 You can check out my post about how to enable ONNX models in Gaffer here: https://www.patreon.com/posts/120644489 

Note that all scene files, scripts and code will be available to download in the GitHub link at the end of this article.

Moving on! We know we need to have a set of 224x224 image tiles. The ONNX model will throw an error with any other size. It's easy enough to divide up an image using oiiotool, which comes included with Gaffer and many other CG apps.

All ONNX Super Resolution output images contain predictable edge artifacts. To compensate, we need to crop the upscaled tile by 9 pixels on each edge. This means we'll have to overlap our tiles since we'll be losing image data as a result of the 18 pixel crop. On our initial 224x224 tile we'll be losing 3 pixels around the edge. The tile offset should therefore be -6 pixels on X and Y, incrementing every subsequent tile. For example, tile 1 we could leave unchanged, tile 2 would be -6 pixels on X, tile 3 -12 on X, etc. The same process would be required on the Y dimension.

This is all easily accomplished using nested loops. I wrote a PowerShell script that runs oiiotool through two loops and saves all the tiles to disk.

Great, we've got our offset tiles! What next? I started by reading them into Gaffer as a sequence of frames. In the case of a 1920x1080 image we have 45 tiles, so 45 frames.

The ONNX example Python script shows that the model needs a YCbCr (luminance, blue-difference, red-difference) image for its input. The model upscales only the Y channel. Then you upscale Cb and Cr using normal means, combine the new ONNX-ified luminance channel and that's your final image.

Gaffer includes fantastic support for OSL (Open Shading Language), so I used this excellent, simple language to perform the RGB to YCbCr translation. Here we put the code needed to get us into the new color space into an OSL Code node.

We then feed this converted image data into the ONNX model and retrieve the upscaled data using a Tensor to Image node. Next, perform a normal upscale on Cb and Cr and pipe in the new luminance (Y) channel with a Copy Channels node. (In this case, selecting "R" simply means the first of the three channels, not "red" specifically.)

We now need to get back to RGB land. I wrote some more OSL code to reverse the trip, easy!

We've now got our 3x upscaled tile, but as promised there are 9 pixels of artifacts around the edge.

We'll use a Crop node to produce our final tile, sized 654x654.

Gaffer provides a great way to reassemble our tiles: a Contact Sheet. This is most often used for creating collections of images for later analysis and inspection. For example, you might have 50 shoe renders with different materials and lighting. You could use a contact sheet to view them all at once, additionally leveraging Gaffer's tools to automatically label each one. It's very powerful.

In our case, we can create a contact sheet sourced from all 45 processed frames/tiles. Gaffer will stitch them together into one image, creating our final high resolution output.

There are some black bars on the right and bottom, so we'll need to do one last small crop to get into a "normal" aspect ratio, then we can do a final resize to standard 4K resolution.

This is all great, but it's not automatic since we need to perform the initial tile creation manually. Daniel on the Gaffer crew came up with an impressive solution that uses Python scripting and another Contact Sheet to create the initial tiles, including offsets.

Check out the complete Gaffer scene file here on GitHub: https://github.com/BrianHanke/GafferML

Manual Mode relies on the included PowerShell script to create the initial tiles, then the rest of the upscaling process happens within Gaffer as I described above.

Magic Mode is courtesy of Daniel and produces the same result, all within Gaffer. Feel free to study the node tree and scripting to see how he did it. It's a great showcase of the power and flexibility of Gaffer for creating advanced tools and workflows. Maybe we'll revisit this topic in particular in a future article.

Until then, happy upscaling!