258 lines
11 KiB
Python
258 lines
11 KiB
Python
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 |