Add files via upload

Author: Dynamo13

dataloader.py

@@ -0,0 +1,31 @@
+import keras
+import numpy as np
+
+class SpinePTXT(keras.utils.Sequence):
+    """Helper to iterate over the data (as Numpy arrays)."""
+
+    def __init__(self, batch_size, img_size, input_img_paths, mask_img_paths,num_classes):
+        self.batch_size = batch_size
+        self.img_size = img_size
+        self.input_img_paths = input_img_paths
+        self.mask_img_paths = mask_img_paths
+        self.num_classes = num_classes
+
+
+    def __len__(self):
+        return len(self.mask_img_paths) // self.batch_size
+
+    def __getitem__(self, idx):
+        """Returns tuple (input, target) correspond to batch #idx."""
+        i = idx * self.batch_size
+        batch_input_img_paths = self.input_img_paths[i : i + self.batch_size]
+        batch_mask_img_paths = self.mask_img_paths[i : i + self.batch_size]
+        x = np.zeros((self.batch_size,) + self.img_size +(3,) , dtype="float32")
+        for j, path in enumerate(batch_input_img_paths):
+            img = np.load(path)
+            x[j]=img
+        y = np.zeros((self.batch_size,) + self.img_size + (self.num_classes,), dtype="uint8")
+        for j, path in enumerate(batch_mask_img_paths):
+            msk = np.load(path)
+            y[j]=msk
+        return x, y

main.py

@@ -0,0 +1,53 @@
+from model import*
+from dataloader import*
+import os
+import random
+
+def main(input_dir, mask_dir, image_height, image_width, image_channel, num_classes, batch_size, epochs, val_samples):
+
+    img_size = (image_height,image_width)
+    input_img_paths = sorted(
+    [
+        os.path.join(input_dir, fname)
+        for fname in os.listdir(input_dir)
+    ]
+    )
+    mask_img_paths = sorted(
+    [
+        os.path.join(mask_dir, fname)
+        for fname in os.listdir(mask_dir)
+    ]
+    )
+    random.Random(1337).shuffle(input_img_paths)
+    random.Random(1337).shuffle(mask_img_paths)
+    train_input_img_paths = input_img_paths[:-val_samples]
+    train_mask_img_paths = mask_img_paths[:-val_samples]
+    val_input_img_paths = input_img_paths[-val_samples:]
+    val_mask_img_paths = mask_img_paths[-val_samples:]
+
+    # Instantiate data Sequences for each split
+    train_gen = SpinePTXT(
+        batch_size, img_size, train_input_img_paths, train_mask_img_paths,num_classes
+    )
+    val_gen = SpinePTXT(batch_size, img_size, val_input_img_paths, val_mask_img_paths,num_classes)
+    
+    model=AW_Net((image_height,image_width,image_channel),num_classes, dropout_rate=0.0, batch_norm=True)
+
+    model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])
+    
+    history = model.fit(train_gen,validation_data=val_gen,epochs=epochs,)
+
+
+if __name__ == "__main__":
+    input_dir = "E:\\BraTS_data\\Image"
+    mask_dir = "E:\\BraTS_data\\Mask"
+    image_height=128
+    image_width=128
+    image_channel=3
+    img_size = (image_height,image_width)
+    num_classes = 4
+    batch_size = 8
+    epochs=200
+    val_samples = 250
+    main(input_dir, mask_dir, image_height, image_width, image_channel, num_classes, batch_size, epochs, val_samples)
+	

model.py

