Streaming RGB-D detection

This example expects a corresponding RGB and depth topic pair. It streams RGB frames through YOLOX, parses any detection’s bounding box and uses the olo.vision.sample_depth_region() helper to sample the distance to the object. Any detections are annotated on the frames and published to a /annotated topic.

Python

"""Continuously run RGB-D YOLOX detection and publish annotated frames."""

import queue
import threading
import time

import cv2

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

RGB_TOPIC = "wrist_rgbd/color_image"
DEPTH_TOPIC = "wrist_rgbd/depth_image"
ANNOTATED_TOPIC = "wrist_rgbd/annotated"
CONFIDENCE = 0.4
MAX_FPS = 10


def annotate(rgb: Image, depth: 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())
        try:
            depth_m = sample_depth_region(depth, detection.box)
            caption = f"{detection.label} {detection.score:.2f} {depth_m:.2f}m"
        except ValueError:
            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/Image")
    depth_topic = robot.core.topic(DEPTH_TOPIC, "sensor_msgs/msg/Image")
    annotated_topic = robot.core.topic(ANNOTATED_TOPIC, "sensor_msgs/msg/Image")
    publisher = robot.core.publisher(annotated_topic)
    rgb_frames = queue.Queue(maxsize=1)
    depth_frames = queue.Queue(maxsize=1)
    errors = queue.Queue()
    for topic, frames in ((rgb_topic, rgb_frames), (depth_topic, depth_frames)):
        threading.Thread(
            target=receive_latest,
            args=(robot, 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:
                if not errors.empty():
                    raise errors.get()

                now = time.monotonic()
                if now < next_frame:
                    time.sleep(min(next_frame - now, 0.01))
                    continue

                try:
                    rgb_message = rgb_frames.get(timeout=0.1)
                    depth_message = depth_frames.get(timeout=0.1)
                except queue.Empty:
                    continue
                next_frame = now + 1 / MAX_FPS

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

TypeScript

/** Continuously run RGB-D YOLOX detection and publish annotated frames. */

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

const RGB_TOPIC = "wrist_rgbd/color_image";
const DEPTH_TOPIC = "wrist_rgbd/depth_image";
const ANNOTATED_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, depth: 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();
    let caption = `${detection.label} ${detection.score.toFixed(2)}`;
    try {
      caption += ` ${sampleDepthRegion(depth, detection.box).toFixed(2)}m`;
    } catch {
      // Keep the label and score when no valid depth is available.
    }

    context.strokeRect(xMin, yMin, xMax - xMin, yMax - yMin);
    context.fillText(caption, 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 depthTopic = robot.core.topic(DEPTH_TOPIC, "sensor_msgs/msg/Image");
const annotatedTopic = robot.core.topic(ANNOTATED_TOPIC, "sensor_msgs/msg/Image");
const publisher = robot.core.publisher(annotatedTopic);
let latestRgb: Image | undefined;
let latestDepth: Image | undefined;
let subscriptionError: unknown;

const rgbSubscription = robot.core.subscribe(rgbTopic);
const depthSubscription = robot.core.subscribe(depthTopic);
void (async () => {
  try {
    for await (const message of rgbSubscription) {
      latestRgb = imageMessageToFrame(message);
    }
  } catch (error) {
    subscriptionError = error;
  }
})();
void (async () => {
  try {
    for await (const message of depthSubscription) {
      latestDepth = 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 || latestDepth === undefined || now < nextFrame) {
      await sleep(10);
      continue;
    }

    const rgb = latestRgb;
    const depth = latestDepth;
    latestRgb = undefined;
    latestDepth = 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, depth, detections) : rgb);
  }
} finally {
  rgbSubscription.close();
  depthSubscription.close();
  await model.unload();
}