DAGM2007数据集转换成VOC格式

DAGM2007数据集下载

链接: DAGM2007.
下载后把每一类的压缩包解压放到文件夹，文件路径放成这样：

• raw_dataset\
  • Class1\
    • Train\
    • -Test\
  • Class2\
    • Train\
    • -Test\
  • …

数据集简单介绍

DAGM2007数据集包含10类缺陷图像，但下载到的数据集大部分是无缺陷图像，而且数据集给出的标签属于弱监督标签，类似于下图。

那么当我们想用目标检测类的方法去做缺陷检测这个任务时，就避免不了将其转换成VOC格式。转换成VOC数据集的话，其中可用图片数目统计如下：

类别	Train	Test
Class1	79	71
Class2	66	84
Class3	66	84
Class4	82	68
Class5	70	80
Class6	83	67
Class7	150	150
Class8	150	150
Class9	150	150
Class10	150	150
总计	1046	1054

转换代码

首先做一个annotations的范例，方便我们后面借用它对图片生成标注文件：
sample.xml:

<annotation>
	<folder>JPEGImages</folder>
	<filename>000001.jpg</filename>
	<path>./DAGM2007_VOC/train/JPEGImages/000001.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>512</width>
		<height>512</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>Class1</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>1</xmin>
			<ymin>1</ymin>
			<xmax>1</xmax>
			<ymax>1</ymax>
		</bndbox>
	</object>
</annotation>

接着我们用数据集自带的椭圆标签生成b-box，然后一些会用到的自定义函数见下面这个脚本：
cover_dataset.py:

import os
import cv2
import math
import sys
import xml.etree.ElementTree as ET


def view_bar(message, num, total):
    """
    Display the schedule
    :param message:
    :param num:
    :param total:
    :return:
    """
    rate = num / total
    rate_num = int(rate * 40)
    rate_nums = math.ceil(rate * 100)
    r = '\r%s:[%s%s]%d%%\t%d/%d' % (message, ">" * rate_num, " " * (40 - rate_num), rate_nums, num, total,)
    sys.stdout.write(r)
    sys.stdout.flush()
    print('\r')


def create_save_folder(save_dir):
    """
    Create the VOC format dataset save dir
    :param save_dir:
    :return:
    """
    voc_train_ann = os.path.join(save_dir, 'train', 'Annotations')
    voc_train_jpg = os.path.join(save_dir, 'train', 'JPEGImages')
    voc_test_ann = os.path.join(save_dir, 'test', 'Annotations')
    voc_test_jpg = os.path.join(save_dir, 'test', 'JPEGImages')
    if not os.path.exists(voc_train_ann):
        os.makedirs(voc_train_ann)
    if not os.path.exists(voc_train_jpg):
        os.makedirs(voc_train_jpg)
    if not os.path.exists(voc_test_ann):
        os.makedirs(voc_test_ann)
    if not os.path.exists(voc_test_jpg):
        os.makedirs(voc_test_jpg)
    return voc_train_ann, voc_train_jpg, voc_test_ann, voc_test_jpg


def get_bbox(img_name):
    """
    Get the bbox from oval
    :param img_name: Label img name in DAGM2007
    :return: xmin, xmax, ymin, ymax
    """
    xmax, xmin, ymax, ymin = 0, 1000, 0, 1000
    src = cv2.imread(img_name)
    for x in range(src.shape[0]):
        for y in range(src.shape[1]):
            if all(src[x, y] == [255, 255, 255]):
                if x > xmax:
                    xmax = x
                if x < xmin:
                    xmin = x
                if y > ymax:
                    ymax = y
                if y < ymin:
                    ymin = y
    ymax, xmax = xmax, ymax
    ymin, xmin = xmin, ymin
    return xmin, xmax, ymin, ymax


def create_xml_file(label_file_path, save_folder, n_id, class_name):
    """
    Write .xml to VOC format dataset
    :param label_file_path: Label img name in DAGM2007
    :param save_folder: Annotation
    :param n_id:
    :param class_name:
    :return:
    """
    img_new_name = '{:06d}'.format(n_id) + '.jpg'
    xmin, xmax, ymin, ymax = get_bbox(label_file_path)
    tree = ET.parse('sample.xml')
    root = tree.getroot()
    for name in root.iter('filename'):
        name.text = '{:06d}'.format(n_id) + '.jpg'
    for name in root.iter('path'):
        name.text = os.path.join(save_folder, img_new_name)
    for obj in root.iter('object'):
        for va in obj.iter('name'):
            va.text = class_name
        for va in obj.iter('xmin'):
            va.text = str(xmin)
        for va in obj.iter('xmax'):
            va.text = str(xmax)
        for va in obj.iter('ymin'):
            va.text = str(ymin)
        for va in obj.iter('ymax'):
            va.text = str(ymax)
    new_xml = '{:06d}'.format(n_id) + '.xml'
    tree.write(os.path.join(save_folder, new_xml))


