Current section

Files

Jump to
nerves_livebook priv samples tflite.livemd
Raw

priv/samples/tflite.livemd

# Use TensorFlow Lite on Nerves Livebook
```elixir
Mix.install([
{:tflite_elixir, "~> 0.3.0"},
{:req, "~> 0.3.0"},
{:progress_bar, "~> 2.0.0"},
{:kino, "~> 0.9.0"}
])
```
## Introduction
TensorFlow Lite is a stripped-down version of
[TensorFlow](https://en.wikipedia.org/wiki/TensorFlow), a free and open-source
software library for machine learning and artificial intelligence.
In Elixir, we can use TensorFlow Lite through the
[`tflite_elixir`](https://github.com/cocoa-xu/tflite_elixir) package, which
does the TensorFlow Lite Elixir bindings with optional [Edge
TPU](https://en.wikipedia.org/wiki/Tensor_Processing_Unit#Edge_TPU) support.
In this notebook, we will perform image classification with pre-trained
[mobilenet_v2_1.0_224_inat_bird_quant.tflite](https://github.com/google-coral/edgetpu/blob/master/test_data/mobilenet_v2_1.0_224_inat_bird_quant.tflite)
model. The example code below is based on the instructions in the
[tflite_elixir
README](https://github.com/cocoa-xu/tflite_elixir/blob/main/README.md). For
more information, check out the [tflite_elixir API
reference](https://hexdocs.pm/tflite_elixir/api-reference.html).
## Prepare helper functions
```elixir
defmodule Utils do
def download!(source_url, req_options \\ []) do
Req.get!(source_url, [finch_request: &finch_request/4] ++ req_options).body
end
defp finch_request(req_request, finch_request, finch_name, finch_options) do
acc = Req.Response.new()
case Finch.stream(finch_request, finch_name, acc, &handle_message/2, finch_options) do
{:ok, response} -> {req_request, response}
{:error, exception} -> {req_request, exception}
end
end
defp handle_message({:status, status}, response), do: %{response | status: status}
defp handle_message({:headers, headers}, response) do
{_, total_size} = Enum.find(headers, &match?({"content-length", _}, &1))
response
|> Map.put(:headers, headers)
|> Map.put(:private, %{total_size: String.to_integer(total_size), downloaded_size: 0})
end
defp handle_message({:data, data}, response) do
new_downloaded_size = response.private.downloaded_size + byte_size(data)
ProgressBar.render(new_downloaded_size, response.private.total_size, suffix: :bytes)
response
|> Map.update!(:body, &(&1 <> data))
|> Map.update!(:private, &%{&1 | downloaded_size: new_downloaded_size})
end
end
```
## Decide on where downloaded files are saved
```elixir
downloads_dir = "/data/tmp"
File.mkdir_p!(downloads_dir)
```
## Download pre-trained model
```elixir
model_source =
"https://raw.githubusercontent.com/google-coral/test_data/master/mobilenet_v2_1.0_224_inat_bird_quant.tflite"
model_file = Path.join(downloads_dir, "mobilenet_v2_1.0_224_inat_bird_quant.tflite")
Utils.download!(model_source, output: model_file)
IO.puts("Model saved to #{model_file}")
```
## Download labels
```elixir
label_source =
"https://raw.githubusercontent.com/google-coral/test_data/master/inat_bird_labels.txt"
labels = String.split(Utils.download!(label_source), "\n", trim: true)
Kino.DataTable.new(Enum.with_index(labels, &%{class_name: &1, class_id: &2}), name: "Labels")
```
## Choose image to be classified
An input image can be uploaded here, or default parrot image will be used.
```elixir
image_input = Kino.Input.image("Image", size: {224, 224})
```
```elixir
uploaded_image = Kino.Input.read(image_input)
default_input_image_url =
"https://raw.githubusercontent.com/google-coral/test_data/master/parrot.jpg"
input_image =
if uploaded_image do
# Build a tensor from the raw pixel data
uploaded_image.data
|> Nx.from_binary(:u8)
|> Nx.reshape({uploaded_image.height, uploaded_image.width, 3})
else
IO.puts("Loading default image from #{default_input_image_url}")
Utils.download!(default_input_image_url)
|> StbImage.read_binary!()
|> StbImage.to_nx()
end
Kino.Image.new(input_image)
```
## Classify image
```elixir
how_many_results = 3
labels = List.to_tuple(labels)
input_nx =
input_image
|> StbImage.from_nx()
|> StbImage.resize(224, 224)
|> StbImage.to_nx()
interpreter = TFLiteElixir.Interpreter.new!(model_file)
[output_tensor_0] = TFLiteElixir.Interpreter.predict(interpreter, input_nx[[.., .., 0..2]])
indices_nx = Nx.flatten(output_tensor_0)
class_ids =
indices_nx
|> Nx.argsort(direction: :desc)
|> Nx.take(Nx.iota({how_many_results}))
|> Nx.to_flat_list()
class_ids
|> Enum.map(fn class_id -> %{class_id: class_id, class_name: elem(labels, class_id)} end)
|> Kino.DataTable.new(name: "Inference results")
```
## Next steps
### Run other models
You can find a variety of pre-trained open-source models in [TensorFlow
Hub](https://tfhub.dev). For Elixir code, check out [example
notebooks](https://github.com/cocoa-xu/tflite_elixir/blob/bda47628e143c860e8cc796f491edd49260b787b/notebooks/README.md)
in `tflite_elixir` repository.
In case some example notebooks require the
[`evision`](https://github.com/cocoa-xu/evision) package for using
[OpenCV](https://opencv.org), add it to your Nerves project's `mix.exs` file
and rebuild Nerves firmware.
### Run inference on Edge TPU
You can speed up model inference time, running a TensorFlow Lite model on the
Edge TPU. Check out `tflite_elixir`'s ["Inference on TPU"
example](https://github.com/cocoa-xu/tflite_elixir/blob/main/notebooks/tpu.livemd).