最近在做目标检测,需要将yolo格式下的txt标签数据贴回原图,跟模型预测的结果进行对比,所以写了一个代码:该代码实现的是标签数据贴回原图。
import cv2
import numpy as np
#------------------------------------
#读取txt文件,返回一个列表
def hcp_txt2list(txtpath):
# txt_path = r'./0.txt'
with open(txtpath,'r') as f :
bbox_list = f.readlines()
# print(bbox_list)
f.close()
return bbox_list
#对读取回来的boxlist做成可以使用box
def hcp_boxlist2box(bboxlist):
# 需要知道图像的高H和宽W
H = 416
W = 416
# nbox =[]
Nbox = []
for box in bboxlist:
cls = (box.split(' ')[0])
x = float(box.split(' ')[1])
y = float(box.split(' ')[2])
h = float(box.split(' ')[3])
w = float(box.split(' ')[4])
# print(cls,x,y,h,w)
# 还原成真实坐标
x0 = x * W
y0 = y * H # 目标中心点(x0,y0)
h = h * H
w = w * W
# 需要几个点?
x0 = int(x0)
y0 = int(y0)
h = int(h)
w = int(w)
nbox = np.array([x0, y0, h, w])
Nbox.append(nbox)
# cv2.rectangle(img,)
# print(Nbox)
return Nbox
#对一张图像画框
def hcp_drawrectangle(img,Nbox):
for box in Nbox:
p1 =(int(box[0]-(box[2]/2)),int(box[1]-(box[3]/2)))
p2 = (int(box[0] + (box[2] / 2)),int( box[1] + (box[3] / 2)))
color =(0,0,255)
print(p1,p2,color)
img = cv2.rectangle(img,p1,p2,color)
return img
#------------------------------------
if __name__ == '__main__':
pass
#读取txt文件
txt_path = r'./7.txt'
bboxlist = hcp_txt2list(txt_path)
Nbox = hcp_boxlist2box(bboxlist)
#读取一张图像
img_path = r'./7.jpg'
img = cv2.imread(img_path)
img = hcp_drawrectangle(img, Nbox)
cv2.imshow('img ',img)
cv2.waitKey(0)