dora/examples/python-dataflow/plot.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
from enum import Enum
from typing import Callable
from dora import Node

import cv2
import numpy as np

from utils import LABELS

CI = os.environ.get("CI")

font = cv2.FONT_HERSHEY_SIMPLEX

class Status(Enum):
    CONTINUE = 0
    STOP = 1


class Plotter:
    """
    Plot image and bounding box
    """

    def __init__(self):
        self.image = []
        self.bboxs = []

    def on_input(
        self,
        dora_input: dict,
    ) -> Status:
        """
        Put image and bounding box on cv2 window.

        Args:
            dora_input["id"] (str): Id of the dora_input declared in the yaml configuration
            dora_input["data"] (bytes): Bytes message of the dora_input
        """
        if dora_input["id"] == "image":
            frame = np.frombuffer(dora_input["data"], dtype="uint8")
            frame = cv2.imdecode(frame, -1)
            self.image = frame

        elif dora_input["id"] == "bbox" and len(self.image) != 0:
            bboxs = np.frombuffer(dora_input["data"], dtype="float32")
            self.bboxs = np.reshape(bboxs, (-1, 6))
        for bbox in self.bboxs:
            [
                min_x,
                min_y,
                max_x,
                max_y,
                confidence,
                label,
            ] = bbox
            cv2.rectangle(
                self.image,
                (int(min_x), int(min_y)),
                (int(max_x), int(max_y)),
                (0, 255, 0),
                2,
            )

            cv2.putText(
                self.image,
                LABELS[int(label)] + f", {confidence:0.2f}",
                (int(max_x), int(max_y)),
                font,
                0.75,
                (0, 255, 0),
                2,
                1,
            )

        if CI != "true":
            cv2.imshow("frame", self.image)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                return Status.STOP

        return Status.CONTINUE

    def __del__(self):
        cv2.destroyAllWindows()


plotter = Plotter()
node = Node()

for event in node:
    match event["type"]:
        case "input":
            status = plotter.on_input(event)
            match status:
                case Status.CONTINUE:
                    pass
                case Status.STOP:
                    print("plotter returned stop status")
                    break
        case "stop":
            print("received stop")
        case other:
            print("received unexpected event:", other)