Python通过YOLO格式TXT标签文件在图像中画框
使用场景
- 检测数据集标注是否有误: 在目标检测算法中需要标注自己的数据集,为了更加方便的检查数据集标注是否有误,可以使用该工具将标注结果绘制在图像中并查看。
- 美化识别结果中的检测框: 在一些目标检测场景中,YOLO检测算法原始的检测框绘制会导致重叠、颜色冲突、字体过大等问题。可以使用该工具进行修改。
代码
import os
import cv2
class check_label:
def __init__(self, classes:list, label_path:str, img_path:str, result_path:str):
self.classes = classes
self.line_width = 5 #线宽
self.rec_color = (0, 0, 255) #颜色
self.font_color = (255, 255, 255) #字体颜色
self.font = cv2.FONT_HERSHEY_SIMPLEX
self.font_size = 5 #字体大小
self.font_thickness = 4
self.font_x_offset = 0#字体x坐标偏移
self.font_y_offset = -15#字体y坐标偏移
self.isDrawFontRec = False#是否绘制字体矩形框
self.isShowFont = False#是否绘制字体
self.isShowConfidence = False#是否绘制置信度
self.label_path = label_path # 数据集标注结果文件(yolo格式)
self.img_path = img_path # 图像文件
self.result_path = result_path # 在图像上画好标注框文件
self.label_files = os.listdir(label_path)
self.img_files = os.listdir(img_path)
self.label_files.sort(key=lambda x: int(x[:-4]))
self.img_files.sort(key=lambda x: int(x[:-4]))
def paint(self, imgName, pos):
img = cv2.imread(self.img_path + "/" + imgName)
size = img.shape
imgW = size[1]
imgH = size[0]
# print("pos:", len(pos))
for pos_i in pos:
# a, x, y, w, h, b = "orange", pos_i[1], pos_i[2], pos_i[3], pos_i[4], pos_i[4]
pos_i = pos_i.split(' ')
x_center = float(pos_i[1]) * imgW + 1
y_center = float(pos_i[2]) * imgH + 1
x_min = int(x_center - 0.5 * float(pos_i[3]) * imgW)
y_min = int(y_center - 0.5 * float(pos_i[4]) * imgH)
x_max = int(x_center + 0.5 * float(pos_i[3]) * imgW)
y_max = int(y_center + 0.5 * float(pos_i[4]) * imgH)
x = x_min
y = y_min
w = x_max - x_min
h = y_max - y_min
# rotate90(imgW, imgH, x_min, y_min, w, h, x_center, y_center)
# rotate180(imgW, imgH, x_min, y_min, w, h, x_center, y_center)
# rotate270(imgW, imgH, x_min, y_min, w, h, x_center, y_center)
# b = float(pos_i[5])
b = 0.5
if self.isShowConfidence:
a = self.classes[int(pos_i[0])]
else:
a = ""
cv2.rectangle(img, (x, y), (x + w, y + h), self.rec_color, self.line_width)
if self.isDrawFontRec:
cv2.rectangle(img, (x + self.font_x_offset, y + self.font_y_offset), (x + w, y + abs(self.font_y_offset)), self.rec_color,
-1)
if self.isShowFont:
cv2.putText(img, '{} {:.3f}'.format(a, b), (x + self.font_x_offset, y + self.font_y_offset), self.font, self.font_size,
self.font_color, self.font_thickness)
cv2.imwrite(self.result_path + "/" + imgName, img)
def process(self):
for label_file, img_file in zip(self.label_files, self.img_files):
print(img_file, label_file)
if not os.path.isdir(label_file):
f = open(self.label_path + "/" + label_file, "r", encoding='utf-8')
result = f.read().splitlines()
# print(result)
self.paint(img_file, result)
f.close()
原文地址:https://blog.csdn.net/weixin_44223883/article/details/142377203
免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!