modified dataset

Author: matteo6198

config/conf.json

@@ -3,7 +3,7 @@
     "lr":1e-3,
     "batch_size":4,
     "optimizer":"adam",
-    "metrics":"mIoU",
+    "metrics":["mIoU"],
     "loss":"dice",
     "monitor_metric":"val_mIoU"
 }

dataset/Defects/Defects/Image0.jpg dataset/Defects/Image0.jpg

@@ -1,2 +1,3 @@
 tensorflow
-glob
+glob
+numpy

dataset/Defects/Defects/Image1.jpg dataset/Defects/Image1.jpg

@@ -1,8 +1,54 @@
+import json
 import tensorflow as tf
 from glob import glob
 import os
+import numpy as np
+from PIL import Image, ImageDraw
 
-REQUIRED_FOLDERS = ['Defects', 'NoDefects', 'Annotations']
+REQUIRED_FOLDERS = ['Defects', 'NoDefects', 'annotations']
+
+LABEL_DICT = {
+    'BG':0,
+    'HOLE':1,
+    'VERTICAL':2,
+    'HORIZONTAL':3,
+    'SPATTERING':4,
+    'INCANDESCENCE':5
+}
+
+all_defects = []
+
+def map_fn(file:str, save_defects=False):
+    # load image
+    img = tf.keras.utils.load_img(file)
+    shape = img.size
+    # build labels
+    label = np.zeros((shape[0], shape[1]))
+    if 'NoDefects' not in file:
+        # load file
+        annotation_file = file.replace('Defects', 'annotations').replace('jpg', 'json')
+        with open(annotation_file, 'r') as f:
+            annotations = json.load(f)
+        # build annotation image & draw poligons
+        img = Image.fromarray(label)
+        for shape in annotations['shapes']:
+            # use only one label name
+            if shape['label'].upper() == 'VERTICAL DEFECT':
+                shape['label'] = 'VERTICAL'                
+            if shape['label'].upper() == 'SPATTING':
+                shape['label'] = 'SPATTERING'
+
+            points = [(x[0], x[1]) for x in shape['points']]
+            label_id = LABEL_DICT[shape['label'].upper()]
+            # append defects to all defects
+            if save_defects:
+                all_defects.append((label_id, points))
+            # draw poligon on image
+            ImageDraw.Draw(img).polygon(points, fill=label_id)
+        label = np.array(img)
+    label = tf.one_hot(label, len(LABEL_DICT))
+    return (img, label)
+    
 
 class AMDdataset():
     '''Additive Manufactoring dataset class'''
@@ -19,3 +65,23 @@ def build(self):
             raise FileNotFoundError(f'Directory {self.path} does not contain correct folders. It must contains {REQUIRED_FOLDERS}')
 
         # TODO split the dataset and load it into 3 tf.dataset: self.train, self.val, self.test
+        train = []
+        test = []
+        val = []
+
+        # split folders
+        for f in folders:
+            files = np.array(glob(os.path.join(self.path, f,'*.jpg')))
+            if len(files) == 0:
+                continue
+            n = len(files) // 3
+            idx = np.random.permutation(np.arange(len(files)))
+
+            test.extend(files[idx[:n]])
+            val.extend(files[idx[n:n*2]])
+            train.extend(files[idx[n*2:]])
+
+        train = [map_fn(x, True) for x in train]
+        test = [map_fn(x) for x in test]
+        val = [map_fn(x, True) for x in val]
+

dataset/Defects/Defects/Image10.jpg dataset/Defects/Image10.jpg

@@ -1,5 +1,5 @@
 import tensorflow as tf
-from tensorflow.keras.applications.mobilenet_v3 import MobileNetV3
+from tensorflow.keras.applications import MobileNetV3Small
 
 class ConvBlock(tf.keras.layers.Layer):
     KERNEL_SIZE = 3
@@ -88,21 +88,49 @@ def call(self, inputs):
         x = self.base_model(inputs)
         x = self.result(x)
         return x
