はじめに 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,) ...