Skip to content

DetectionProcessor class

DetectionProcessor

A class for processing detection results and transforming them using available utility methods. This class takes the output from an object detection model (e.g., YOLOv8) and provides functionality to separate, process, and visualize detected color checker cards and their patches.

Source code in color_correction/processor/detection.py
 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
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
class DetectionProcessor:
    """
    A class for processing detection results and transforming them using
    available utility methods. This class takes the output from an object
    detection model (e.g., YOLOv8) and provides functionality to separate,
    process, and visualize detected color checker cards and their patches.
    """

    @staticmethod
    def get_each_class_box(
        prediction: DetectionResult,
    ) -> tuple[list[BoundingBox | SegmentPoint], list[BoundingBox | SegmentPoint]]:
        """
        Separates detection boxes into card boxes and patch boxes.

        Parameters
        ----------
        prediction : DetectionResult
            The detection output that includes bounding boxes and class IDs.

        Returns
        -------
        tuple[list[BoundingBox | SegmentPoint], list[BoundingBox | SegmentPoint]]
            A tuple containing:

            - a list of card boxes (class id 1)
            - a list of patch boxes (class id 0).
        """

        if prediction.boxes is not None:
            ls_cards = [
                box
                for box, class_id in zip(
                    prediction.boxes,
                    prediction.class_ids,
                    strict=False,
                )
                if class_id == 1
            ]
            ls_patches = [
                box
                for box, class_id in zip(
                    prediction.boxes,
                    prediction.class_ids,
                    strict=False,
                )
                if class_id == 0
            ]
        if prediction.segment is not None:
            ls_cards = [
                segment
                for segment, class_id in zip(
                    prediction.segment,
                    prediction.class_ids,
                    strict=False,
                )
                if class_id == 1
            ]
            ls_patches = [
                segment
                for segment, class_id in zip(
                    prediction.segment,
                    prediction.class_ids,
                    strict=False,
                )
                if class_id == 0
            ]
        return ls_cards, ls_patches

    @staticmethod
    def print_summary(prediction: DetectionResult) -> None:
        """
        Prints a summary of the detected cards and patches.

        Parameters
        ----------
        prediction : DetectionResult
            The detection result to summarize.
        """
        ls_cards, ls_patches = DetectionProcessor.get_each_class_box(prediction)
        print(f"Number of cards detected: {len(ls_cards)}")
        print(f"Number of patches detected: {len(ls_patches)}")

    @staticmethod
    def process_patches(
        input_image: ImageBGR,
        ordered_patches: list[tuple[BoundingBox, tuple[int, int]] | None],
    ) -> tuple[list[ColorPatchType], ImageBGR]:
        """
        Processes each detected patch by extracting its region from the image,
        computing its mean BGR color, and building a visualization grid.

        Parameters
        ----------
        input_image : ImageBGR
            The original image containing detected patches.
        ordered_patches : list[tuple[BoundingBox, tuple[int, int]] | None]
            The list of ordered patch coordinates paired with their center
            points, possibly with missing entries as None.

        Returns
        -------
        tuple[list[ColorPatchType], ImageBGR]
            a tuple containing:

            - list of mean BGR color values for each patch.
            - an image visualizing these patches in a grid layout.
        """
        patch_size = (50, 50, 1)
        ls_bgr_mean_patch = []
        ls_horizontal_patch = []
        ls_vertical_patch = []

        for idx, coord_patch in enumerate(ordered_patches, start=1):
            if coord_patch is None:
                continue

            bbox_patch, _ = coord_patch

            # Extract and process each patch
            cropped_patch = crop_region_with_margin(
                image=input_image,
                coordinates=bbox_patch,
                margin_ratio=0.2,
            )
            bgr_mean_patch = calc_mean_color_patch(cropped_patch)
            ls_bgr_mean_patch.append(bgr_mean_patch)

            # Build visualization
            patch_viz = np.tile(bgr_mean_patch, patch_size)
            ls_horizontal_patch.append(patch_viz)
            if idx % 6 == 0:
                ls_vertical_patch.append(np.hstack(ls_horizontal_patch))
                ls_horizontal_patch = []

        patches_image = np.vstack(ls_vertical_patch)
        return ls_bgr_mean_patch, patches_image

    @staticmethod
    def extract_color_patches(
        input_image: ImageBGR,
        prediction: DetectionResult,
        draw_processed_image: bool = False,
    ) -> tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]:
        """
        Extracts and processes color patches from detected color checker cards,
        transforming the detection results with available methods.

        The method first separates detected cards and patches, generates an expected
        patch grid, and then matches the detected patches with this grid. If patches
        are missing, it attempts to auto-fill them with suggested coordinates.
        Finally, it computes the mean color for each patch and builds a visualization.

        Parameters
        ----------
        input_image : ImageBGR
            The original image containing the color checker card.
        prediction : DetectionResult
            The detection result output from the model.
        draw_processed_image : bool, optional
            If True, returns an additional image with visualized detections.
            Otherwise, only patch processing is performed. Defaults to False.

        Returns
        -------
        tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]
            a tuple containing:

            - a list of mean BGR color values for each patch.
            - an image visualizing these patches in a grid layout.
            - an optional image with visualized detection results.

        Raises
        ------
        ValueError
            If no cards or patches are detected.
        """
        ls_cards, ls_patches = DetectionProcessor.get_each_class_box(prediction)

        if not ls_cards:
            raise ValueError("No cards detected")
        if not ls_patches:
            raise ValueError("No patches detected")

        ls_bgr_mean_patch, grid_patch_img, detection_viz = None, None, None
        if prediction.segment is not None:
            print("Using segmentation results")
            ls_bgr_mean_patch, grid_patch_img, detection_viz = DetectionProcessor.extract_patches_from_segment(
                input_image,
                prediction,
                draw_processed_image,
                ls_cards,
                ls_patches,
            )
            print(f"grid_patch_img shape: {grid_patch_img}")

        if prediction.boxes is not None:
            print("Using box results")
            ls_bgr_mean_patch, grid_patch_img, detection_viz = DetectionProcessor.extract_patches_from_boxes(
                input_image,
                prediction,
                draw_processed_image,
                ls_cards,
                ls_patches,
            )
        return ls_bgr_mean_patch, grid_patch_img, detection_viz

    @staticmethod
    def extract_patches_from_segment(
        input_image: ImageBGR,
        prediction: DetectionResult,  # noqa: ARG004
        draw_processed_image: bool,
        ls_cards: list[SegmentPoint],
        ls_patches: list[SegmentPoint],
    ) -> tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]:
        patch_size = (50, 50, 1)
        ls_bgr_mean_segment = []
        ls_horizontal_patch = []
        ls_vertical_patch = []

        # we assume that white patch is bottom left, before this we need to arrange the patches
        for idx, segment in enumerate(ls_patches, start=1):
            cropped_segment = crop_segment_straighten(image=input_image, segment=segment)
            bgr_mean_segment = calc_mean_color_patch(cropped_segment)
            ls_bgr_mean_segment.append(bgr_mean_segment)

            # Build visualization
            patch_viz = np.tile(bgr_mean_segment, patch_size)  # Updated to use bgr_mean_segment
            ls_horizontal_patch.append(patch_viz)
            if idx % 6 == 0:
                ls_vertical_patch.append(np.hstack(ls_horizontal_patch))
                ls_horizontal_patch = []

        image_det_viz = None
        grid_patch_img = None

        if draw_processed_image:
            image_det_viz = input_image.copy()
            # draw cards
            for idx, card in enumerate(ls_cards):
                random_color = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
                # draw card
                cv2.polylines(image_det_viz, [np.int32(card)], True, random_color, 2)
                # draw text
                cv2.putText(
                    image_det_viz,
                    str(idx) + " card",
                    (int(card[0][0]), int(card[0][1]) - 4),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1,
                    random_color,
                    2,
                )

            # draw patches
            for idx, group_patch in enumerate(ls_patches):
                random_color = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
                # draw patch
                cv2.polylines(image_det_viz, [np.int32(group_patch)], True, random_color, 2)
                # draw text
                cv2.putText(
                    image_det_viz,
                    str(idx),
                    (int(group_patch[0][0]), int(group_patch[0][1]) - 4),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    1,
                    random_color,
                    2,
                )

            grid_patch_img = np.vstack(ls_vertical_patch)
        return ls_bgr_mean_segment, grid_patch_img, image_det_viz

    @staticmethod
    def extract_patches_from_boxes(
        input_image: np.ndarray,
        prediction: DetectionResult,
        draw_processed_image: bool,
        ls_cards: list[BoundingBox],
        ls_patches: list[BoundingBox],
    ) -> tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]:
        card_box = ls_cards[0]
        ls_grid_card = generate_expected_patches(card_box)

        # Match detected patches with grid
        ls_ordered_patch_bbox = extract_intersecting_patches(
            ls_patches=ls_patches,
            ls_grid_card=ls_grid_card,
        )

        # Handle missing patches
        d_suggest = None
        if None in ls_ordered_patch_bbox:
            print("Auto filling missing patches...")
            ls_ordered_bbox_only = [patch[0] if patch is not None else None for patch in ls_ordered_patch_bbox]
            d_suggest = suggest_missing_patch_coordinates(ls_ordered_bbox_only)
            for idx, patch in d_suggest.items():
                cxpatch = (patch[0] + patch[2]) // 2
                cypatch = (patch[1] + patch[3]) // 2
                ls_ordered_patch_bbox[idx] = (patch, (cxpatch, cypatch))

        # Process patches and create visualizations
        ls_bgr_mean_patch, grid_patch_img = DetectionProcessor.process_patches(
            input_image=input_image,
            ordered_patches=ls_ordered_patch_bbox,
        )

        detection_viz = None
        if draw_processed_image:
            detection_viz = DetectionProcessor.draw_preprocess(
                image=input_image,
                expected_boxes=ls_grid_card,
                prediction=prediction,
                ls_ordered_patch_bbox=ls_ordered_patch_bbox,
                suggested_patches=d_suggest,
            )

        return ls_bgr_mean_patch, grid_patch_img, detection_viz

    @staticmethod
    def draw_preprocess(
        image: ImageBGR,
        expected_boxes: list[BoundingBox],
        prediction: DetectionResult,
        ls_ordered_patch_bbox: list[BoundingBox | None],
        suggested_patches: dict[int, BoundingBox] | None = None,
    ) -> ImageBGR:
        """
        Draws visualizations on the input image to compare the detected patches with
        the expected positions. It overlays expected boxes, connects detected patches to
        their corresponding expected boxes, and highlights individual patch detections.

        Parameters
        ----------
        image : ImageBGR
            The original image to draw on.
        expected_boxes : list[BoundingBox]
            List of expected bounding boxes for patches.
        prediction : DetectionResult
            The detection results containing predicted bounding boxes.
        ls_ordered_patch_bbox : list[BoundingBox | None]
            The list of ordered patch bounding boxes.
        suggested_patches : dict[int, BoundingBox], optional
            Dictionary of suggested patch coordinates for missing patches.

        Returns
        -------
        ImageBGR
            The image with drawn detection and patch visualizations.
        """
        color_green = (0, 255, 0)
        color_cyan = (255, 255, 10)
        color_violet = (255, 0, 255)
        color_red = (0, 0, 255)
        color_blue = (255, 0, 0)

        result_image = image.copy()

        # Draw all expected boxes
        for idx_b, box in enumerate(expected_boxes):
            cv2.rectangle(
                img=result_image,
                pt1=(box[0], box[1]),
                pt2=(box[2], box[3]),
                color=color_green,
                thickness=2,
            )

            # Draw connection lines between expected and detected patch boxes
            patch = ls_ordered_patch_bbox[idx_b]
            if patch is None:
                continue
            cx, cy = patch[1]
            crefx, crefy = (box[0] + box[2]) // 2, (box[1] + box[3]) // 2
            cv2.line(
                img=result_image,
                pt1=(cx, cy),
                pt2=(crefx, crefy),
                color=color_blue,
                thickness=1,
            )

        # Draw all predicted boxes
        for pbox, pids, pscore in zip(
            prediction.boxes,
            prediction.class_ids,
            prediction.scores,
            strict=False,
        ):
            if pids == 1:
                continue
            cv2.rectangle(
                img=result_image,
                pt1=(pbox[0], pbox[1]),
                pt2=(pbox[2], pbox[3]),
                color=color_cyan,
                thickness=2,
            )
            cv2.putText(
                img=result_image,
                text=f"{pids} {pscore:.2f}",
                org=(pbox[0] + 3, pbox[1] + 12),
                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                fontScale=0.4,
                color=color_red,
                thickness=1,
                lineType=cv2.LINE_AA,
            )

        # Draw suggested patches if provided
        if suggested_patches:
            for box in suggested_patches.values():
                cv2.rectangle(
                    img=result_image,
                    pt1=(box[0], box[1]),
                    pt2=(box[2], box[3]),
                    color=color_violet,
                    thickness=2,
                )

        return result_image

