Streaming RGB detection

Stream RGB frames from a camera topic through YOLOX. Any detections are annotated on the frames and published to a /annotated topic.

Python

"""Continuously run YOLOX on a camera stream."""

import queue
import threading
import time

import cv2

from olo import Client
from olo.ai.yolo import YoloDetector
from olo.core.images import Image, compressed_message_to_frame
# @olo-playground: autofill-sim-topics

BASE_CAMERA_TOPIC = "wrist_rgbd"
RGB_TOPIC = f"{BASE_CAMERA_TOPIC}/compressed_image"
OUTPUT_TOPIC = f"{BASE_CAMERA_TOPIC}/annotated"
CONFIDENCE = 0.4
MAX_FPS = 10


def annotate(rgb: Image, detections) -> Image:
    canvas = rgb.array.copy()

    for detection in detections:
        x_min, y_min, x_max, y_max = map(int, detection.box.to_xyxy())
        caption = f"{detection.label} {detection.score:.2f}"
        cv2.rectangle(canvas, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
        cv2.putText(
            canvas,
            caption,
            (x_min, max(y_min - 6, 14)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 255, 0),
            1,
            cv2.LINE_AA,
        )

    return Image(array=canvas, encoding="bgr8", header=rgb.header)


def receive_latest(robot, topic, frames, errors):
    try:
        with robot.core.subscribe(topic) as sub:
            for message in sub:
                try:
                    frames.get_nowait()
                except queue.Empty:
                    pass
                frames.put_nowait(message)
    except Exception as error:
        errors.put(error)


with Client() as client:
    robot = client.robot()
    rgb_topic = robot.core.topic(RGB_TOPIC, "sensor_msgs/msg/CompressedImage")
    output_topic = robot.core.topic(OUTPUT_TOPIC, "sensor_msgs/msg/Image")
    publisher = robot.core.publisher(output_topic)
    frames = queue.Queue(maxsize=1)
    errors = queue.Queue(maxsize=1)
    threading.Thread(
        target=receive_latest,
        args=(robot, rgb_topic, frames, errors),
        daemon=True,
    ).start()

    with client.ai.load_model("yolox_s.onnx") as model:
        detector = YoloDetector(model)
        print("Running continuously; press Ctrl+C to stop.")
        next_frame = time.monotonic()

        try:
            while True:
                try:
                    message = frames.get(timeout=0.1)
                except queue.Empty:
                    if not errors.empty():
                        raise errors.get()
                    continue

                now = time.monotonic()
                if now < next_frame:
                    continue
                next_frame = now + 1 / MAX_FPS

                rgb = compressed_message_to_frame(message)
                detections = []
                try:
                    detections = detector.detect(rgb, conf=CONFIDENCE)
                except Exception:
                    pass
                publisher.publish(annotate(rgb, detections) if detections else rgb)
        except KeyboardInterrupt:
            print("Stopping...")

TypeScript

/** Continuously run YOLOX on a camera stream. */

import { createCanvas, ImageData } from "@napi-rs/canvas";
import { YoloDetector } from "olo/ai";
import { imageMessageToFrame, type Image } from "olo/core";
import { connect } from "olo/web";
// @olo-playground: autofill-sim-topics

const RGB_TOPIC = "wrist_rgbd/color_image";
const OUTPUT_TOPIC = "wrist_rgbd/annotated";
const CONFIDENCE = 0.4;
const MAX_FPS = 10;

const sleep = (milliseconds: number): Promise<void> =>
  new Promise((resolve) => setTimeout(resolve, milliseconds));

function annotate(rgb: Image, detections: Awaited<ReturnType<YoloDetector["detect"]>>): Image {
  const { data, width, height, step } = rgb;
  const rgba = new Uint8ClampedArray(width * height * 4);
  const rgbOrder = rgb.encoding.toLowerCase() === "rgb8";

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const source = y * step + x * 3;
      const target = (y * width + x) * 4;
      rgba[target] = data[source + (rgbOrder ? 0 : 2)]!;
      rgba[target + 1] = data[source + 1]!;
      rgba[target + 2] = data[source + (rgbOrder ? 2 : 0)]!;
      rgba[target + 3] = 255;
    }
  }

  const canvas = createCanvas(width, height);
  const context = canvas.getContext("2d");
  context.putImageData(new ImageData(rgba, width, height), 0, 0);
  context.strokeStyle = "#00ff00";
  context.fillStyle = "#00ff00";
  context.lineWidth = 2;
  context.font = "14px sans-serif";

  for (const detection of detections) {
    const [xMin, yMin, xMax, yMax] = detection.box.toXyxy();
    context.strokeRect(xMin, yMin, xMax - xMin, yMax - yMin);
    context.fillText(
      `${detection.label} ${detection.score.toFixed(2)}`,
      xMin,
      Math.max(yMin - 4, 14),
    );
  }

  const annotatedRgba = context.getImageData(0, 0, width, height).data;
  // Preserve the source channel order because no-detection frames are
  // published unchanged on the same topic.
  const annotatedInSourceOrder = new Uint8Array(width * height * 3);
  for (let pixel = 0; pixel < width * height; pixel += 1) {
    annotatedInSourceOrder[pixel * 3] =
      annotatedRgba[pixel * 4 + (rgbOrder ? 0 : 2)]!;
    annotatedInSourceOrder[pixel * 3 + 1] = annotatedRgba[pixel * 4 + 1]!;
    annotatedInSourceOrder[pixel * 3 + 2] =
      annotatedRgba[pixel * 4 + (rgbOrder ? 2 : 0)]!;
  }

  return {
    data: annotatedInSourceOrder,
    // The byte order above matches the source, so do not hardcode "bgr8".
    encoding: rgb.encoding || "bgr8",
    format: "",
    width,
    height,
    step: width * 3,
    header: rgb.header,
  };
}

const client = connect();
const robot = await client.robot();
const rgbTopic = robot.core.topic(RGB_TOPIC, "sensor_msgs/msg/Image");
const outputTopic = robot.core.topic(OUTPUT_TOPIC, "sensor_msgs/msg/Image");
const publisher = robot.core.publisher(outputTopic);
let latestRgb: Image | undefined;
let subscriptionError: unknown;

const subscription = robot.core.subscribe(rgbTopic);
void (async () => {
  try {
    for await (const message of subscription) {
      latestRgb = imageMessageToFrame(message);
    }
  } catch (error) {
    subscriptionError = error;
  }
})();

const model = await client.ai.loadModel("yolox_s.onnx");
try {
  const detector = new YoloDetector(model);
  console.log("Running continuously; press Ctrl+C to stop.");
  let nextFrame = Date.now();

  while (true) {
    if (subscriptionError !== undefined) throw subscriptionError;

    const now = Date.now();
    if (latestRgb === undefined || now < nextFrame) {
      await sleep(10);
      continue;
    }

    const rgb = latestRgb;
    latestRgb = undefined;
    nextFrame = now + 1_000 / MAX_FPS;

    let detections: Awaited<ReturnType<YoloDetector["detect"]>> = [];
    try {
      detections = await detector.detect(rgb, { conf: CONFIDENCE });
    } catch {
      // Publish the unannotated frame if detection fails.
    }
    await publisher.publish(detections.length > 0 ? annotate(rgb, detections) : rgb);
  }
} finally {
  subscription.close();
  await model.unload();
}