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.

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 a multiple of the model dimensions since other dimensions can cause the model to fail.

img = MiniMagick::Image.open("lynx.jpg")
img.resize "448x448^", "-gravity", "center", "-extent", "448x448"
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[*img.get_pixels]
pixels = pixels.transpose(2, 0, 1)
pixels = pixels.expand_dims(0)

Run the model

result = model.predict(input1: pixels.to_a)

Perform the postprocessing steps

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

And save the image

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

Lynx with Rain Princess style

Four other styles are also available.

Lynx with Udnie style Lynx with Candy style Lynx with Mosaic style Lynx with Pointilism style

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

Published September 11, 2019 · Tweet


You might also enjoy

Introducing Archer: Rails Console History for Heroku, Docker, and More

Learn Data Science

Gem Patterns


All code examples are public domain.
Use them however you’d like (licensed under CC0).