tuoheng_alg_move_20240715/util/PlotsUtils.py

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import unicodedata
FONT_PATH = "../AIlib2/conf/platech.ttf"

zhFont = ImageFont.truetype(FONT_PATH, 20, encoding="utf-8") 

def get_label_array(color=None, label=None, font=None, fontSize=40, unify=False):
    if unify:
        x, y, width, height = font.getbbox("标")    # 统一数组大小
    else:
        x, y, width, height = font.getbbox(label)
    text_image = np.zeros((height, width, 3), dtype=np.uint8)
    text_image = Image.fromarray(text_image)
    draw = ImageDraw.Draw(text_image)
    draw.rectangle((0, 0, width, height), fill=tuple(color))
    draw.text((0, -1), label, fill=(255, 255, 255), font=font)
    im_array = np.asarray(text_image)
    # scale = fontSize / height
    # im_array = cv2.resize(im_array, (0, 0), fx=scale, fy=scale)
    scale = height / fontSize
    im_array = cv2.resize(im_array, (0, 0), fx=scale, fy=scale)
    return im_array


def get_label_arrays(labelNames, colors, fontSize=40, fontPath="platech.ttf"):
    font = ImageFont.truetype(fontPath, fontSize, encoding='utf-8')
    label_arraylist = [get_label_array(colors[i % 20], label_name, font, fontSize) for i, label_name in
                       enumerate(labelNames)]
    return label_arraylist

def get_label_array_dict(colors, fontSize=40, fontPath="platech.ttf"):
    font = ImageFont.truetype(fontPath, fontSize, encoding='utf-8')
    all_chinese_characters = []
    for char in range(0x4E00, 0x9FFF + 1):  # 中文
        chinese_character = chr(char)
        if unicodedata.category(chinese_character) == 'Lo':
            all_chinese_characters.append(chinese_character)
    for char in range(0x0041, 0x005B):  # 大写字母
        all_chinese_characters.append(chr(char))
    for char in range(0x0061, 0x007B):  # 小写字母
        all_chinese_characters.append(chr(char))
    for char in range(0x0030, 0x003A):  # 数字
        all_chinese_characters.append(chr(char))
    zh_dict = {}
    for code in all_chinese_characters:
        arr = get_label_array(colors[2], code, font, fontSize, unify=True)
        zh_dict[code] = arr
    return zh_dict


def xywh2xyxy(box):
    if not isinstance(box[0], (list, tuple, np.ndarray)):
        xc, yc, w, h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
        bw, bh = int(w / 2), int(h / 2)
        lt, yt, rt, yr = xc - bw, yc - bh, xc + bw, yc + bh
        box = [(lt, yt), (rt, yt), (rt, yr), (lt, yr)]
    return box

def xywh2xyxy2(param):
    if not isinstance(param[0], (list, tuple, np.ndarray)):
        xc, yc, x2, y2 = int(param[0]), int(param[1]), int(param[2]), int(param[3])
        return [(xc, yc), (x2, yc), (x2, y2), (xc, y2)], float(param[4]), int(param[5])
        # bw, bh = int(w / 2), int(h / 2)
        # lt, yt, rt, yr = xc - bw, yc - bh, xc + bw, yc + bh
        # return [(lt, yt), (rt, yt), (rt, yr), (lt, yr)]
    return np.asarray(param[0][0:4], np.int32), float(param[1]), int(param[2])


