image classification

File.cd!(__DIR__)
# for windows JP
System.shell("chcp 65001")

Mix.install([
  {:onnx_interp, "~> 0.1.8"},
  {:cimg, "~> 0.1.14"},
  {:nx, "~> 0.4.0"},
  {:kino, "~> 0.7.0"}
])

Original work

• ResNet18:
  ONNX Model Zoo/ResNet - https://github.com/onnx/models/tree/main/vision/classification/resnet

• VGG16
  ONNX Model Zoo/VGG - https://github.com/onnx/models/tree/main/vision/classification/vgg

• ViT
  Vision Transformer (ViT) in PyTorch - https://github.com/lukemelas/PyTorch-Pretrained-ViT
  (convert this model to ONNX with torch.onnx.export())

Thanks a lot!!!

Implementation with OnnxInterp in Elixir

Please select a DNN model. {“resnet18”, “vgg16”, “vit”}

System.put_env("MODEL", "resnet18")

# select dnn model in {"vit", "vgg16", "resnet18"} at compile time.
{model_name, model_path, {model_width, model_height}, model_url} =
  System.get_env("MODEL", "resnet18")
  |> String.downcase()
  |> tap(&IO.puts("-- COMPLE the deme for \"#{&1}\"."))
  |> case do
    "vit" ->
      {"ViT", "./model/vit.onnx", {384, 384},
       "https://drive.google.com/uc??authuser=0&export=download&confirm=t&id=1L2sDNeK7CXdiMYflgCN8dWkrP2397ULv"}

    "vgg16" ->
      {"VGG16", "./model/vgg16-7.onnx", {224, 224},
       "https://github.com/onnx/models/raw/main/vision/classification/vgg/model/vgg16-7.onnx"}

    "resnet18" ->
      {"ResNet18", "./model/resnet18-v1-7.onnx", {224, 224},
       "https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet18-v1-7.onnx"}

    unknown ->
      raise(
        "** COMILE ERROR: unknon dnn model: '#{unknown}'. it must be one of {vit, vgg16, resnet18} **"
      )
  end

defmodule ImageClassify do
  @name model_name
  @model model_path
  @url model_url
  @width model_width
  @height model_height

  alias OnnxInterp, as: NNInterp
  use NNInterp, model: @model, url: @url

  @imagenet1000 (for item <- File.stream!("./model/imagenet1000.label") do
                   String.trim_trailing(item)
                 end)
                |> Enum.with_index(&amp;{&amp;2, &amp;1})
                |> Enum.into(%{})

  def apply(img, top \\ 1) do
    # preprocess
    input0 =
      img
      |> CImg.resize({@width, @height})
      |> CImg.to_binary([{:gauss, {{123.7, 58.4}, {116.3, 57.1}, {103.5, 57.4}}}, :nchw])

    # prediction
    output0 =
      session()
      |> NNInterp.set_input_tensor(0, input0)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)
      |> Nx.from_binary(:f32)
      |> Nx.reshape({1000})

    # postprocess
    # softmax
    Nx.exp(output0)
    |> then(fn exp -> Nx.divide(exp, Nx.sum(exp)) end)
    |> Nx.argsort(direction: :desc)
    |> Nx.slice([0], [top])
    |> Nx.to_flat_list()
    |> Enum.map(&amp;@imagenet1000[&amp;1])
  end

  def info() do
    %{name: @name, path: @model, shape: {@width, @height}, url: @url}
  end
end

Launch ImageClassify.

# OnnxInterp.stop(ImageClassify)
ImageClassify.start_link([])
ImageClassify.info()

Let’s try it

Load a photo and apply ImageClassify to it.

filename = "lion.jpg"

CImg.load(filename)
# display the image.
|> tap(fn img -> Kino.render(CImg.display_kino(img, :jpeg)) end)

# inference.
|> ImageClassify.apply()
|> tap(fn _ -> IO.puts("#{ImageClassify.info().name} answers:") end)
|> IO.inspect(label: "'#{filename}' is ")

:ok

TIL ;-)

Appendix

How to get pre-trained ViT onnx model.

You need install “pytorch_pretrained_vit” and covert the model beforehand.

> $ pip install pytorch_pretrained_vit

# convert Pytorch ViT model to ONNX.

import torch
from pytorch_pretrained_vit import ViT

model = ViT('B_16_imagenet1k', pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 384, 384)

os.makedirs("data", exist_ok=True)
torch.onnx.export(model,
   dummy_input,
   "vit.onnx",
   verbose=False,
   input_names=["input.0"],
   output_names=["output.0"],
   export_params=True
   )

□