Functions

get_each_class_box staticmethod

get_each_class_box(prediction: DetectionResult) -> tuple[list[BoundingBox | SegmentPoint], list[BoundingBox | SegmentPoint]]

Separates detection boxes into card boxes and patch boxes.

Parameters:

Name Type Description Default
prediction DetectionResult

The detection output that includes bounding boxes and class IDs.

required

Returns:

Type Description
tuple[list[BoundingBox | SegmentPoint], list[BoundingBox | SegmentPoint]]

A tuple containing:

  • a list of card boxes (class id 1)
  • a list of patch boxes (class id 0).
Source code in color_correction/processor/detection.py
@staticmethod
def get_each_class_box(
    prediction: DetectionResult,
) -> tuple[list[BoundingBox | SegmentPoint], list[BoundingBox | SegmentPoint]]:
    """
    Separates detection boxes into card boxes and patch boxes.

    Parameters
    ----------
    prediction : DetectionResult
        The detection output that includes bounding boxes and class IDs.

    Returns
    -------
    tuple[list[BoundingBox | SegmentPoint], list[BoundingBox | SegmentPoint]]
        A tuple containing:

        - a list of card boxes (class id 1)
        - a list of patch boxes (class id 0).
    """

    if prediction.boxes is not None:
        ls_cards = [
            box
            for box, class_id in zip(
                prediction.boxes,
                prediction.class_ids,
                strict=False,
            )
            if class_id == 1
        ]
        ls_patches = [
            box
            for box, class_id in zip(
                prediction.boxes,
                prediction.class_ids,
                strict=False,
            )
            if class_id == 0
        ]
    if prediction.segment is not None:
        ls_cards = [
            segment
            for segment, class_id in zip(
                prediction.segment,
                prediction.class_ids,
                strict=False,
            )
            if class_id == 1
        ]
        ls_patches = [
            segment
            for segment, class_id in zip(
                prediction.segment,
                prediction.class_ids,
                strict=False,
            )
            if class_id == 0
        ]
    return ls_cards, ls_patches