def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=None, isNew=False):
    # 识别问题描述图片的高、宽
    lh, lw = label_array.shape[0:2]
    # 图片的长度和宽度
    imh, imw = img.shape[0:2]
    box = xywh2xyxy(box)
    # 框框左上的位置
    x0, y1 = box[0][0], box[0][1]
    # if score_location == 'leftTop':
    #     x0, y1 = box[0][0], box[0][1]
    # # 框框左下的位置
    # elif score_location == 'leftBottom':
    #     x0, y1 = box[3][0], box[3][1]
    # else:
    #     x0, y1 = box[0][0], box[0][1]
    # x1 框框左上x位置 + 描述的宽
    # y0 框框左上y位置 - 描述的高
    x1, y0 = x0 + lw, y1 - lh
    # 如果y0小于0, 说明超过上边框
    if y0 < 0:
        y0 = 0
        # y1等于文字高度
        y1 = y0 + lh
    # 如果y1框框的高大于图片高度
    if y1 > imh:
        # y1等于图片高度
        y1 = imh
        # y0等于y1减去文字高度
        y0 = y1 - lh
    # 如果x0小于0
    if x0 < 0:
        x0 = 0
        x1 = x0 + lw
    if x1 > imw:
        x1 = imw
        x0 = x1 - lw
    # box_tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1
    '''
    1. img（array） 为ndarray类型（可以为cv.imread）直接读取的数据
    2. box（array）：为所画多边形的顶点坐标
    3. 所画四边形是否闭合，通常为True
    4. color（tuple）：BGR三个通道的值
    5. thickness（int）：画线的粗细
    6. shift：顶点坐标中小数的位数
    '''
    tl = config[0]
    box1 = np.asarray(box, np.int32)
    cv2.polylines(img, [box1], True, color, tl)
    img[y0:y1, x0:x1, :] = label_array
    pts_cls = [(x0, y0), (x1, y1)]
    # 把英文字符score画到类别旁边
    # tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1
    label = ' %.2f' % score
    # tf = max(tl, 1)
    # fontScale = float(format(imw / 1920 * 1.1, '.2f')) or tl * 0.33
    # fontScale = tl * 0.33
    '''
    1. text：要计算大小的文本内容，类型为字符串。
    2. fontFace：字体类型，例如cv2.FONT_HERSHEY_SIMPLEX等。
    3. fontScale：字体大小的缩放因子，例如1.2表示字体大小增加20%。
    4. thickness：文本线条的粗细，以像素为单位。
    5. (text_width, text_height)：给定文本在指定字体、字体大小、线条粗细下所占用的像素宽度和高度。
    '''
    # t_size = cv2.getTextSize(label, 0, fontScale=fontScale, thickness=tf)[0]
    t_size = (config[1], config[2])
    # if   socre_location=='leftTop':
    p1, p2 = (pts_cls[1][0], pts_cls[0][1]), (pts_cls[1][0] + t_size[0], pts_cls[1][1])
    '''
    1. img：要绘制矩形的图像
    2. pt1：矩形框的左上角坐标，可以是一个包含两个整数的元组或列表，例如(x1, y1)或[x1, y1]。
    3. pt2：矩形框的右下角坐标，可以是一个包含两个整数的元组或列表，例如(x2, y2)或[x2, y2]。
    4. color：矩形框的颜色，可以是一个包含三个整数的元组或列表，例如(255, 0, 0)表示蓝色，或一个标量值，例如255表示白色。颜色顺序为BGR。
    5. thickness：线条的粗细，以像素为单位。如果为负值，则表示要绘制填充矩形。默认值为1。
    6. lineType：线条的类型，可以是cv2.LINE_AA表示抗锯齿线条，或cv2.LINE_4表示4连通线条，或cv2.LINE_8表示8连通线条。默认值为cv2.LINE_8。
    7. shift：坐标点小数点位数。默认值为0。
    '''
    cv2.rectangle(img, p1, p2, color, -1, cv2.LINE_AA)
    p3 = pts_cls[1][0], pts_cls[1][1] - (lh - t_size[1]) // 2
    '''
    1. img：要在其上绘制文本的图像
    2. text：要绘制的文本内容，类型为字符串
    3. org：文本起始位置的坐标，可以是一个包含两个整数的元组或列表，例如(x, y)或[x, y]。
    4. fontFace：字体类型，例如cv2.FONT_HERSHEY_SIMPLEX等。
    5. fontScale：字体大小的缩放因子，例如1.2表示字体大小增加20%。
    6. color：文本的颜色，可以是一个包含三个整数的元组或列表，例如(255, 0, 0)表示蓝色，或一个标量值，例如255表示白色。颜色顺序为BGR。
    7. thickness：文本线条的粗细，以像素为单位。默认值为1。
    8. lineType：线条的类型，可以是cv2.LINE_AA表示抗锯齿线条，或cv2.LINE_4表示4连通线条，或cv2.LINE_8表示8连通线条。默认值为cv2.LINE_8。
    9. bottomLeftOrigin：文本起始位置是否为左下角。如果为True，则文本起始位置为左下角，否则为左上角。默认值为False。
    '''
    if isNew:
        cv2.putText(img, label, p3, 0, config[3], [0, 0, 0], thickness=config[4], lineType=cv2.LINE_AA)
    else:
        cv2.putText(img, label, p3, 0, config[3], [225, 255, 255], thickness=config[4], lineType=cv2.LINE_AA)
    return img, box

