通用目标检测标签可视化工具
在目标检测的任务中,我们通常需要对检测结果进行可视化,以便更好地理解模型的性能。本文将介绍一个通用的目标检测标签可视化工具,该工具支持读取VOC格式(XML)、COCO格式(JSON)和YOLO格式(TXT)的标签,并将这些标签以边界框的形式绘制在对应的图像上。
工具简介
本文提供的代码可以自动读取指定目录下的图片和标签文件,并将标签信息绘制在图片上,然后将处理后的图片保存到指定目录。代码使用了OpenCV库进行图像处理,并提供了颜色管理类Colors
,可以方便地为不同类别的对象分配不同的颜色。
代码实现
依赖库
首先,需要确保安装了以下依赖库:
pip install opencv-python
代码实现
以下是实现代码:
import os
import cv2
import xml.etree.ElementTree as ET
import json
class Colors:
def __init__(self):
hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
self.palette = [self.hex2rgb('#' + c) for c in hex]
self.n = len(self.palette)
def __call__(self, i, bgr=False):
c = self.palette[int(i) % self.n]
return (c[2], c[1], c[0]) if bgr else c
@staticmethod
def hex2rgb(h):
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
def plot_one_box(x, im, color=(128, 128, 128), label=None, line_thickness=3):
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
if label:
tf = max(tl - 1, 1)
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA)
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
def parse_voc_annotation(xml_file):
tree = ET.parse(xml_file)
root = tree.getroot()
objects = []
for obj in root.findall('object'):
cls = obj.find('name').text
bbox = obj.find('bndbox')
x1 = int(bbox.find('xmin').text)
y1 = int(bbox.find('ymin').text)
x2 = int(bbox.find('xmax').text)
y2 = int(bbox.find('ymax').text)
objects.append((cls, x1, y1, x2, y2))
return objects
def parse_coco_annotation(json_file, img_id):
with open(json_file, 'r') as f:
data = json.load(f)
objects = []
for ann in data['annotations']:
cls_id = ann['category_id']
bbox = ann['bbox']
x1 = int(bbox[0])
y1 = int(bbox[1])
x2 = int(bbox[0] + bbox[2])
y2 = int(bbox[1] + bbox[3])
objects.append((cls_id, x1, y1, x2, y2))
return objects
def parse_yolo_annotation(txt_file, img_shape):
with open(txt_file, 'r') as f:
lines = f.readlines()
w, h = img_shape[1], img_shape[0]
objects = []
for line in lines:
parts = line.strip().split()
cls_id = int(parts[0])
x_center = float(parts[1]) * w
y_center = float(parts[2]) * h
width = float(parts[3]) * w
height = float(parts[4]) * h
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
objects.append((cls_id, x1, y1, x2, y2))
return objects
def show_label(path_root_imgs, path_root_labels, save_path_root_imgs):
class_categories = ['Abrasion', 'Crazing', 'Patches', 'Inclusion', 'Uneven', 'Blowhole', 'Break', 'Crack', 'Crescent_Gap', 'Crease', 'Silk-Spot', 'Water-Spot', 'Weld-Line', 'GC-Inclusion', 'Oil-Spot', 'Rolled-Pit', 'Punching', 'Waist-Folding', 'Bruise', 'Pitted_Surface', 'Rolled-in_Scale', 'Scratches', 'Bubble']
if not os.path.exists(save_path_root_imgs):
os.makedirs(save_path_root_imgs)
colors = Colors()
for root, _, files in os.walk(path_root_imgs):
for file in files:
if file.endswith('.jpg'):
img_path = os.path.join(root, file)
label_path = os.path.join(path_root_labels, file.replace('.jpg', '.txt'))
save_path = os.path.join(save_path_root_imgs, file)
img = cv2.imread(img_path)
if img is None:
continue
objects = []
if os.path.exists(label_path):
objects = parse_yolo_annotation(label_path, img.shape)
else:
xml_path = os.path.join(path_root_labels, file.replace('.jpg', '.xml'))
json_path = os.path.join(path_root_labels, file.replace('.jpg', '.json'))
if os.path.exists(xml_path):
objects = parse_voc_annotation(xml_path)
elif os.path.exists(json_path):
img_id = int(file.split('.')[0])
objects = parse_coco_annotation(json_path, img_id)
img_tmp = img.copy()
for obj in objects:
if isinstance(obj[0], int):
label = class_categories[obj[0]]
else:
label = obj[0]
plot_one_box(obj[1:], img_tmp, color=colors(obj[0]), label=label)
cv2.imwrite(save_path, img_tmp)
print(f"Saved: {save_path}")
if __name__ == '__main__':
path_root_labels = r'路径到标签文件夹'
path_root_imgs = r'路径到图片文件夹'
save_path_root_imgs = r'保存结果的文件夹路径'
show_label(path_root_imgs, path_root_labels, save_path_root_imgs)
代码使用说明
准备工作
- 安装依赖:确保安装了OpenCV库,使用以下命令安装:
pip install opencv-python
- 文件组织:将图片文件和标签文件分别放置在两个不同的文件夹中。
参数设置
在运行代码前,需要根据自己的文件路径进行参数设置:
path_root_labels
:标签文件夹路径。path_root_imgs
:图片文件夹路径。save_path_root_imgs
:保存结果的文件夹路径。
运行代码
配置好路径后,直接运行代码即可:
python your_script_name.py
功能说明
- Colors类:用于管理颜色的分配,每个类别的对象会分配一个唯一的颜色。
- plot_one_box函数:在图像上绘制一个边界框,并在边界框上显示标签信息。
- parse_voc_annotation函数:解析VOC格式(XML)的标签文件。
- parse_coco_annotation函数:解析COCO格式(JSON)的标签文件。
- parse_yolo_annotation函数:解析YOLO格式(TXT)的标签文件。
- show_label函数:读取指定目录下的图片和标签文件,调用上述函数解析标签并绘制在图片上,最后将结果保存到指定目录。