1. 前言

最近在弄关于目标检测的任务，因为检测的图片和标签是分开的，可视化效果不明显，也不知道随便下载的数据集，标注信息对不对。网上看了好多代码，代码风格和本人平时不同，看起来麻烦，也不知道怎么更改，于是按照自己的编码风格，写了个可视化脚本，这里做下记录

本章可视化的数据格式是VOC格式，即标注文件是XML文件，边界框是左上角和右下角的两个点

下面是代码的目录结构，1.jpeg 是数据图像，1.xml是对应的标签信息

result.png 是保存的绘制好边界框的图像

2. 关于代码

代码部分

2.1 自定义函数传入的路径

如下，根据上图摆放好的数据目录，将img_path 传入图片。

对应的 xml 这里是代码生成的，因为大部分数据的标签和训练图像仅仅是目录和后缀的不同，文件名是相同的

2.2 自定义可视化函数

自定义函数传入的三个参数就是原始图片、图片对应的xml标签，save 是展示后是否保存，这里默认保存

如下，parse_xml_to_dict 也是自定义函数，大概就是将xml文件按照树状图的形式解析成层层的字典。看控制台的输出部分，data 就是xml 文件的内容

接下来就是根据解析的字典，读取里面的边界框和检测类别

XML标注格式的文件，目标在<object>下面

因为每张图片的目标不是只有一个，所有用for循环遍历，如下每个i就是一个完整的目标

通过字典的key将对应的value读取出来，然后用列表仅仅保存边界框和分类的名称

值得注意的是，XML标注的目标类别就是真实的类别，而非yolo格式的0 1 2 索引，所以这里绘制边界框是不需要json文件的

xml生成json类别字典文件，可以查看：目标检测篇：如何根据xml标注文件生成类别classes的json文件

关键的信息提取好了，就可以绘制边界框了，通过打印，可以发现ob里面存放的就是我们需要的部分（name、xmin、ymin、xmax、ymax）

绘制的部分很简单，利用cv就行了

2.3 绘制结果

因为大部分训练图像size很大，cv窗口展示不出来，所以这里查看生成的result.png文件

• 想要同时展示多张绘制好的图片，可以利用for循环嵌套，或者用dataloader和plt结合等等
• 待会看看能不能把YOLO格式txt的可视化数据写出来，因为yolo标注的相对信息和类别索引，可能会相对复杂一些

 

3. 完整代码

如下：

from lxml import etree
import cv2


# 读取 xml 文件信息，并返回字典形式
def parse_xml_to_dict(xml):
    if len(xml) == 0:  # 遍历到底层，直接返回 tag对应的信息
        return {xml.tag: xml.text}

    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)  # 递归遍历标签信息
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:  # 因为object可能有多个，所以需要放入列表里
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


# xml 标注文件的可视化
def xmlShow(img,xml,save = True):
    image = cv2.imread(img)
    with open(xml, encoding='gb18030', errors='ignore') as fid:  # 防止出现非法字符报错
        xml_str = fid.read()
    xml = etree.fromstring(xml_str)
    data = parse_xml_to_dict(xml)["annotation"]  # 读取 xml文件信息

    ob = []         # 存放目标信息
    for i in data['object']:        # 提取检测框
        name = str(i['name'])        # 检测的目标类别

        bbox = i['bndbox']
        xmin = int(bbox['xmin'])
        ymin = int(bbox['ymin'])
        xmax = int(bbox['xmax'])
        ymax = int(bbox['ymax'])

        tmp = [name,xmin,ymin,xmax,ymax]    # 单个检测框
        ob.append(tmp)

    # 绘制检测框
    for name,x1,y1,x2,y2 in ob:
        cv2.rectangle(image,(x1,y1),(x2,y2),color=(255,0,0),thickness=2)    # 绘制矩形框
        cv2.putText(image,name,(x1,y1-10),fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.5,thickness=1,color=(0,0,255))

    # 保存图像
    if save:
        cv2.imwrite('result.png',image)

    # 展示图像
    cv2.imshow('test',image)
    cv2.waitKey()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    img_path = './my_xml_dataset/train/images/1.jpeg'  # 传入图片

    labels_path = img_path.replace('images', 'labels')       # 自动获取对应的 xml 标注文件
    labels_path = labels_path.replace('.jpeg', '.xml')

    xmlShow(img=img_path, xml=labels_path,save=True)