@@ -35,7 +35,7 @@ def map_fn(file:str, save_defects=False):
3535 for shape in annotations ['shapes' ]:
3636 # use only one label name
3737 if shape ['label' ].upper () == 'VERTICAL DEFECT' :
38- shape ['label' ] = 'VERTICAL'
38+ shape ['label' ] = 'VERTICAL'
3939 if shape ['label' ].upper () == 'SPATTING' :
4040 shape ['label' ] = 'SPATTERING'
4141
@@ -57,7 +57,7 @@ def make_dataset(tuples):
5757 return tf .data .Dataset .from_tensor_slices ((x , y ))
5858
5959class AMDdataset ():
60- ''' Additive Manufactoring dataset class'''
60+ """ Additive Manufactoring dataset class"""
6161
6262 def __init__ (self , path , image_shape = (1280 , 1024 , 3 )):
6363 self .path = path
@@ -70,7 +70,6 @@ def build(self):
7070 if len (folders ) != len (REQUIRED_FOLDERS ):
7171 raise FileNotFoundError (f'Directory { self .path } { REQUIRED_FOLDERS } )
7272
73- # TODO split the dataset and load it into 3 tf.dataset: self.train, self.val, self.test
7473 train = []
7574 test = []
7675 val = []
@@ -80,19 +79,16 @@ def build(self):
8079 files = np .array (glob (os .path .join (self .path , f ,'*.jpg' )))
8180 if len (files ) == 0 :
8281 continue
83- n = len (files ) // 3
82+ n = len (files ) // 9
8483 idx = np .random .permutation (np .arange (len (files )))
8584
86- test .extend (files [idx [:n ]])
87- val .extend (files [idx [n : n * 2 ]])
88- train .extend (files [idx [n * 2 :]])
85+ test .extend (files [idx [:3 * n ]]) # 3/9
86+ val .extend (files [idx [3 * n : 5 * n ]]) # 2/9
87+ train .extend (files [idx [5 * n :]]) # 4/9
8988
9089 train = [map_fn (x , True ) for x in train ]
91- test = [map_fn (x ) for x in test ]
92- val = [map_fn (x , True ) for x in val ]
93-
94- # img, mask = train[0]
95- # apply_mask_on_image(img, mask, "../visualize")
90+ test = [map_fn (x ) for x in test ]
91+ val = [map_fn (x , True ) for x in val ]
9692
9793 # build the datasets
9894 self .train = make_dataset (train )
0 commit comments