print_summary staticmethod

print_summary(prediction: DetectionResult) -> None

Prints a summary of the detected cards and patches.

Parameters:

Name Type Description Default
prediction DetectionResult

The detection result to summarize.

required
Source code in color_correction/processor/detection.py
@staticmethod
def print_summary(prediction: DetectionResult) -> None:
    """
    Prints a summary of the detected cards and patches.

    Parameters
    ----------
    prediction : DetectionResult
        The detection result to summarize.
    """
    ls_cards, ls_patches = DetectionProcessor.get_each_class_box(prediction)
    print(f"Number of cards detected: {len(ls_cards)}")
    print(f"Number of patches detected: {len(ls_patches)}")

process_patches staticmethod

process_patches(input_image: ImageBGR, ordered_patches: list[tuple[BoundingBox, tuple[int, int]] | None]) -> tuple[list[ColorPatchType], ImageBGR]

Processes each detected patch by extracting its region from the image, computing its mean BGR color, and building a visualization grid.

Parameters:

Name Type Description Default
input_image ImageBGR

The original image containing detected patches.

required
ordered_patches list[tuple[BoundingBox, tuple[int, int]] | None]

The list of ordered patch coordinates paired with their center points, possibly with missing entries as None.

required

Returns:

Type Description
tuple[list[ColorPatchType], ImageBGR]