# 动态标签
def draw_name_joint(box, img, label_array_dict, score=0.5, color=None, config=None, name=""):
    label_array = None
    for zh in name:
        if zh in label_array_dict:
            if label_array is None:
                label_array = label_array_dict[zh]
            else:
                label_array = np.concatenate((label_array,label_array_dict[zh]), axis= 1)
    # 识别问题描述图片的高、宽
    if label_array is None:
        lh, lw = 0, 0
    else:
        lh, lw = label_array.shape[0:2]
    # 图片的长度和宽度
    imh, imw = img.shape[0:2]
    box = xywh2xyxy(box)
    # 框框左上的位置
    x0, y1 = box[0][0], box[0][1]
    x1, y0 = x0 + lw, y1 - lh
    # 如果y0小于0, 说明超过上边框
    if y0 < 0:
        y0 = 0
        # y1等于文字高度
        y1 = y0 + lh
    # 如果y1框框的高大于图片高度
    if y1 > imh:
        # y1等于图片高度
        y1 = imh
        # y0等于y1减去文字高度
        y0 = y1 - lh
    # 如果x0小于0
    if x0 < 0:
        x0 = 0
        x1 = x0 + lw
    if x1 > imw:
        x1 = imw
        x0 = x1 - lw
    tl = config[0]
    box1 = np.asarray(box, np.int32)
    cv2.polylines(img, [box1], True, color, tl)
    if label_array is not None:
        img[y0:y1, x0:x1, :] = label_array
    pts_cls = [(x0, y0), (x1, y1)]
    # 把英文字符score画到类别旁边
    # tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1
    label = ' %.2f' % score
    t_size = (config[1], config[2])
    # if   socre_location=='leftTop':
    p1, p2 = (pts_cls[1][0], pts_cls[0][1]), (pts_cls[1][0] + t_size[0], pts_cls[1][1])
    cv2.rectangle(img, p1, p2, color, -1, cv2.LINE_AA)
    p3 = pts_cls[1][0], pts_cls[1][1] - (lh - t_size[1]) // 2
    cv2.putText(img, label, p3, 0, config[3], [225, 255, 255], thickness=config[4], lineType=cv2.LINE_AA)
    return img, box


def filterBox(det0, det1, pix_dis):
    # det0为 (m1, 11) 矩阵
    # det1为 (m2, 12) 矩阵
    if len(det0.shape) == 1:
        det0 = det0[np.newaxis,...]
    if len(det1.shape) == 1:
        det1 = det1[np.newaxis,...]
    det1 = det1[...,0:11].copy()
    m, n = det0.size, det1.size
    if not m:
        return det0
    # 在det0的列方向加一个元素flag代表该目标框中心点是否在之前目标框内(0代表不在，其他代表在)
    flag = np.zeros([len(det0), 1])
    det0 = np.concatenate([det0, flag], axis=1)
    det0_copy = det0.copy()
    # det1_copy = det1.copy()
    if not n:
        return det0
    # det0转成 (m1, m2, 12) 的矩阵
    # det1转成 (m1, m2, 12) 的矩阵
    # det0与det1在第3维方向上拼接(6 + 7 = 13)
    det0 = det0[:, np.newaxis, :].repeat(det1.shape[0], 1)
    det1 = det1[np.newaxis, ...].repeat(det0.shape[0], 0)
    joint_det = np.concatenate((det1, det0), axis=2)
    # 分别求det0和det1的x1, y1, x2, y2(水平框的左上右下角点)
    x1, y1, x2, y2 = joint_det[..., 0], joint_det[..., 1], joint_det[..., 4], joint_det[..., 5]
    x3, y3, x4, y4 = joint_det[..., 11], joint_det[..., 12], joint_det[..., 15], joint_det[..., 16]

    x2_c, y2_c = (x1+x2)//2, (y1+y2)//2
    x_c, y_c = (x3+x4)//2, (y3+y4)//2
    dis = (x2_c - x_c)**2 + (y2_c - y_c)**2
    mask = (joint_det[..., 9] == joint_det[..., 20]) & (dis <= pix_dis**2)
    
    # 类别相同 & 中心点在上一帧的框内 判断为True
    res = np.sum(mask, axis=1)
    det0_copy[..., -1] = res
    return det0_copy