Saltar a contenido

Plot Module

Plot

Class for plotting detections on images.

Source code in devices\raspberry_pi_5\src\plot\__init__.py
 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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
class Plot:
    """
    Class for plotting detections on images.
    """

    @staticmethod
    def draw_detection(
        image: np.ndarray,
        box: List,
        class_name: str,
        score: float,
        color: tuple,
        scale_factor: float
    ) -> None:
        """
        Draw box and label for one detection.

        Args:
            image (np.ndarray): Image to draw on.
            box (List): Bounding box coordinates.
            class_name (str): Class label of the detection.
            score (float): Detection score.
            color (tuple): Color for the bounding box.
            scale_factor (float): Scale factor for coordinates.
        """
        label = f"{class_name}: {score:.2f}%"
        ymin, xmin, ymax, xmax = box
        ymin, xmin, ymax, xmax = int(ymin * scale_factor), int(
            xmin * scale_factor
        ), int(ymax * scale_factor), int(
            xmax * scale_factor
        )
        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
        cv2.putText(
            image, label, (xmin + 4, ymin + 20), FONT, 0.5, color, 1,
            cv2.LINE_AA
        )

    @staticmethod
    def denormalize_and_remove_padding(
        box: List,
        size: int,
        padding_length: int,
        input_height: int,
        input_width: int
    ) -> List:
        """
        Denormalize bounding box coordinates and remove padding.

        Args:
            box (List): Normalized bounding box coordinates.
            size (int): Size to scale the coordinates.
            padding_length (int): Length of padding to remove.
            input_height (int): Height of the input image.
            input_width (int): Width of the input image.
        Returns:
            List: Denormalized bounding box coordinates with padding removed.
        """
        for i, x in enumerate(box):
            box[i] = int(x * size)
            if (input_width != size) and (i % 2 != 0):
                box[i] -= padding_length
            if (input_height != size) and (i % 2 == 0):
                box[i] -= padding_length

        return box

    @classmethod
    def draw_detections(
        cls,
        colors: Dict[int, tuple[int, int, int]],
        image_bounding_boxes: ImageBoundingBoxes,
        image: np.ndarray,
        min_score: float = 0.45,
        scale_factor: float = 1
    ):
        """
        Draw detections on the image.

        Args:
            colors (Dict[int, tuple[int, int, int]): Dictionary mapping class names to RGB colors.
            image_bounding_boxes (ImageBoundingBoxes): Object containing bounding boxes, classes, and scores.
            image (np.ndarray): Image to draw on.
            min_score (float): Minimum score threshold. Defaults to 0.45.
            scale_factor (float): Scale factor for coordinates. Defaults to 1.
        Returns:
            np.ndarray: Image with detections drawn.
        """
        # Values used for scaling coords and removing padding
        img_height, img_width = image.shape[:2]
        size = max(img_height, img_width)
        padding_length = int(abs(img_height - img_width) / 2)

        # Get the required values from the image bounding boxes
        boxes = image_bounding_boxes.get_xyxy()
        classes = image_bounding_boxes.get_classes()
        scores = image_bounding_boxes.get_confidences()
        for idx in range(image_bounding_boxes.get_number_of_objects()):
            if scores[idx] >= min_score:
                class_name = classes[idx]
                color = colors.get(idx, UNUSED_COLOR)
                scaled_box = cls.denormalize_and_remove_padding(
                    boxes[idx],
                    size,
                    padding_length,
                    img_height,
                    img_width
                )
                cls.draw_detection(
                    image,
                    scaled_box,
                    class_name,
                    scores[idx] * 100.0,
                    color,
                    scale_factor
                )

        return image

    @staticmethod
    def display_detections(
        class_names: Dict[int, str],
        preprocessed_image: List[np.ndarray],
        image_bounding_boxes: ImageBoundingBoxes,
        draw_labels_name=False,
        font=FONT,
        font_x_diff=0,
        font_y_diff=-10,
        font_scale=0.9,
        thickness=2,
        rgb_colors: tuple[tuple[int, int, int]] = None
    ) -> None:
        """
        Function to display the preprocessed image and the image with detections.

        Args:
            class_names (Dict[int, str]): Dictionary mapping class numbers to class names.
            preprocessed_image (List[np.ndarray]): Preprocessed image in CHW format.
            image_bounding_boxes (ImageBoundingBoxes): Object containing bounding boxes, classes, and scores.
            draw_labels_name (bool): Whether to draw class names instead of class numbers. Defaults to False.
            font (int): Font type for text. Defaults to FONT.
            font_x_diff (int): Horizontal offset for text placement. Defaults to 0.
            font_y_diff (int): Vertical offset for text placement. Defaults to -10.
            font_scale (float): Scale factor for text size. Defaults to 0.9.
            thickness (int): Thickness of bounding box and text. Defaults to 2.
            rgb_colors (tuple[tuple[int, int, int]]): Tuple mapping class indices to RGB colors. Defaults to None.
        Returns:
            None: Displays the preprocessed image and the image with detections using matplotlib.
        """
        # Convert the image back to HWC format
        preprocessed_image_hwc = np.transpose(preprocessed_image[0], (1, 2, 0))

        # Convert the image to uint8
        preprocessed_image_uint8 = (preprocessed_image_hwc * 255).astype(
            np.uint8
        )

        # Display the preprocessed image
        plt.subplot(1, 2, 1)
        plt.imshow(preprocessed_image_uint8)
        plt.title('Preprocessed Image')

        # Get the image with detections
        image_with_detections = preprocessed_image_uint8.copy()

        # Draw bounding boxes with class numbers
        xyxy = image_bounding_boxes.get_xyxy()
        class_numbers = image_bounding_boxes.get_classes()
        n = image_bounding_boxes.get_number_of_objects()

        if draw_labels_name:
            for i in range(n):
                x1, y1, x2, y2 = xyxy[i].astype(int)
                class_number = int(class_numbers[i])
                class_name = class_names[class_number]
                color = OpenCV.get_rgb_color(class_number, rgb_colors)
                cv2.rectangle(
                    image_with_detections, (x1, y1), (x2, y2), color,
                    thickness
                )
                cv2.putText(
                    image_with_detections, class_name,
                    (x1 + font_x_diff, y1 + font_y_diff), font,
                    font_scale,
                    color, thickness
                )

        else:
            for i in range(n):
                x1, y1, x2, y2 = xyxy[i].astype(int)
                class_number = int(class_numbers[i])
                color = OpenCV.get_rgb_color(class_number, rgb_colors)
                cv2.rectangle(
                    image_with_detections, (x1, y1), (x2, y2), color,
                    thickness
                )
                cv2.putText(
                    image_with_detections, str(class_number),
                    (x1 + font_x_diff, y1 + font_y_diff), font,
                    font_scale, color,
                    thickness
                )

        # Convert the image back to HWC format
        plt.subplot(1, 2, 2)
        plt.imshow(image_with_detections)
        plt.title('Image with Detections')
        plt.show()

