将图像分割数据集转为目标检测数据集（外接矩形坐标保存到XML文件中）

一、前言

        通过借助cv2.findContours（）函数来获取轮廓的坐标，再对坐标进行排序从而画出外接矩形，再将坐标转存到XML文件中，实现目标检测数据集构建（VOC格式），并根据xml文件在原图画出矩形框验证转换结果。

二、代码

        原图

        1、将外接矩形坐标保存到xml文件中（注意修改相应的地址以及图片名的后缀）

import cv2
import numpy as np
import pandas as pd
import os



def get_coor(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # 变为灰度图
    ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)  ## 阈值分割得到二值化图片
    contours, heriachy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    print(len(contours))
    for i, contour in enumerate(contours):
        if i == len(contours)-1:
            print('i:', i)
            a = sorted(contour[:, 0], key=lambda x: x[0])  # 所有坐标按x轴从小到大排序
            x_min = a[0][0]
            x_max = a[-1][0]
            b = sorted(contour[:, 0], key=lambda x: x[1])  # 所有坐标按y轴从小到大排序
            y_min = b[0][1]
            y_max = b[-1][1]
            rec = img
                # cv2.drawContours(img, contours, i, (0, 0, 255), 5)
                # 第一个参数指在哪幅图上绘制轮廓信息，第二个参数是轮廓本身，第三个参数是指定绘制哪条轮廓
                # 第四个参数是绘图的颜色，第五个参数是绘制的线宽 输入-1则表示填充
            cv2.rectangle(rec, (x_min, y_min), (x_max, y_max), (0, 255, 0), 3)
            # cv2.imshow('rectangle', rec)
            # cv2.waitKey()
            return x_min, y_min, x_max, y_max

def save_xml(src_xml_dir, img_name, h, w, x1, y1, x2, y2):

    xml_file = open((src_xml_dir + '/' + img_name + '.xml'), 'w')
    xml_file.write('\n')
    xml_file.write('    VOC2007\n')


    xml_file.write('    ' + str(img_name) + '.jpg' + '\n')


    xml_file.write('    \n')
    xml_file.write('        ' + str(w) + '\n')
    xml_file.write('        ' + str(h) + '\n')
    xml_file.write('        3\n')
    xml_file.write('    \n')

    # write the region of image on xml file
    xml_file.write('    \n')
    xml_file.write('')



file_dir = r'D:\YQ\General_code\Json2mask\2_preprocessed_data_stage2_1\datasets\SegmentationClass'
save_xml_dir = r'D:\YQ\General_code\Json2mask\2_preprocessed_data_stage2_1\datasets\rectangle_xml'
for name in os.listdir(file_dir):
    print(name)
    img_path = os.path.join(file_dir, name)
    img = cv2.imread(img_path)
    # img = cv2.imdecode(np.fromfile(img_path), dtype=np.uint8), -1)
    h, w = img.shape[:-1]
    x1, y1, x2, y2 = get_coor(img)
    img_name = name.split('.')[0]
    save_xml(save_xml_dir,img_name, h, w, x1, y1, x2, y2)

 

        生成的XML文件

        2、读取xml文件，将矩形画在原图上（注意修改路径）

import os
import xml.etree.ElementTree  as ET
import cv2
img_root = ""  # 图片地址
xml_root = ""  # xml文件地址
dstRoot = ""  # 保存画完矩形框后得到的图片的地址

for img in os.listdir(img_root):
    iname = img[:-4]
    for xml in os.listdir(xml_root):
        xname = xml[:-4]
        if iname == xname:
            xmlpath = os.path.join(xml_root, xml)
            xmlfile = open(xmlpath, encoding='utf-8')
            tree = ET.parse(xmlfile)
            root = tree.getroot()
            imgpath = os.path.join(img_root, img)
            mat = cv2.imread(imgpath)
            darwpath = os.path.join(dstRoot, img)
            size = root.find('size')
            w = int(size.find('width').text)
            h = int(size.find('height').text)
            for obj in root.iter('object'):
                # difficult = obj.find('difficult').text
                cls = obj.find('name').text
                xmlbox = obj.find('bndbox')
                b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
                     float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
                x1, x2, y1, y2 = b
                # 标注越界修正
                if x2 > w:
                    x2 = w
                if y2 > h:
                    y2 = h
                b = (x1, x2, y1, y2)
                # print('%d, %d, %d, %d' % (x1, x2, y1, y2))
                pt1 = (int(x1), int(y1))
                pt2 = (int(x2), int(y2))
                color = (255, 0, 255)
                cv2.rectangle(mat, pt1, pt2, color, 4, 4)
            # cv2.imshow('show', mat)
            # cv2.waitKey(33)
            cv2.imwrite(darwpath, mat)

        结果