Commit d4ecbefe authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add new params for visualize

parent 75edcff7

imgutils/detect/visual.py

+22 −9
Original line number Diff line number Diff line
@@ -4,6 +4,7 @@ Overview: 


    See :func:`imgutils.detect.head.detect_heads` and :func:`imgutils.detect.person.detect_person` for examples.

"""

import math

from typing import List, Tuple, Optional



from PIL import ImageFont, ImageDraw
@@ -27,7 +28,7 @@ def _try_get_font_from_matplotlib(fp=None, fontsize: int = 12): 


def detection_visualize(image: ImageTyping, detection: List[Tuple[Tuple[float, float, float, float], str, float]],

                        labels: Optional[List[str]] = None, text_padding: int = 6, fontsize: int = 12,

                        fp=None, no_label: bool = False):

                        max_short_edge_size: Optional[int] = None, fp=None, no_label: bool = False):

    """

    Overview:

        Visualize the results of the object detection.
@@ -49,27 +50,39 @@ def detection_visualize(image: ImageTyping, detection: List[Tuple[Tuple[float, f 
        See :func:`imgutils.detect.head.detect_heads` and :func:`imgutils.detect.person.detect_person` for examples.

    """

    image = load_image(image, force_background=None, mode='RGBA')

    original_width, original_height = image.width, image.height

    if max_short_edge_size is not None and max_short_edge_size < min(original_height, original_width):

        r = max_short_edge_size / min(original_height, original_width)

        new_width = int(math.ceil(original_width * r))

        new_height = int(math.ceil(original_height * r))

    else:

        new_width, new_height = original_width, original_height



    visual_image = image.copy()

    if (new_width, new_height) != (original_width, original_height):

        visual_image = visual_image.resize((new_width, new_height))

    draw = ImageDraw.Draw(visual_image, mode='RGBA')

    font = _try_get_font_from_matplotlib(fp, fontsize) or ImageFont.load_default()



    labels = sorted(labels or {label for _, label, _ in detection})

    _colors = list(map(str, rnd_colors(len(labels))))

    _color_map = dict(zip(labels, _colors))

    for _, ((xmin, ymin, xmax, ymax), label, score) in sorted(enumerate(detection), key=lambda x: (x[1][2], x[0])):

    for _, ((x0, y0, x1, y1), label, score) in sorted(enumerate(detection), key=lambda x: (x[1][2], x[0])):

        x0, y0 = int(x0 * new_width / original_width), int(y0 * new_height / original_height)

        x1, y1 = int(x1 * new_width / original_width), int(y1 * new_height / original_height)

        box_color = _color_map[label]

        draw.rectangle((xmin, ymin, xmax, ymax), outline=box_color, width=2)

        draw.rectangle((x0, y0, x1, y1), outline=box_color, width=2)



        if not no_label:

            label_text = f'{label}: {score * 100:.2f}%'

            _t_x0, _t_y0, _t_x1, _t_y1 = draw.textbbox((xmin, ymin), label_text, font=font)

            _t_x0, _t_y0, _t_x1, _t_y1 = draw.textbbox((x0, y0), label_text, font=font)

            _t_width, _t_height = _t_x1 - _t_x0, _t_y1 - _t_y0

            if ymin - _t_height - text_padding < 0:

                _t_text_rect = (xmin, ymin, xmin + _t_width + text_padding * 2, ymin + _t_height + text_padding * 2)

                _t_text_co = (xmin + text_padding, ymin + text_padding)

            if y0 - _t_height - text_padding < 0:

                _t_text_rect = (x0, y0, x0 + _t_width + text_padding * 2, y0 + _t_height + text_padding * 2)

                _t_text_co = (x0 + text_padding, y0 + text_padding)

            else:

                _t_text_rect = (xmin, ymin - _t_height - text_padding * 2, xmin + _t_width + text_padding * 2, ymin)

                _t_text_co = (xmin + text_padding, ymin - _t_height - text_padding)

                _t_text_rect = (x0, y0 - _t_height - text_padding * 2, x0 + _t_width + text_padding * 2, y0)

                _t_text_co = (x0 + text_padding, y0 - _t_height - text_padding)



            draw.rectangle(_t_text_rect, fill=str(Color(box_color, alpha=0.5)))

            draw.text(_t_text_co, label_text, fill="black", font=font)

test/detect/test_visual.py

+15 −0
Original line number Diff line number Diff line
@@ -21,3 +21,18 @@ class TestDetectVisual: 
            visual.convert('RGB'),

            throw_exception=False

        ) < 1e-2



    def test_detection_visualize_480(self, image_diff):

        image = get_testfile('genshin_post.jpg')

        visual = detection_visualize(image, [

            ((202, 155, 356, 294), 'first', 0.878),

            ((938, 87, 1121, 262), 'second', 0.846),

            ((652, 440, 725, 514), 'third', 0.839),

            ((464, 250, 535, 326), 'fourth', 0.765)

        ], fontsize=24, max_short_edge_size=480)



        assert image_diff(

            load_image(get_testfile('genshin_post_face_visual_480.jpg'), mode='RGB'),

            visual.convert('RGB'),

            throw_exception=False

        ) < 1e-2

test/testfile/genshin_post_face_visual_480.jpg

0 → 100644
+91.5 KiB
Loading image diff...