YOLOv8 on a Nerves device

Setup

This notebook runs against a device already running the nerves_example firmware, with a Hailo accelerator and a camera attached. It drives the functions in NervesExample, so the code lives in the firmware and only the loop lives here.

Get the model onto the device first — see livebooks/download_models.livemd in the nx_hailo repository.

Then start distribution from the device's IEx prompt, over SSH:

cmd("epmd -daemon")
Node.start(:"nerves@<IP OF THE DEVICE>")
Node.set_cookie(:nerves_example_cookie)

Connect this notebook to it with Runtime → Attached Node, using that node name and cookie.

{:ok, {model, classes, input_name, output_name}} = NervesExample.load()
{capture, device} = NervesExample.find_capture()

input_shape = {trunc(capture.frame_height), trunc(capture.frame_width)}
padded_shape = {640, 640}

IO.puts("camera #{device}, frames of #{inspect(input_shape)}")

Inference loop

Each frame is padded to a square, resized to what the model takes, and run through the accelerator. Detections come back normalized against the padded square, so postprocess/2 maps them onto the camera frame before anything is drawn.

frame_interval_ms = 20
score_threshold = 0.5

Kino.animate(frame_interval_ms, fn _ ->
  frame = NervesExample.frame(capture)

  input_tensor =
    frame
    |> NervesExample.resize_and_pad(input_shape, padded_shape)
    |> Evision.Mat.to_nx()
    |> Nx.backend_transfer()

  {:ok, raw_detections} =
    NxHailo.infer(model, %{input_name => input_tensor}, NxHailo.Parsers.YoloV8,
      classes: classes,
      key: output_name
    )

  detections =
    raw_detections
    |> Enum.reject(&(&1.score < score_threshold))
    |> NxHailo.Parsers.YoloV8.postprocess(input_shape)

  NervesExample.YOLODraw.draw_detected_objects(
    frame,
    detections,
    "#{length(detections)} objects at #{div(1000, frame_interval_ms)} fps"
  )
end)