MVTecAD データセットローダー TensorFlow

はじめに

MVTecADデータセットローダープログラム。

プログラム①：MVTecAD クラス

Data_Loader.py

# MVTec AD Data Loader

# -----------------------------------------------
# Import
# -----------------------------------------------
# 01. Built-in
import os

# 02. Third-party modules
import cv2
import numpy as np
import tensorflow as tf




# -----------------------------------------------
# Class
# -----------------------------------------------

class MVtecADLoader():
    def __init__(self, img_size = 224):

        # Datasets Path

        path_notebook = 'C:/Users/user_name/'
        path_nokori = 'Google ドライブ/datasets/mvtec_anomaly_detection'
        self.base_path = os.path.join(path_notebook, path_nokori)

        # Image Resolution
        self.img_size = img_size

        # number
        self.num_train = None
        # Category NG List

        self.category = {'bottle': ['good', 'broken_large', 'broken_small', 'contamination'],
                         'cable': ['good', 'bent_wire', 'cable_swap', 'combined', 'cut_inner_insulation',
                                   'cut_outer_insulation',
                                   'missing_cable', 'missing_wire', 'poke_insulation'],
                         'capsule': ['good', 'crack', 'faulty_imprint', 'poke', 'scratch', 'squeeze'],
                         'carpet': ['good', 'color', 'cut', 'hole', 'metal_contamination', 'thread'],
                         'grid': ['good', 'bent', 'broken', 'glue', 'metal_contamination', 'thread'],
                         'hazelnut': ['good', 'crack', 'cut', 'hole', 'print'],
                         'leather': ['good', 'color', 'cut', 'fold', 'glue', 'poke'],
                         'metal_nut': ['good', 'bent', 'color', 'flip', 'scratch'],
                         'pill': ['good', 'color', 'combined', 'contamination', 'crack', 'faulty_imprint', 'pill_type',
                                  'scratch'],
                         'screw': ['good', 'manipulated_front', 'scratch_head', 'scratch_neck', 'thread_side',
                                   'thread_top'],
                         'tile': ['good', 'crack', 'glue_strip', 'gray_stroke', 'oil', 'rough'],
                         'toothbrush': ['good', 'defective'],
                         'transistor': ['good', 'bent_lead', 'cut_lead', 'damaged_case', 'misplaced'],
                         'wood': ['good', 'color', 'combined', 'good', 'hole', 'liquid', 'scratch'],
                         'zipper': ['good', 'broken_teeth', 'combined', 'fabric_border', 'fabric_interior', 'rough',
                                    'split_teeth', 'squeezed_teeth']}



    def load(self, category, repeat = 4, max_rot = 10):

        #---------------------------------------------
        # 01. Load Train Dataset (OK)
        # ---------------------------------------------
        #data, mask, binary anomaly label ( 0 for anomaly, 1 for good)
        x, y, z = [], [], []

        path = os.path.join(os.path.join(self.base_path, category), 'train/good')
        files = os.listdir(path)

        zero_mask = np.zeros(shape = (self.img_size, self.img_size), dtype = np.float32)

        for rdx in range(repeat):
            for file in files:
                # x, data
                img_full_path = os.path.join(path, file)
                img = self._read_image(img_path= img_full_path) #(RGB)

                if not max_rot == 0:
                    img = tf.keras.preprocessing.image.random_rotation(img, max_rot)

                # y, mask
                img_mask = zero_mask
                # z, binary ( 0 for anomaly, 1 for good)
                z_element = 1

                x.append(img)
                y.append(img_mask)
                z.append(z_element)

        x = np.asarray(x)
        y = np.asarray(y)
        z = np.asarray(z)
        self.num_train = len(x)

        x = tf.data.Dataset.from_tensor_slices(tf.convert_to_tensor(x, dtype = tf.float32))
        y = tf.data.Dataset.from_tensor_slices(tf.convert_to_tensor(y, dtype = tf.int32))
        z = tf.data.Dataset.from_tensor_slices(tf.convert_to_tensor(z, dtype = tf.int32))


        self.train = tf.data.Dataset.zip((x,y,z))

        # ---------------------------------------------
        # 02. Load Test Dataset (OK, NG)
        # ---------------------------------------------

        #data, mask, binary anomaly label ( 0 for anomaly, 1 for good)
        x, y, z = [], [], []

        for label in self.category[category]:
            path = os.path.join(os.path.join(self.base_path, category),'test/{}'.format(label))
            files = os.listdir(path)

            for file in files:
                # x, data
                img_full_path = os.path.join(path, file)
                img = self._read_image(img_path = img_full_path)

                if label =='good':
                    mask = zero_mask
                    idx_z = 1 #label = 1 for good
                else:
                    mask_path = os.path.join(os.path.join(self.base_path, category), 'ground_truth/{}'.format(label))
                    mask_path_file = os.path.join(mask_path, '{}_mask.png'.format(file.split('.')[0]))

                    mask = self._read_image(img_path= mask_path_file, flags = cv2.IMREAD_GRAYSCALE)
                    mask = mask/255
                    idx_z = 0  # label = 0 for NG
                # mask = mask[16:-16, 16:-16]
                mask = tf.convert_to_tensor(mask, dtype = tf.int32)


            x.append(img)
            y.append(img_mask)
            z.append(int(idx_z))

        x = np.asarray(x)
        y = np.asarray(y)
        z = np.asarray(z)
        self.num_test = len(x)


        x = tf.data.Dataset.from_tensor_slices(tf.convert_to_tensor(x, dtype = tf.float32))
        y = tf.data.Dataset.from_tensor_slices(tf.convert_to_tensor(y, dtype = tf.int32))
        z = tf.data.Dataset.from_tensor_slices(tf.convert_to_tensor(z, dtype = tf.int32))

        self.test = tf.data.Dataset.zip((x, y, z))


    def _read_image(self, img_path , flags=cv2.IMREAD_COLOR):

        n = np.fromfile(img_path, dtype = np.uint8)
        img = cv2.imdecode(n, flags)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, dsize = (self.img_size, self.img_size))
        # img = img[16:-16, 16:-16, :] #why? rotateするから？

        return img

プログラム②：テストプログラム

test.py

# -----------------------------------------------
# Import
# -----------------------------------------------
# 01. Built-in
import os

# 02. Third-party modules
import cv2
import numpy as np
import tensorflow as tf

# 03. My Module
from Data_Loader import MVtecADLoader


# Variables
img_size = 224

#Instance
loader= MVtecADLoader(com = 'note', img_size = img_size)

#------------------------------------------------------
# Data Load
#------------------------------------------------------
loader.load(category = 'bottle')

#Train Data and Test Data

train_set = loader.train.batch(batch_size = 16, drop_remainder = True).shuffle(buffer_size = loader.num_train,
                                                                                    reshuffle_each_iteration = True)
test_set = loader.test.batch(batch_size = 1, drop_remainder = False)

print('train set')
for x, y, z in train_set:

    print(x.shape)
    print(y.shape)
    print(z.shape)

print('test set')
for x, y, z in test_set:
    print(x.shape)
    print(y.shape)
    print(z.shape)

結果

train set
(16, 224, 224, 3) #x 
(16, 224, 224) #y
(16,) #z
...

test set
(16, 224, 224, 3)
(16, 224, 224)
(16,)
...