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%
效果不錯，詳細地後面可能會再寫一個博客。