def write_img(original_folder, name, save_folder, n_id):
    """
    Write image to VOC format dataset
    :param original_folder:
    :param name:
    :param save_folder:
    :param n_id:
    :return:
    """
    img_name = name + '.PNG'
    img_new_name = '{:06d}'.format(n_id) + '.jpg'
    src = cv2.imread(os.path.join(original_folder, img_name))
    cv2.imwrite(os.path.join(save_folder, img_new_name), src)

最后是主函数，修改data_dir为你下载好的路径，save_dir为你保存的路径就行了：
main.py:

import os
import glob
import ntpath
import cover_dataset

# The filename of the download dataset
data_dir = 'raw_dataset'
# The folder name you want to save the VOC format dataset
save_dir = 'DAGM2007_VOC'

EXPECTED_DEFECTIVE_SAMPLES_PER_CLASS = {
    "Train": {
        1: 79, 2: 66, 3: 66, 4: 82, 5: 70,
        6: 83, 7: 150, 8: 150, 9: 150, 10: 150,
    },
    "Test": {
        1: 71, 2: 84, 3: 84, 4: 68, 5: 80,
        6: 67, 7: 150, 8: 150, 9: 150, 10: 150,
    }
}


if __name__ == '__main__':
    voc_train_ann, voc_train_jpg, voc_test_ann, voc_test_jpg = cover_dataset.create_save_folder(save_dir=save_dir)
    train_num = 0
    test_num = 0
    for class_id in range(10):
        # It depends on your own folder
        class_name = "Class%d" % (class_id + 1)
        class_folder_path = os.path.join(data_dir, class_name)

        print("\n[DAGM Preprocessing] Parsing Class ID: %02d ..." % (class_id + 1))

        for data_set in ["Train", "Test"]:

            num_sample = 0

            class_set_folder_path = os.path.join(class_folder_path, data_set)

            img_files = glob.glob(os.path.join(class_set_folder_path, "*.PNG"))

            for file in img_files:

                filepath, fullname = ntpath.split(file)

                filename, extension = os.path.splitext(os.path.basename(fullname))

                lbl_filename = "%s_label.PNG" % filename
                lbl_filepath = os.path.join(filepath, "Label", lbl_filename)

                # We ignore the images without defects.
                # Only if the image contains defects, we use it to construct dataset
                if os.path.exists(lbl_filepath):
                    num_sample += 1
                    img_id = lbl_filepath[-14:-10]
                    if data_set == "Train":
                        train_num += 1
                        cover_dataset.write_img(class_set_folder_path, img_id, voc_train_jpg, train_num)
                        cover_dataset.create_xml_file(lbl_filepath, voc_train_ann, train_num, class_name)
                    else:
                        test_num += 1
                        cover_dataset.write_img(class_set_folder_path, img_id, voc_test_jpg, test_num)
                        cover_dataset.create_xml_file(lbl_filepath, voc_test_ann, test_num, class_name)

            # Ensure the num of defect images is correct
            if num_sample != EXPECTED_DEFECTIVE_SAMPLES_PER_CLASS[data_set][class_id + 1]:
                raise RuntimeError("There should be `%d` defective samples instead of `%d` in challenge (%s): %d" % (
                    EXPECTED_DEFECTIVE_SAMPLES_PER_CLASS[data_set][class_id + 1],
                    num_sample, data_set, class_id + 1))

            # Display the proposs
            cover_dataset.view_bar([data_set + ' Conversion progress:'], num_sample,
                                   EXPECTED_DEFECTIVE_SAMPLES_PER_CLASS[data_set][class_id + 1])

转换结果

最后就可以在save_dir拿到VOC格式的DAGM2007数据集了：

不过有一些b-box过大，我直接拿FPN目标检测网络去做缺陷检测训练测试，10类缺陷的mAP：96.99%
效果不错，详细地后面可能会再写一个博客。