+
 class GAN(tf.keras.Model):
     def __init__(
         self,
         name="gan",
-        input_shape=(1024, 800, 1),
+        # input_shape=(1024, 800, 1),
         **kwargs
     ):
         super(GAN, self).__init__(name=name,**kwargs)
         self.generator = Generator()
         self.discriminator = Discriminator()
 
     def train_step(self, inputs):
-        # TODO
-        return 0
+        inputs_with_defects = inputs    # TODO add function to insert defects
+        # train the discriminator
+        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
+            y_true_pred = self.discriminator(inputs, training=True)
+            gen_images = self.generator(inputs_with_defects, training=True)
+            y_fake_pred = self.discriminator(gen_images, training=True)
+
+            y_true = tf.ones_like(y_true_pred)
+            y_fake = tf.zeros_like(y_fake_pred)
+
+            # (the loss function is configured in compile())
+            loss_d = self.loss['discriminator'](tf.concat(y_true, y_fake), tf.concat(y_true_pred, y_fake_pred), regularization_losses=self.losses)
+            loss_g = self.loss['generator'](y_true, y_fake_pred, regularization_losses=self.losses)
+            loss = tf.reduce_sum(loss_g, 0.5 * loss_d)
+
+        # compute gradients
+        gen_gradients = gen_tape.gradient(loss_g, self.generator.trainable_variables)
+        disc_gradients = disc_tape.gradient(loss_d, self.discriminator.trainable_variables)
+
+        # apply gradients
+        # (the optimizer should be configured in compile())
+        self.optimizer['generator'].apply_gradients(zip(gen_gradients, self.generator.trainable_variables))
+        self.optimizer['discriminator'].apply_gradients(zip(disc_gradients, self.discriminator.trainable_variables))
+
+        # compute metrics
+        out = {m.name: m.result() for m in self.metrics}
+        out['loss'] = loss
+        return out
 
     def call(self, inputs):
-        # TODO
-        return {}
+        gen_out = self.generator(inputs)
+        disc_out = self.discriminator(gen_out)
+        return [gen_out, disc_out]

dataset/Defects/Defects/Image11.jpg dataset/Defects/Image11.jpg

@@ -1,27 +1,34 @@
 from telnetlib import GA
-import tensorflow
+import tensorflow as tf
 from glob import glob
 from utils.parser import parse_arguments
 from datasets.dataset import AMDdataset
 from network.unet import UNet
 from network.gan import GAN
 from utils.train_utils import TrainWrapper
 from utils.conf_reader import read_conf
+import numpy as np
+
 
 MODELS={
     'unet':UNet,
     'gan':GAN
 }
 
 if __name__ == '__main__':
+    # make deterministic
+    np.random.seed(0)
+    tf.keras.utils.set_random_seed(1)
+    tf.config.experimental.enable_op_determinism()
     # parsing the arguments
     args = parse_arguments()
     # read config file
     conf = read_conf(args.conf)
     # opening the dataset
     datasets = AMDdataset(args.dataset_folder)
+    datasets.build()
     # define the model
-    net = MODELS(args.model)
+    net = MODELS[args.model]()
     # train
     train_wrapper = TrainWrapper(net, conf=conf, train_dataset=datasets.train, val_dataset=datasets.val, test_dataset=datasets.test)
 

dataset/Defects/Defects/Image12.jpg dataset/Defects/Image12.jpg

@@ -1,6 +1,6 @@
 import json
 import tensorflow as tf
-from losses import LOSSES
+from utils.losses import LOSSES
 
 OPTIMIZERS = {
     'adam':tf.keras.optimizers.Adam,

dataset/Defects/Defects/Image13.jpg dataset/Defects/Image13.jpg

@@ -17,5 +17,6 @@ def call(self, y_true, y_pred):
 
 LOSSES = {
     'dice':DiceLoss,
-    'cross_entropy':tf.keras.losses.CategoricalCrossEntropy
+    'categorical_cross_entropy':tf.keras.losses.CategoricalCrossentropy,
+    'binary_crossentropy':tf.keras.losses.BinaryCrossentropy
 }