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 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")

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

Gem Patterns

Scaling Reads

Argon2 with Devise


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