Add training code

.pylintrc

@@ -59,7 +59,7 @@ confidence=
 #
 # Kubeflow disables string-interpolation because we are starting to use f
 # style strings
-disable=useless-suppression,suppressed-message,missing-docstring,invalid-name,no-member,locally-disabled,fixme,import-error,too-many-locals,no-name-in-module,too-many-instance-attributes,logging-fstring-interpolation
+disable=too-many-arguments,unused-argument,inconsistent-return-statements,useless-suppression,suppressed-message,missing-docstring,invalid-name,no-member,locally-disabled,fixme,import-error,too-many-locals,no-name-in-module,too-many-instance-attributes,logging-fstring-interpolation
 
 
 [REPORTS]

config.yml

@@ -13,19 +13,19 @@ preprocessing:
   img_scaling: [3, 3] # downsampling in preprocessing
 
 train:
-  learning_rate: 0.003
+  learning_rate: 0.01
   batch_size: 64
   epochs: 100
   shuffle: True
   do_augmentation: True
   reduceLROnPlat:
     factor: 0.33
-    patience: 2
+    patience: 8
     min_delta: 0.0001
     min_lr: 0.00000001
     cooldown: 0
   EarlyStopping:
-    patience: 10
+    patience: 20
   checkpoint:
     save_best_only: True
     save_weights_only: True

core/data/preprocessing.py

@@ -1,9 +1,9 @@
 import os
-from typing import Optional, Tuple
+from typing import Tuple
 
 import numpy as np
 import pandas as pd
-import albumentations as A
+import tensorflow as tf
 from skimage.io import imread
 from sklearn.model_selection import train_test_split
 
@@ -42,6 +42,9 @@ def __len__(self):
         return len(self.in_df) // self.batch_size
 
     def __getitem__(self, index):
+        if self.shuffle_batches:
+            np.random.shuffle(self.all_batches)
+
         for c_img_id, c_masks in self.all_batches:
             rgb_path = os.path.join(self.img_dir, c_img_id)
             c_img = imread(rgb_path)
@@ -60,11 +63,11 @@ def __getitem__(self, index):
             self.out_mask.append(c_mask)
 
             if len(self.out_rgb) >= self.batch_size:
-                batch_rgb = tf.convert_to_tensor(np.stack(self.out_rgb, 0) / 255.0, dtype=tf.float32)
-                batch_mask = tf.convert_to_tensor(np.stack(self.out_mask, 0), dtype=tf.float32)
-                self.out_rgb.clear(), self.out_mask.clear()
+                batch_rgb = np.stack(self.out_rgb, 0) / 255.0
+                batch_mask = np.stack(self.out_mask, 0)
+                self.out_rgb, self.out_mask = [], []
                 return batch_rgb, batch_mask
 
-        def on_epoch_end(self):
-            if self.shuffle_batches:
-                np.random.shuffle(self.all_batches)
+    def on_epoch_end(self):
+        if self.shuffle_batches:
+            np.random.shuffle(self.all_batches)

core/data/read_data.py

