utils/viz.py

import os
from typing import List
import logging
import PIL
from PIL import Image 
from PIL import ImageDraw
from .annotation import Annotation


def draw_box(im:Image.Image, result, lables, threshold=0.5):
    im = im.copy()
    draw_thickness = min(im.size) // 320
    draw = ImageDraw.Draw(im)
    color_list = get_color_map_list(len(lables))
    clsid2color = {n.lower():color_list[i] for i,n in enumerate(lables)}
    result = [r for r in result if r["score"] >= threshold]

    for dt in result:
        color = tuple(clsid2color[dt["type"]])
        xmin, ymin, xmax, ymax = dt["bbox"]
        
        #Draws a line forming a rectangle around the detected object.

        draw.line(
            [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
             (xmin, ymin)],
            width=draw_thickness,
            fill=color)

        #Prepares the text for the label (class type and score).
        
        
        # draw label
        text = "{} {:.4f}".format(dt["type"], dt["score"])
        #Computes the size of the text using imagedraw_textsize_c.
        tw, th = imagedraw_textsize_c(draw, text)
        #Draws a filled rectangle for the text background.
        draw.rectangle(
            [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color)
        #Draws the text on top of the rectangle.
        draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
    return im

def draw_only_box(im:Image.Image, result):
    im = im.copy()
    draw_thickness = min(im.size) // 400
    draw = ImageDraw.Draw(im)
    result = [r for r in result]

    for dt in result:
        xmin, ymin, xmax, ymax = dt
        xmin = int(xmin)
        ymin = int(ymin)
        xmax = int(xmax)
        ymax = int(ymax)

        draw.line(
            [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
             (xmin, ymin)],
            width=draw_thickness,
            fill="red")

    return im

def draw_box_with_text(im:Image.Image, result:List[Annotation], threshold=0.5):
    im = im.copy()
    draw_thickness = min(im.size) // 320
    draw = ImageDraw.Draw(im)

    result = [r for r in result if r.score >= threshold]

    for dt in result:

        xmin, ymin, xmax, ymax = dt.box
        draw.line(
            [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
             (xmin, ymin)],
            width=draw_thickness,
            fill="red")

        # draw label
        text = "{:.4f}".format(dt.score)
        tw, th = imagedraw_textsize_c(draw, text)
        draw.rectangle(
            [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill="green")
        draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
        
    return im

def get_color_map_list(num_classes):
    """
    Args:
        num_classes (int): number of class
    Returns:
        color_map (list): RGB color list
    """
    color_map = num_classes * [0, 0, 0]
    for i in range(0, num_classes):
        j = 0
        lab = i
        while lab:
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
            j += 1
            lab >>= 3
    color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
    return color_map


def imagedraw_textsize_c(draw, text):
    if int(PIL.__version__.split('.')[0]) < 10:
        tw, th = draw.textsize(text)
    else:
        left, top, right, bottom = draw.textbbox((0, 0), text)
        tw, th = right - left, bottom - top

    return tw, th