a tuple containing:

  • list of mean BGR color values for each patch.
  • an image visualizing these patches in a grid layout.
Source code in color_correction/processor/detection.py
@staticmethod
def process_patches(
    input_image: ImageBGR,
    ordered_patches: list[tuple[BoundingBox, tuple[int, int]] | None],
) -> tuple[list[ColorPatchType], ImageBGR]:
    """
    Processes each detected patch by extracting its region from the image,
    computing its mean BGR color, and building a visualization grid.

    Parameters
    ----------
    input_image : ImageBGR
        The original image containing detected patches.
    ordered_patches : list[tuple[BoundingBox, tuple[int, int]] | None]
        The list of ordered patch coordinates paired with their center
        points, possibly with missing entries as None.

    Returns
    -------
    tuple[list[ColorPatchType], ImageBGR]
        a tuple containing:

        - list of mean BGR color values for each patch.
        - an image visualizing these patches in a grid layout.
    """
    patch_size = (50, 50, 1)
    ls_bgr_mean_patch = []
    ls_horizontal_patch = []
    ls_vertical_patch = []

    for idx, coord_patch in enumerate(ordered_patches, start=1):
        if coord_patch is None:
            continue

        bbox_patch, _ = coord_patch

        # Extract and process each patch
        cropped_patch = crop_region_with_margin(
            image=input_image,
            coordinates=bbox_patch,
            margin_ratio=0.2,
        )
        bgr_mean_patch = calc_mean_color_patch(cropped_patch)
        ls_bgr_mean_patch.append(bgr_mean_patch)

        # Build visualization
        patch_viz = np.tile(bgr_mean_patch, patch_size)
        ls_horizontal_patch.append(patch_viz)
        if idx % 6 == 0:
            ls_vertical_patch.append(np.hstack(ls_horizontal_patch))
            ls_horizontal_patch = []

    patches_image = np.vstack(ls_vertical_patch)
    return ls_bgr_mean_patch, patches_image