@@ -20,6 +20,7 @@ def get_metadata(
     unique_img_ids['file_size_kb'] = unique_img_ids['ImageId'].map(
         lambda c_img_id: os.stat(os.path.join(train_image_dir, c_img_id)).st_size / 1024, meta=('ImageId', int)
     )
+    unique_img_ids = unique_img_ids[unique_img_ids['file_size_kb'] > 50] # keep only +50kb files
 
     unique_img_ids = unique_img_ids.compute()
     masks = masks.compute()

core/inference.py

@@ -13,6 +13,8 @@ def create_full_res_model(model, img_scaling: tuple[int, int]):
         fullres_model.add(layers.AvgPool2D(img_scaling, input_shape=(None, None, 3)))
         fullres_model.add(model)
         fullres_model.add(layers.UpSampling2D(img_scaling))
+        return fullres_model
+
     return model
 
 

core/losses.py

@@ -1,11 +1,10 @@
-import numpy as np
-import tensorflow.keras.backend as K
+import tensorflow as tf
+from tensorflow.keras.losses import binary_crossentropy
+from core.metrics import dice_coef
 
 
-def IoU(y_true, y_pred, eps=1e-6):
-    if np.max(y_true) == 0.:
-        return IoU(1-y_true, 1-y_pred) ## empty image; calc IoU of zeros
+def dice_p_bce(y_true, y_pred):
+    y_true = tf.cast(y_true, dtype=tf.float32)
+    y_pred = tf.cast(y_pred, dtype=tf.float32)
 
-    intersection = K.sum(y_true * y_pred, axis=[1, 2, 3])
-    union = K.sum(y_true, axis=[1, 2, 3]) + K.sum(y_pred, axis=[1, 2, 3]) - intersection
-    return -K.mean( (intersection + eps) / (union + eps), axis=0)
+    return 1e-3 * binary_crossentropy(y_true, y_pred) - dice_coef(y_true, y_pred)

core/metrics.py

@@ -1,9 +1,22 @@
 import tensorflow as tf
+import tensorflow.keras.backend as K
 
 
-def dice(y_true, y_pred):
-    y_true = tf.cast(y_true, tf.float32)
-    y_pred = tf.math.sigmoid(y_pred)
-    numerator = 2 * tf.reduce_sum(y_true * y_pred)
-    denominator = tf.reduce_sum(y_true + y_pred)
-    return numerator / denominator
+def dice_coef(y_true, y_pred, smooth=1):
+    y_true = tf.cast(y_true, dtype=tf.float32)
+    y_pred = tf.cast(y_pred, dtype=tf.float32)
+
+    intersection = K.sum(y_true * y_pred, axis=[1, 2, 3])
+    union = K.sum(y_true, axis=[1, 2, 3]) + K.sum(y_pred, axis=[1, 2, 3])
+    return K.mean((2. * intersection + smooth) / (union + smooth), axis=0)
+
+
+def POD(y_true, y_pred):
+    y_true = tf.cast(y_true, dtype=tf.float32)
+    y_pred = tf.cast(y_pred, dtype=tf.float32)
+
+    y_true_pos = K.flatten(y_true)
+    y_pred_pos = K.flatten(y_pred)
+    true_pos = K.sum(y_true_pos * y_pred_pos)
+    false_neg = K.sum(y_true_pos * (1 - y_pred_pos))
+    return true_pos / (true_pos + false_neg)

core/model/UNet.py

@@ -2,7 +2,7 @@
 from core.model.blocks import upsample_conv, upsample_simple, encoder_block, decoder_block
 
 
-def init_model(config: dict, input_shape: tuple=(None, 256, 256, 3)) -> models.Model:
+def init_model(config: dict, input_shape: tuple=(256, 256, 3)) -> models.Model:
     n_filters = config['n_filters']
 
     if config['upsample_mode'] == 'DECONV':
@@ -23,16 +23,18 @@ def init_model(config: dict, input_shape: tuple=(None, 256, 256, 3)) -> models.M
     enc2 = encoder_block(enc1[0], n_filters * 2)
     enc3 = encoder_block(enc2[0], n_filters * 4)
     enc4 = encoder_block(enc3[0], n_filters * 8)
+    enc5 = encoder_block(enc4[0], n_filters * 16)
 
-    c5 = layers.Conv2D(n_filters * 16, (3, 3), activation='relu', padding='same')(enc4[0])
-    c5 = layers.Conv2D(n_filters * 16, (3, 3), activation='relu', padding='same')(c5)
+    c6 = layers.Conv2D(n_filters * 32, (3, 3), activation='relu', padding='same')(enc5[0])
+    c6 = layers.Conv2D(n_filters * 32, (3, 3), activation='relu', padding='same')(c6)
 
-    dec1 = decoder_block(c5, enc4[1], n_filters * 8, upsample)
-    dec2 = decoder_block(dec1, enc3[1], n_filters * 4, upsample)
-    dec3 = decoder_block(dec2, enc2[1], n_filters * 2, upsample)
-    dec4 = decoder_block(dec3, enc1[1], n_filters, upsample)
+    dec1 = decoder_block(c6, enc5[1], n_filters * 16, upsample)
+    dec2 = decoder_block(dec1, enc4[1], n_filters * 8, upsample)
+    dec3 = decoder_block(dec2, enc3[1], n_filters * 4, upsample)
+    dec4 = decoder_block(dec3, enc2[1], n_filters * 2, upsample)
+    dec5 = decoder_block(dec4, enc1[1], n_filters, upsample)
 
-    dec = layers.Conv2D(1, (1, 1), activation='sigmoid')(dec4)
+    dec = layers.Conv2D(1, (1, 1), activation='sigmoid')(dec5)
 
     if config['net_scaling'] is not None:
         dec = layers.UpSampling2D(config['net_scaling'])(dec)

core/model/blocks.py

@@ -24,8 +24,10 @@ def upsample_simple(
 
 def encoder_block(prev_layer_inputs: tf.Tensor, n_filters: int) -> Tuple[tf.Tensor, tf.Tensor]:
     conv = layers.Conv2D(n_filters, (3, 3), activation='relu', padding='same')(prev_layer_inputs)
+    conv = layers.BatchNormalization()(conv)
     skip_connection = layers.Conv2D(n_filters, (3, 3), activation='relu', padding='same')(conv)
     pool = layers.MaxPooling2D((2, 2))(skip_connection)
+    pool = layers.BatchNormalization()(pool)
     return pool, skip_connection
 
 
@@ -35,5 +37,7 @@ def decoder_block(
     up = upsample(n_filters, (2, 2), strides=(2, 2), padding='same')(prev_layer_input)
     up = layers.concatenate([up, skip_layer_input])
     up = layers.Conv2D(n_filters, (3, 3), activation='relu', padding='same')(up)
+    up = layers.BatchNormalization()(up)
     up = layers.Conv2D(n_filters, (3, 3), activation='relu', padding='same')(up)
+    up = layers.BatchNormalization()(up)
     return up

core/train.py

@@ -0,0 +1,27 @@
+from tensorflow.keras.optimizers import AdamW
+from core.data.preprocessing import SemanticSegmentationDataGenerator
+from core.losses import dice_p_bce
+from core.metrics import POD, dice_coef
+
+
+def train(model, img_dir, config, img_scaling, callbacks, train_df, valid_df, transform):
+    batch_size = config['batch_size']
+    model.compile(optimizer=AdamW(learning_rate=config['learning_rate']), loss=dice_p_bce,
+                      metrics=['binary_accuracy', dice_coef, POD], run_eagerly=True)
+
+    train_data_generator = SemanticSegmentationDataGenerator(
+        train_df, img_dir, batch_size, img_scaling, config['do_augmentation'], transform
+    )
+    val_data_generator = SemanticSegmentationDataGenerator(valid_df, img_dir, batch_size, img_scaling)
+    loss_history = model.fit(
+        train_data_generator,
+        steps_per_epoch=len(train_data_generator) // train_data_generator.batch_size,
+        epochs=config['epochs'],
+        validation_data=val_data_generator,
+        validation_steps=len(val_data_generator) // val_data_generator.batch_size,
+        callbacks=callbacks,
+        verbose=1,
+        workers=1,
+    )
+
+    return model, loss_history