add text label to plot

examples/python-dataflow/plot.py

@@ -5,8 +5,12 @@ from typing import Callable
 import cv2
 import numpy as np
 
+from utils import LABELS
+
 CI = os.environ.get("CI")
 
+font = cv2.FONT_HERSHEY_SIMPLEX
+
 
 class DoraStatus(Enum):
     CONTINUE = 0
@@ -49,8 +53,8 @@ class Operator:
                     min_y,
                     max_x,
                     max_y,
-                    _confidence,
-                    _class_label,
+                    confidence,
+                    label,
                 ] = bbox
                 cv2.rectangle(
                     self.image,
@@ -59,6 +63,18 @@ class Operator:
                     (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"):

examples/python-dataflow/utils.py

@@ -0,0 +1,82 @@
+LABELS = [
+    "ABC",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "backpack",
+    "umbrella",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "dining table",
+    "toilet",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+]