Deploy machine learning models with R Shiny and ONNX

Export image classification model to ONNX

To showcase how ONNX models can be served in Shiny, we first have to obtain an image classification model. Here, I have used PyTorch to get ResNet18 model pre-trained on the ImageNet dataset, which classifies the input image as one of 1000 categories. TensorFlow could just as well have been used to get a model as all major frameworks can export models to the ONNX format.

#Import libraries
import torch
from torch import nn
from torchvision.models import resnet18, ResNet18_Weights
from torchvision import transforms

#Initialize pretrained model
#https://pytorch.org/vision/stable/models.html
weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
model.eval()

#Get transforms and class labels
preprocess = weights.transforms()
print(preprocess)

#Write class labels to text file
class_labels = weights.meta["categories"]

with open("classes.txt", "w") as file:
    file.write('\n'.join(class_labels))

For convenience, two additional layers are added to the model, a normalization layer with the ImageNet stats and a final softmax layer that converts model logits to ‘probabilities’ (model outputs summing to 1). The layers are assembled and the model is set to inference mode. Finally, the model is exported to the ONNX format.

#Add normalization and softmax layers to model
normalization_layer = transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                           std=[0.229, 0.224, 0.225])
softmax_layer = nn.Softmax(dim=1) 

#Assemble the final model
model_onnx = nn.Sequential(normalization_layer, model, softmax_layer).eval()

#Export to model to onnx
torch.onnx.export(model_onnx,
                  torch.randn(1, 3, 224, 224), 
                  "model.onnx",
                  input_names = ['input'],
                  output_names = ['output']) 

Adding the Shiny interface

Here, a very simple Shiny interface is created, it only has an upload button and a placeholder for the image .

library(shiny)
library(OpenImageR)
library(reticulate)

#Define UI with upload button and display for image
ui <- fluidPage(
  
  titlePanel("Image classification app"),
  
  sidebarLayout(
    sidebarPanel(
      fileInput("file", "Choose a file", accept = c('image/png', 'image/jpeg'))
    ),
    
    mainPanel(
      imageOutput("image", height = "224px", width = "224px"),
      textOutput("image_label")
      
    )
  )
)

Following the UI, the server part of the application is defined. This part sets up a Python virtual environment and installs the onnxruntime Python library. The model is loaded and the server then waits for an image to be uploaded. Once uploaded, the image is resized using the same preprocessing steps as used to train the PyTorch model using the OpenImageR package. The steps are as follows:

• resizing such the shorter side is 256 pixels
• center crop to final image size 224 by 224 pixel
• re-arrange array dimensions to channel, height, width
• add empty batch dimensions

The preprocessed image is used as input for the model which outputs class a vector of 1000 ‘probabilities’, from which the maximum probability and corresponding class can be determined. A character string containing the result and the image is returned to the UI.

server <- function(input, output) {
  
  #Create virtual env and install dependencies
  virtualenv_dir = Sys.getenv('VIRTUALENV_NAME')
  python_path = Sys.getenv('PYTHON_PATH')
  PYTHON_DEPENDENCIES = c('onnxruntime')
  
  virtualenv_create(envname = virtualenv_dir, python = python_path)
  
  virtualenv_install(virtualenv_dir, packages = PYTHON_DEPENDENCIES, 
                     ignore_installed=TRUE)
  
  use_virtualenv(virtualenv_dir, required = TRUE)
  
  #Load model using onnxruntime
  ort <- import("onnxruntime")
  model_path <- "model.onnx"
  ort_sess <- ort$InferenceSession(model_path)
  
  #Read classes from text file
  classes <- readLines("classes.txt")

  observeEvent(input$file, {
    
    inFile <- input$file
    
    if (is.null(inFile)){
      return()
    }
    
    #Read input image
    img <- readImage(input$file$datapath)

    img_h <- dim(img)[2]
    img_w <- dim(img)[1]
  
    #Determine size for preprocessing (similar to Pytorch preprocessing steps)
    init_size <- 256
    final_size <- 224
    
    if(img_h > img_w){
      new_h <- init_size * (img_h/img_w)
      new_w <- init_size
    }else{
      new_h <- init_size 
      new_w <- init_size / (img_h/img_w)
    }
    
    #Resize image
    img_resize <- resizeImage(img, width = new_w, height = new_h, 
                              method="bilinear", normalize_pixels = TRUE)
    
    #Center crop image
    img_crop <- cropImage(img_resize, type = "equal_spaced", 
                          new_width = final_size, new_height = final_size)
    
    #Re-arrange array axes
    img_perm <- aperm(img_crop, c(3, 1, 2))
    
    #Expand array dimension
    dim(img_perm) <- c(1, dim(img_perm))
    
    #Prepare input images in Python dictionary
    image_dict <- dict(list('input' = np_array(img_perm, dtype="float32")))
    
    #Get model predictions
    predictions <- ort_sess$run(py_none(),  image_dict)
    predictions <- unlist(predictions)
    
    #Get predicted class index, probability, and label
    class_ind <- which.max(predictions)
    class_prob <- max(predictions)
    class_label <- classes[class_ind]
    
    #Render result as text
    output$image_label <- renderText(paste0("Predicted class: ", 
                                            class_label, 
                                            " (", 
                                            round(class_prob*100, digits = 1), 
                                            "%)"))
    
    #Write image to temporary file and return to UI display
    outfile <- tempfile(fileext='.png')
    
    writeImage(img_crop, outfile)
    
    output$image <- renderImage(list(src=outfile), deleteFile=TRUE)
    
  })
}

The Shiny application can be run locally using:

#Run the application 
shinyApp(ui = ui, server = server)

A simple Shiny application using the ONNX Python runtime to perform image classification.

I have compared the outputs obtained from the raw PyTorch model, ONNX model, and Shiny-ONNX runtime combination in a notebook. The outputs agree well and some minor differences should be expected due to differences in image preprocessing e.g. resampling.

Deployment notes

One thing to note is on setting up the Python virtual environment and making sure that reticulate use it, especially if making apps hosted online e.g. shinyapps.io. To make this easy, the appropriate environment variables can be defined in the ‘.Rprofile’ file, as described here.

VIRTUALENV_NAME = 'onnx_runtime'

if (Sys.info()[['user']] == 'shiny'){
  
  # Running on shinyapps.io
  Sys.setenv(PYTHON_PATH = 'python3')
  Sys.setenv(VIRTUALENV_NAME = VIRTUALENV_NAME) 
  Sys.setenv(RETICULATE_PYTHON = paste0('/home/shiny/.virtualenvs/', VIRTUALENV_NAME, '/bin/python'))
  
} else {
  
  # Running locally
  options(shiny.port = 7450)
  Sys.setenv(PYTHON_PATH = 'python')
  Sys.setenv(VIRTUALENV_NAME = VIRTUALENV_NAME)
  
}

Concluding remarks

R Shiny is a great framework for quickly building web applications, and with the reticulate R package, Python code can easily be injected for further functionality. Here, I have shown how an image classification model can be deployed within an R Shiny app which can be deployed online using any ONNX model.