@@ -0,0 +1,77 @@
+import tensorflow as tf
+from keras.models import Model
+from keras.layers import Input
+from model_blocks import*
+
+def AW_Net(input_shape, NUM_CLASSES=4, dropout_rate=0.0, batch_norm=True):
+    '''
+    Attention UNet, 
+    
+    '''
+    # network structure
+    FILTER_NUM = 16 # number of basic filters for the first layer
+    FILTER_SIZE = 3 # size of the convolutional filter
+    UP_SAMP_SIZE = 2 # size of upsampling filters
+    
+    inputs = layers.Input(input_shape, dtype=tf.float32)
+
+    # Downsampling layers
+    # DownRes 1, convolution + pooling
+    conv_128 = conv_block(inputs, FILTER_SIZE, FILTER_NUM, dropout_rate, 1, batch_norm)
+    pool_64 = layers.MaxPooling2D(pool_size=(2,2))(conv_128)
+    # DownRes 2
+    conv_64 = conv_block(pool_64, FILTER_SIZE, 2*FILTER_NUM, dropout_rate, 2, batch_norm)
+    pool_32 = layers.MaxPooling2D(pool_size=(2,2))(conv_64)
+    # DownRes 3
+    conv_32 = conv_block(pool_32, FILTER_SIZE, 4*FILTER_NUM, dropout_rate, 3, batch_norm)
+    pool_16 = layers.MaxPooling2D(pool_size=(2,2))(conv_32)
+    # DownRes 4
+    conv_16 = conv_block(pool_16, FILTER_SIZE, 8*FILTER_NUM, dropout_rate,4, batch_norm)
+    pool_8 = layers.MaxPooling2D(pool_size=(2,2))(conv_16)
+    # DownRes 5, convolution only
+    conv_8 = reg_conv_block(pool_8, FILTER_SIZE, 16*FILTER_NUM, dropout_rate,5, batch_norm)
+
+    # W-net layers
+    gatingw_16 = gating_signal(conv_8, 8*FILTER_NUM, batch_norm)
+    attw_16 = attention_block(conv_16, gatingw_16, 8*FILTER_NUM)
+    upw_16 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(conv_8)
+    upw_16 = layers.concatenate([upw_16, attw_16], axis=3)
+    up_convw_16 = reg_conv_block(upw_16, FILTER_SIZE, 8*FILTER_NUM, dropout_rate,6, batch_norm)
+
+    poolw_8 = layers.MaxPooling2D(pool_size=(2,2))(up_convw_16)
+    ct_16 = layers.concatenate([conv_8, poolw_8], axis=3)
+    convw_16 = reg_conv_block(ct_16, FILTER_SIZE, 16*FILTER_NUM, dropout_rate,7, batch_norm)
+    
+    # UpRes 6, attention gated concatenation + upsampling + double residual convolution
+    gating_16 = gating_signal(convw_16, 8*FILTER_NUM, batch_norm)
+    att_16 = attention_block(up_convw_16, gating_16, 8*FILTER_NUM)
+    up_16 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(convw_16)
+    up_16 = layers.concatenate([up_16, att_16], axis=3)
+    up_conv_16 = conv_block(up_16, FILTER_SIZE, 8*FILTER_NUM, dropout_rate,8, batch_norm)
+    # UpRes 7
+    gating_32 = gating_signal(up_conv_16, 4*FILTER_NUM, batch_norm)
+    att_32 = attention_block(conv_32, gating_32, 4*FILTER_NUM)
+    up_32 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(up_conv_16)
+    up_32 = layers.concatenate([up_32, att_32], axis=3)
+    up_conv_32 = conv_block(up_32, FILTER_SIZE, 4*FILTER_NUM, dropout_rate,9, batch_norm)
+    # UpRes 8
+    gating_64 = gating_signal(up_conv_32, 2*FILTER_NUM, batch_norm)
+    att_64 = attention_block(conv_64, gating_64, 2*FILTER_NUM)
+    up_64 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(up_conv_32)
+    up_64 = layers.concatenate([up_64, att_64], axis=3)
+    up_conv_64 = conv_block(up_64, FILTER_SIZE, 2*FILTER_NUM, dropout_rate,10, batch_norm)
+    # UpRes 9
+    gating_128 = gating_signal(up_conv_64, FILTER_NUM, batch_norm)
+    att_128 = attention_block(conv_128, gating_128, FILTER_NUM)
+    up_128 = layers.UpSampling2D(size=(UP_SAMP_SIZE, UP_SAMP_SIZE), data_format="channels_last")(up_conv_64)
+    up_128 = layers.concatenate([up_128, att_128], axis=3)
+    up_conv_128 = conv_block(up_128, FILTER_SIZE, FILTER_NUM,dropout_rate,11, batch_norm)
+
+    # 1*1 convolutional layers
+    conv_final = layers.Conv2D(NUM_CLASSES, name='conv12', kernel_size=(1,1))(up_conv_128)
+    conv_final = layers.BatchNormalization(axis=3)(conv_final)
+    conv_final = layers.Activation('softmax')(conv_final)  #Change to softmax for multichannel
+
+    # Model integration
+    model = models.Model(inputs, conv_final, name="AW-Net")
+    return model

