[二] pytorch + PyQt5 制作一个图像检测及保存数据集功能

目的

在完成demo的基本功能后，开始想着拓展其他功能，考虑到之前有调试yolov5（pytorch版本)的经验，如果能用UI更加直接的展示其结果，那调试起来也会相对舒服一点。此外，如果可以把得到了检测结果再保存成数据集，那就可以大大减少人力。

需要完成的目标：
一、可以加载模型文件(.pt)，并检测图像内容。
二、可以调整后处理参数(置信区间和交并比阈值)
三、对于多分类模型，可以自定义选择想要识别出的物体类别，并将其再次存储为数据集(.xml)

---

一、使用效果

后处理演示

视频里展示了该功能的演示
流程为: 导入模型和打开图片->选择想要识别的类->点击识别->调整置信区间和交并比阈值->保存数据

二、部分说明

这里重点介绍下保存数据的过程, 如果有调试过yolov5，那应该就会知道数据集的预处理是先得把图像和对应的xml进行转化，同理，我们也可以把检测的结果再转回图像和对应的xml文件.

部分代码说明:

def create_xml(xml_template, save_path, img, bbox, abspath=True):
    if img is not None:
        if abspath:
            save_path = os.path.abspath(save_path)
        data_path = os.path.join(save_path, 'data')
        annotations_path = os.path.join(save_path, 'Annotations')

        if not os.path.exists(data_path):
            os.mkdir(data_path)
        if not os.path.exists(annotations_path):
            os.mkdir(annotations_path)

        img_name = len(os.listdir(data_path))

        dom = xml.dom.minidom.parse(xml_template)
        folder = dom.getElementsByTagName("folder")[0]
        folder.childNodes[0].data = save_path

        filename = dom.getElementsByTagName("filename")[0]
        filename.childNodes[0].data = '{}.jpg'.format(img_name)

        path = dom.getElementsByTagName("path")[0]
        path.childNodes[0].data = os.path.join(data_path, '{}.jpg'.format(img_name))

        database = dom.getElementsByTagName("database")[0]
        database.childNodes[0].data = 'Unknown'

        width = dom.getElementsByTagName("width")[0]
        height = dom.getElementsByTagName("height")[0]
        depth = dom.getElementsByTagName("depth")[0]
        if len(img.shape) < 3:
            width.childNodes[0].data, height.childNodes[0].data, depth.childNodes[0].data = img.shape[0], img.shape[
                1], 1
        else:
            width.childNodes[0].data, height.childNodes[0].data, depth.childNodes[0].data = img.shape

        segmented = dom.getElementsByTagName("segmented")[0]
        segmented.childNodes[0].data = 0

        root = dom.getElementsByTagName("annotation")[0]
        for i in range(len(bbox)):
            o = dom.createElement("object")
            name = dom.createElement("name")
            name.appendChild(dom.createTextNode(str(bbox[i].get('name', 'Unknown'))))
            o.appendChild(name)
            pose = dom.createElement("pose")
            pose.appendChild(dom.createTextNode(str(bbox[i].get('pose', 'Unspecified'))))
            o.appendChild(pose)
            truncated = dom.createElement("truncated")
            truncated.appendChild(dom.createTextNode(str(bbox[i].get('truncated', '0'))))
            o.appendChild(truncated)
            difficult = dom.createElement("difficult")
            difficult.appendChild(dom.createTextNode(str(bbox[i].get('difficult', '0'))))
            o.appendChild(difficult)

            bndbox = dom.createElement("bndbox")
            xmin = dom.createElement("xmin")
            xmax = dom.createElement("xmax")
            ymin = dom.createElement("ymin")
            ymax = dom.createElement("ymax")
            xmin.appendChild(dom.createTextNode(str(bbox[i].get('xmin', '0'))))
            xmax.appendChild(dom.createTextNode(str(bbox[i].get('xmax', '0'))))
            ymin.appendChild(dom.createTextNode(str(bbox[i].get('ymin', '0'))))
            ymax.appendChild(dom.createTextNode(str(bbox[i].get('ymax', '0'))))
            bndbox.appendChild(xmin)
            bndbox.appendChild(ymin)
            bndbox.appendChild(xmax)
            bndbox.appendChild(ymax)
            o.appendChild(bndbox)

            root.appendChild(o)

        cv2.imwrite(os.path.join(data_path, '{}.jpg'.format(img_name)), img)

        with open(os.path.join(annotations_path, "{}.xml".format(img_name)), 'w') as fh:
            dom.writexml(fh)
            
if __name__ == '__main__':
    img = cv2.imread('15.jpg')
    bbox = [{'name': 'halobolite', 'xmin': 8, 'xmax': 493, 'ymin': 161, 'ymax': 452},
            {'name': 'halobolite', 'xmin': 507, 'xmax': 969, 'ymin': 175, 'ymax': 528}]
    create_xml('./example.xml', '.', img, bbox)  # 需要一个模板文件