denormalize_and_remove_padding(box, size, padding_length, input_height, input_width) staticmethod

Denormalize bounding box coordinates and remove padding.

Parameters:

Name Type Description Default
box List

Normalized bounding box coordinates.

required
size int

Size to scale the coordinates.

required
padding_length int

Length of padding to remove.

required
input_height int

Height of the input image.

required
input_width int

Width of the input image.

required

Returns:
List: Denormalized bounding box coordinates with padding removed.

Source code in devices\raspberry_pi_5\src\plot\__init__.py
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
@staticmethod
def denormalize_and_remove_padding(
    box: List,
    size: int,
    padding_length: int,
    input_height: int,
    input_width: int
) -> List:
    """
    Denormalize bounding box coordinates and remove padding.

    Args:
        box (List): Normalized bounding box coordinates.
        size (int): Size to scale the coordinates.
        padding_length (int): Length of padding to remove.
        input_height (int): Height of the input image.
        input_width (int): Width of the input image.
    Returns:
        List: Denormalized bounding box coordinates with padding removed.
    """
    for i, x in enumerate(box):
        box[i] = int(x * size)
        if (input_width != size) and (i % 2 != 0):
            box[i] -= padding_length
        if (input_height != size) and (i % 2 == 0):
            box[i] -= padding_length

    return box