extract_color_patches staticmethod

extract_color_patches(input_image: ImageBGR, prediction: DetectionResult, draw_processed_image: bool = False) -> tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]

Extracts and processes color patches from detected color checker cards, transforming the detection results with available methods.

The method first separates detected cards and patches, generates an expected patch grid, and then matches the detected patches with this grid. If patches are missing, it attempts to auto-fill them with suggested coordinates. Finally, it computes the mean color for each patch and builds a visualization.

Parameters:

Name Type Description Default
input_image ImageBGR

The original image containing the color checker card.

required
prediction DetectionResult

The detection result output from the model.

required
draw_processed_image bool

If True, returns an additional image with visualized detections. Otherwise, only patch processing is performed. Defaults to False.

False

Returns:

Type Description
tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]

a tuple containing:

  • a list of mean BGR color values for each patch.
  • an image visualizing these patches in a grid layout.
  • an optional image with visualized detection results.

Raises:

Type Description
ValueError

If no cards or patches are detected.

Source code in color_correction/processor/detection.py
@staticmethod
def extract_color_patches(
    input_image: ImageBGR,
    prediction: DetectionResult,
    draw_processed_image: bool = False,
) -> tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]:
    """
    Extracts and processes color patches from detected color checker cards,
    transforming the detection results with available methods.

    The method first separates detected cards and patches, generates an expected
    patch grid, and then matches the detected patches with this grid. If patches
    are missing, it attempts to auto-fill them with suggested coordinates.
    Finally, it computes the mean color for each patch and builds a visualization.

    Parameters
    ----------
    input_image : ImageBGR
        The original image containing the color checker card.
    prediction : DetectionResult
        The detection result output from the model.
    draw_processed_image : bool, optional
        If True, returns an additional image with visualized detections.
        Otherwise, only patch processing is performed. Defaults to False.

    Returns
    -------
    tuple[list[ColorPatchType], ImageBGR, ImageBGR | None]
        a tuple containing:

        - a list of mean BGR color values for each patch.
        - an image visualizing these patches in a grid layout.
        - an optional image with visualized detection results.

    Raises
    ------
    ValueError
        If no cards or patches are detected.
    """
    ls_cards, ls_patches = DetectionProcessor.get_each_class_box(prediction)

    if not ls_cards:
        raise ValueError("No cards detected")
    if not ls_patches:
        raise ValueError("No patches detected")

    ls_bgr_mean_patch, grid_patch_img, detection_viz = None, None, None
    if prediction.segment is not None:
        print("Using segmentation results")
        ls_bgr_mean_patch, grid_patch_img, detection_viz = DetectionProcessor.extract_patches_from_segment(
            input_image,
            prediction,
            draw_processed_image,
            ls_cards,
            ls_patches,
        )
        print(f"grid_patch_img shape: {grid_patch_img}")

    if prediction.boxes is not None:
        print("Using box results")
        ls_bgr_mean_patch, grid_patch_img, detection_viz = DetectionProcessor.extract_patches_from_boxes(
            input_image,
            prediction,
            draw_processed_image,
            ls_cards,
            ls_patches,
        )
    return ls_bgr_mean_patch, grid_patch_img, detection_viz

