adora/examples/python-operator-dataflow/plot.py at main · dora-rs/adora · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""TODO: Add docstring."""

import os

import cv2
from adora import AdoraStatus
from utils import LABELS

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

CAMERA_WIDTH = 640
CAMERA_HEIGHT = 480

FONT = cv2.FONT_HERSHEY_SIMPLEX


class Operator:
    """Plot image and bounding box."""

    def __init__(self):
        """TODO: Add docstring."""
        self.bboxs = []
        self.buffer = ""
        self.submitted = []
        self.lines = []

    def on_event(
        self,
        adora_event,
        send_output,
    ):
        """TODO: Add docstring."""
        if adora_event["type"] == "INPUT":
            id = adora_event["id"]
            value = adora_event["value"]
            if id == "image":

                image = (
                    value.to_numpy().reshape((CAMERA_HEIGHT, CAMERA_WIDTH, 3)).copy()
                )

                for bbox in self.bboxs:
                    [
                        min_x,
                        min_y,
                        max_x,
                        max_y,
                        confidence,
                        label,
                    ] = bbox
                    cv2.rectangle(
                        image,
                        (int(min_x), int(min_y)),
                        (int(max_x), int(max_y)),
                        (0, 255, 0),
                    )
                    cv2.putText(
                        image,
                        f"{LABELS[int(label)]}, {confidence:0.2f}",
                        (int(max_x), int(max_y)),
                        FONT,
                        0.5,
                        (0, 255, 0),
                    )

                cv2.putText(
                    image, self.buffer, (20, 14 + 21 * 14), FONT, 0.5, (190, 250, 0), 1,
                )

                i = 0
                for text in self.submitted[::-1]:
                    color = (
                        (0, 255, 190)
                        if text["role"] == "user_message"
                        else (0, 190, 255)
                    )
                    cv2.putText(
                        image,
                        text["content"],
                        (
                            20,
                            14 + (19 - i) * 14,
                        ),
                        FONT,
                        0.5,
                        color,
                        1,
                    )
                    i += 1

                for line in self.lines:
                    cv2.line(
                        image,
                        (int(line[0]), int(line[1])),
                        (int(line[2]), int(line[3])),
                        (0, 0, 255),
                        2,
                    )

                if CI != "true":
                    cv2.imshow("frame", image)
                    if cv2.waitKey(1) & 0xFF == ord("q"):
                        return AdoraStatus.STOP
            elif id == "bbox":
                self.bboxs = value.to_numpy().reshape((-1, 6))
            elif id == "keyboard_buffer":
                self.buffer = value[0].as_py()
            elif id == "line":
                self.lines += [value.to_pylist()]
            elif "message" in id:
                self.submitted += [
                    {
                        "role": id,
                        "content": value[0].as_py(),
                    },
                ]

        return AdoraStatus.CONTINUE