display_detections(class_names, preprocessed_image, image_bounding_boxes, draw_labels_name=False, font=FONT, font_x_diff=0, font_y_diff=-10, font_scale=0.9, thickness=2, rgb_colors=None) staticmethod

Function to display the preprocessed image and the image with detections.

Parameters:

Name Type Description Default
class_names Dict[int, str]

Dictionary mapping class numbers to class names.

required
preprocessed_image List[ndarray]

Preprocessed image in CHW format.

required
image_bounding_boxes ImageBoundingBoxes

Object containing bounding boxes, classes, and scores.

required
draw_labels_name bool

Whether to draw class names instead of class numbers. Defaults to False.

False
font int

Font type for text. Defaults to FONT.

FONT
font_x_diff int

Horizontal offset for text placement. Defaults to 0.

0
font_y_diff int

Vertical offset for text placement. Defaults to -10.

-10
font_scale float

Scale factor for text size. Defaults to 0.9.

0.9
thickness int

Thickness of bounding box and text. Defaults to 2.

2
rgb_colors tuple[tuple[int, int, int]]

Tuple mapping class indices to RGB colors. Defaults to None.

None

Returns:
None: Displays the preprocessed image and the image with detections using matplotlib.

Source code in devices\raspberry_pi_5\src\plot\__init__.py
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
@staticmethod
def display_detections(
    class_names: Dict[int, str],
    preprocessed_image: List[np.ndarray],
    image_bounding_boxes: ImageBoundingBoxes,
    draw_labels_name=False,
    font=FONT,
    font_x_diff=0,
    font_y_diff=-10,
    font_scale=0.9,
    thickness=2,
    rgb_colors: tuple[tuple[int, int, int]] = None
) -> None:
    """
    Function to display the preprocessed image and the image with detections.

    Args:
        class_names (Dict[int, str]): Dictionary mapping class numbers to class names.
        preprocessed_image (List[np.ndarray]): Preprocessed image in CHW format.
        image_bounding_boxes (ImageBoundingBoxes): Object containing bounding boxes, classes, and scores.
        draw_labels_name (bool): Whether to draw class names instead of class numbers. Defaults to False.
        font (int): Font type for text. Defaults to FONT.
        font_x_diff (int): Horizontal offset for text placement. Defaults to 0.
        font_y_diff (int): Vertical offset for text placement. Defaults to -10.
        font_scale (float): Scale factor for text size. Defaults to 0.9.
        thickness (int): Thickness of bounding box and text. Defaults to 2.
        rgb_colors (tuple[tuple[int, int, int]]): Tuple mapping class indices to RGB colors. Defaults to None.
    Returns:
        None: Displays the preprocessed image and the image with detections using matplotlib.
    """
    # Convert the image back to HWC format
    preprocessed_image_hwc = np.transpose(preprocessed_image[0], (1, 2, 0))

    # Convert the image to uint8
    preprocessed_image_uint8 = (preprocessed_image_hwc * 255).astype(
        np.uint8
    )

    # Display the preprocessed image
    plt.subplot(1, 2, 1)
    plt.imshow(preprocessed_image_uint8)
    plt.title('Preprocessed Image')

    # Get the image with detections
    image_with_detections = preprocessed_image_uint8.copy()

    # Draw bounding boxes with class numbers
    xyxy = image_bounding_boxes.get_xyxy()
    class_numbers = image_bounding_boxes.get_classes()
    n = image_bounding_boxes.get_number_of_objects()

    if draw_labels_name:
        for i in range(n):
            x1, y1, x2, y2 = xyxy[i].astype(int)
            class_number = int(class_numbers[i])
            class_name = class_names[class_number]
            color = OpenCV.get_rgb_color(class_number, rgb_colors)
            cv2.rectangle(
                image_with_detections, (x1, y1), (x2, y2), color,
                thickness
            )
            cv2.putText(
                image_with_detections, class_name,
                (x1 + font_x_diff, y1 + font_y_diff), font,
                font_scale,
                color, thickness
            )

    else:
        for i in range(n):
            x1, y1, x2, y2 = xyxy[i].astype(int)
            class_number = int(class_numbers[i])
            color = OpenCV.get_rgb_color(class_number, rgb_colors)
            cv2.rectangle(
                image_with_detections, (x1, y1), (x2, y2), color,
                thickness
            )
            cv2.putText(
                image_with_detections, str(class_number),
                (x1 + font_x_diff, y1 + font_y_diff), font,
                font_scale, color,
                thickness
            )

    # Convert the image back to HWC format
    plt.subplot(1, 2, 2)
    plt.imshow(image_with_detections)
    plt.title('Image with Detections')
    plt.show()

