Artistic Style Transfer in Ruby

The ONNX Model Zoo has a number of interesting pretrained deep learning models. Thanks to the ONNX Runtime, we can run them in Ruby.

Today, we’ll look at artistic style transfer. Here’s the model we’ll use.

First, download the pretrained model and this awesome shot of a lynx.

Photo from the U.S. Fish and Wildlife Service

Install the ONNX Runtime, MiniMagick, and Numo::NArray gems. MiniMagick allows us to manipulate images, and Numo::NArray makes it easier to work with multi-dimensional arrays.

gem "onnxruntime"
gem "mini_magick"
gem "numo-narray"

Next, load the image. We resize it to be the model dimensions.

img = MiniMagick::Image.open("lynx.jpg")
img.resize "224x224^", "-gravity", "center", "-extent", "224x224"
pixels = img.get_pixels

And load the model

model = OnnxRuntime::Model.new("rain_princess.onnx")

Perform the preprocessing steps from the model docs

pixels = Numo::NArray.cast(img.get_pixels)
pixels = pixels.transpose(2, 0, 1)
pixels = pixels.expand_dims(0)

Run the model

result = model.predict(input1: pixels)

Perform the postprocessing steps

out_pixels = Numo::NArray.cast(result["output1"].first)
out_pixels = out_pixels.clip(0, 255).cast_to(Numo::UInt8)
out_pixels = out_pixels.transpose(1, 2, 0)

And save the image

img = MiniMagick::Image.import_pixels(out_pixels.to_binary, img.width, img.height, 8, "rgb", "jpg")
img.write("output.jpg")

Four other styles are also available.

Here’s the complete code. Now go out and try it with your own images!