目标检测后的图像上绘制边界框和标签
效果如图所示,有个遗憾就是CV2在图像上显示中文有点难,也不想用别的了,所以改成了英文,代码在下面了,一定要注意一点,就是标注文件的读取一定要根据自己的实际情况改一下,我的所有图像的标注文件是一个XML文件。
import cv2
import os
import numpy as np
def draw_label_type(draw_img,bbox,label_color):
label = str(bbox[-1])
labelSize = cv2.getTextSize(label + '0', cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
if bbox[1] - labelSize[1] - 3 < 0:
# 在图像上绘制边界框
cv2.rectangle(draw_img,
(bbox[0], bbox[1] + 2),
(bbox[0] + labelSize[0], bbox[1] + labelSize[1] + 3),
color=label_color,
thickness=-1
)
# 在图像中的边界框中打上标签
cv2.putText(draw_img, label,
(bbox[0], bbox[1] + labelSize[1] + 3),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 0, 0),
thickness=1
)
else:
# 在图像上绘制边界框
cv2.rectangle(draw_img,
(bbox[0], bbox[1] - labelSize[1] - 3),
(bbox[0] + labelSize[0], bbox[1] - 3),
color=label_color,
thickness=-1
)
# 在图像中的边界框中打上标签
cv2.putText(draw_img, label,
(bbox[0], bbox[1] - 3),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 0, 0),
thickness=1
)
cv2.rectangle(draw_img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color=label_color, thickness=1)
return draw_img
# 读取标注文件
def read_data(data_name):
image_label=[]
with open(data_name, 'r') as f:
for line in f:
image_label.append(line)
return image_label
def spli_lab(word):
labs = []
while (len(word) > 10):
tem = [int(word[-9]),int(word[-8]),int(word[-5]),int(word[-4]),word[-1]]
labs.append(tem)
word = word[:-10]
return labs
def img_ann_ply(label):
for lab in label:
word = lab.split()
#获取一张图象中的标签及位置
# !!!!!怎们分离需要根据自己存储格式改变
img_box = spli_lab(word)
img_name = word[0][:-2]
# 图像文件存储为.bmp。这里因为发现有的标间存储有bug设置了一个筛选
if img_name[-1] != 'p':
img_name = img_name[:-1]
image = os.path.join(inputPath, img_name)
# img = cv2.imread(image)
img = cv2.imdecode(np.fromfile(image, dtype=np.uint8), -1)
# 根据数据集中缺陷的不同设置边界框的颜色
for box in img_box:
if box[-1] == '"虫烂"':
box_color = (255, 0, 0)
box[-1] = 'Insect rot'
elif box[-1] == '"内皮"':
box_color = (0, 0, 255)
box[-1] = 'endothelium'
else:
box_color = (0, 255, 0)
box[-1] = 'charring'
img = draw_label_type(img, box, box_color)
#展示图像
cv2.imshow("banliquexain", img)
# 延时显示,如果想要键盘控制窗口的切换可将int数字改成0
cv2.waitKey(60)
#为了使窗口变得连续,我们将窗口销毁注销
# cv2.destroyAllWindows()
if __name__ == '__main__':
inputPath = r"F:\project\*****\datasets_2000"
dataset_root = r"F:\project\**\datasets_2000\DetectTrainData.txt"
# 读取标注文件,注意!!!!!!
# 这里的标注文件读取会因文件存储格式不同需要自己改动
label=read_data(data_name=dataset_root)
#把标注文件中每张图像分别标注并显示
img_ann_ply(label)