draw_detection(image, box, class_name, score, color, scale_factor) staticmethod

Draw box and label for one detection.

Parameters:

Name Type Description Default
image ndarray

Image to draw on.

required
box List

Bounding box coordinates.

required
class_name str

Class label of the detection.

required
score float

Detection score.

required
color tuple

Color for the bounding box.

required
scale_factor float

Scale factor for coordinates.

required
Source code in devices\raspberry_pi_5\src\plot\__init__.py
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
@staticmethod
def draw_detection(
    image: np.ndarray,
    box: List,
    class_name: str,
    score: float,
    color: tuple,
    scale_factor: float
) -> None:
    """
    Draw box and label for one detection.

    Args:
        image (np.ndarray): Image to draw on.
        box (List): Bounding box coordinates.
        class_name (str): Class label of the detection.
        score (float): Detection score.
        color (tuple): Color for the bounding box.
        scale_factor (float): Scale factor for coordinates.
    """
    label = f"{class_name}: {score:.2f}%"
    ymin, xmin, ymax, xmax = box
    ymin, xmin, ymax, xmax = int(ymin * scale_factor), int(
        xmin * scale_factor
    ), int(ymax * scale_factor), int(
        xmax * scale_factor
    )
    cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
    cv2.putText(
        image, label, (xmin + 4, ymin + 20), FONT, 0.5, color, 1,
        cv2.LINE_AA
    )

draw_detections(colors, image_bounding_boxes, image, min_score=0.45, scale_factor=1) classmethod

Draw detections on the image.

Parameters:

Name Type Description Default
colors Dict[int, tuple[int, int, int]

Dictionary mapping class names to RGB colors.

required
image_bounding_boxes ImageBoundingBoxes

Object containing bounding boxes, classes, and scores.

required
image ndarray

Image to draw on.

required
min_score float

Minimum score threshold. Defaults to 0.45.

0.45
scale_factor float

Scale factor for coordinates. Defaults to 1.

1

Returns:
np.ndarray: Image with detections drawn.

Source code in devices\raspberry_pi_5\src\plot\__init__.py
 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
119
120
121
122
123
124
125
126
127
128
129
@classmethod
def draw_detections(
    cls,
    colors: Dict[int, tuple[int, int, int]],
    image_bounding_boxes: ImageBoundingBoxes,
    image: np.ndarray,
    min_score: float = 0.45,
    scale_factor: float = 1
):
    """
    Draw detections on the image.

    Args:
        colors (Dict[int, tuple[int, int, int]): Dictionary mapping class names to RGB colors.
        image_bounding_boxes (ImageBoundingBoxes): Object containing bounding boxes, classes, and scores.
        image (np.ndarray): Image to draw on.
        min_score (float): Minimum score threshold. Defaults to 0.45.
        scale_factor (float): Scale factor for coordinates. Defaults to 1.
    Returns:
        np.ndarray: Image with detections drawn.
    """
    # Values used for scaling coords and removing padding
    img_height, img_width = image.shape[:2]
    size = max(img_height, img_width)
    padding_length = int(abs(img_height - img_width) / 2)

    # Get the required values from the image bounding boxes
    boxes = image_bounding_boxes.get_xyxy()
    classes = image_bounding_boxes.get_classes()
    scores = image_bounding_boxes.get_confidences()
    for idx in range(image_bounding_boxes.get_number_of_objects()):
        if scores[idx] >= min_score:
            class_name = classes[idx]
            color = colors.get(idx, UNUSED_COLOR)
            scaled_box = cls.denormalize_and_remove_padding(
                boxes[idx],
                size,
                padding_length,
                img_height,
                img_width
            )
            cls.draw_detection(
                image,
                scaled_box,
                class_name,
                scores[idx] * 100.0,
                color,
                scale_factor
            )

    return image