model_blocks.py

@@ -0,0 +1,86 @@
+from tensorflow.keras import models, layers, regularizers
+from tensorflow.keras import backend as K
+
+def conv_block(x, filter_size, size, dropout,num, batch_norm=False):
+    
+    conv = layers.Conv2D(size, (filter_size, filter_size), padding="same")(x)
+    if batch_norm is True:
+        conv = layers.BatchNormalization(axis=3)(conv)
+    conv = layers.Activation("relu")(conv)
+
+    conv = layers.Conv2D(size, (filter_size, filter_size), padding="same",name="conv"+str(num))(conv)
+    if batch_norm is True:
+        conv = layers.BatchNormalization(axis=3)(conv)
+    conv = layers.Activation("relu")(conv)
+    
+    if dropout > 0:
+        conv = layers.Dropout(dropout)(conv)
+
+    return conv
+
+def reg_conv_block(x, filter_size, size, dropout,num, batch_norm=False):
+    
+    conv = layers.Conv2D(size, (filter_size, filter_size),kernel_regularizer='l1', padding="same")(x)
+    if batch_norm is True:
+        conv = layers.BatchNormalization(axis=3)(conv)
+    conv = layers.Activation("relu")(conv)
+
+    conv = layers.Conv2D(size, (filter_size, filter_size),kernel_regularizer='l1', padding="same",name="conv"+str(num))(conv)
+    if batch_norm is True:
+        conv = layers.BatchNormalization(axis=3)(conv)
+    conv = layers.Activation("relu")(conv)
+    
+    if dropout > 0:
+        conv = layers.Dropout(dropout)(conv)
+
+    return conv
+
+def gating_signal(input, out_size, batch_norm=False):
+    """
+    resize the down layer feature map into the same dimension as the up layer feature map
+    using 1x1 conv
+    :return: the gating feature map with the same dimension of the up layer feature map
+    """
+    x = layers.Conv2D(out_size, (1, 1), padding='same')(input)
+    if batch_norm:
+        x = layers.BatchNormalization()(x)
+    x = layers.Activation('relu')(x)
+    return x
+
+def attention_block(x, gating, inter_shape):
+    shape_x = K.int_shape(x)
+    shape_g = K.int_shape(gating)
+
+# Getting the x signal to the same shape as the gating signal
+    theta_x = layers.Conv2D(inter_shape, (2, 2), strides=(2, 2), padding='same')(x)  # 16
+    shape_theta_x = K.int_shape(theta_x)
+
+# Getting the gating signal to the same number of filters as the inter_shape
+    phi_g = layers.Conv2D(inter_shape, (1, 1), padding='same')(gating)
+    upsample_g = layers.Conv2DTranspose(inter_shape, (3, 3),
+                                 strides=(shape_theta_x[1] // shape_g[1], shape_theta_x[2] // shape_g[2]),
+                                 padding='same')(phi_g)  # 16
+
+    concat_xg = layers.add([upsample_g, theta_x])
+    act_xg = layers.Activation('relu')(concat_xg)
+    psi = layers.Conv2D(1, (1, 1), padding='same')(act_xg)
+    sigmoid_xg = layers.Activation('sigmoid')(psi)
+    shape_sigmoid = K.int_shape(sigmoid_xg)
+    upsample_psi = layers.UpSampling2D(size=(shape_x[1] // shape_sigmoid[1], shape_x[2] // shape_sigmoid[2]))(sigmoid_xg)  # 32
+
+    upsample_psi = repeat_elem(upsample_psi, shape_x[3])
+
+    y = layers.multiply([upsample_psi, x])
+
+    result = layers.Conv2D(shape_x[3], (1, 1), padding='same')(y)
+    result_bn = layers.BatchNormalization()(result)
+    return result_bn
+
+def repeat_elem(tensor, rep):
+    # lambda function to repeat Repeats the elements of a tensor along an axis
+    #by a factor of rep.
+    # If tensor has shape (None, 256,256,3), lambda will return a tensor of shape 
+    #(None, 256,256,6), if specified axis=3 and rep=2.
+
+     return layers.Lambda(lambda x, repnum: K.repeat_elements(x, repnum, axis=3),
+                          arguments={'repnum': rep})(tensor)