draw_preprocess staticmethod

draw_preprocess(image: ImageBGR, expected_boxes: list[BoundingBox], prediction: DetectionResult, ls_ordered_patch_bbox: list[BoundingBox | None], suggested_patches: dict[int, BoundingBox] | None = None) -> ImageBGR

Draws visualizations on the input image to compare the detected patches with the expected positions. It overlays expected boxes, connects detected patches to their corresponding expected boxes, and highlights individual patch detections.

Parameters:

Name Type Description Default
image ImageBGR

The original image to draw on.

required
expected_boxes list[BoundingBox]

List of expected bounding boxes for patches.

required
prediction DetectionResult

The detection results containing predicted bounding boxes.

required
ls_ordered_patch_bbox list[BoundingBox | None]

The list of ordered patch bounding boxes.

required
suggested_patches dict[int, BoundingBox]

Dictionary of suggested patch coordinates for missing patches.

None

Returns:

Type Description
ImageBGR

The image with drawn detection and patch visualizations.

Source code in color_correction/processor/detection.py
@staticmethod
def draw_preprocess(
    image: ImageBGR,
    expected_boxes: list[BoundingBox],
    prediction: DetectionResult,
    ls_ordered_patch_bbox: list[BoundingBox | None],
    suggested_patches: dict[int, BoundingBox] | None = None,
) -> ImageBGR:
    """
    Draws visualizations on the input image to compare the detected patches with
    the expected positions. It overlays expected boxes, connects detected patches to
    their corresponding expected boxes, and highlights individual patch detections.

    Parameters
    ----------
    image : ImageBGR
        The original image to draw on.
    expected_boxes : list[BoundingBox]
        List of expected bounding boxes for patches.
    prediction : DetectionResult
        The detection results containing predicted bounding boxes.
    ls_ordered_patch_bbox : list[BoundingBox | None]
        The list of ordered patch bounding boxes.
    suggested_patches : dict[int, BoundingBox], optional
        Dictionary of suggested patch coordinates for missing patches.

    Returns
    -------
    ImageBGR
        The image with drawn detection and patch visualizations.
    """
    color_green = (0, 255, 0)
    color_cyan = (255, 255, 10)
    color_violet = (255, 0, 255)
    color_red = (0, 0, 255)
    color_blue = (255, 0, 0)

    result_image = image.copy()

    # Draw all expected boxes
    for idx_b, box in enumerate(expected_boxes):
        cv2.rectangle(
            img=result_image,
            pt1=(box[0], box[1]),
            pt2=(box[2], box[3]),
            color=color_green,
            thickness=2,
        )

        # Draw connection lines between expected and detected patch boxes
        patch = ls_ordered_patch_bbox[idx_b]
        if patch is None:
            continue
        cx, cy = patch[1]
        crefx, crefy = (box[0] + box[2]) // 2, (box[1] + box[3]) // 2
        cv2.line(
            img=result_image,
            pt1=(cx, cy),
            pt2=(crefx, crefy),
            color=color_blue,
            thickness=1,
        )

    # Draw all predicted boxes
    for pbox, pids, pscore in zip(
        prediction.boxes,
        prediction.class_ids,
        prediction.scores,
        strict=False,
    ):
        if pids == 1:
            continue
        cv2.rectangle(
            img=result_image,
            pt1=(pbox[0], pbox[1]),
            pt2=(pbox[2], pbox[3]),
            color=color_cyan,
            thickness=2,
        )
        cv2.putText(
            img=result_image,
            text=f"{pids} {pscore:.2f}",
            org=(pbox[0] + 3, pbox[1] + 12),
            fontFace=cv2.FONT_HERSHEY_SIMPLEX,
            fontScale=0.4,
            color=color_red,
            thickness=1,
            lineType=cv2.LINE_AA,
        )

    # Draw suggested patches if provided
    if suggested_patches:
        for box in suggested_patches.values():
            cv2.rectangle(
                img=result_image,
                pt1=(box[0], box[1]),
                pt2=(box[2], box[3]),
                color=color_violet,
                thickness=2,
            )

    return result_image