{"metadata":{"colab":{"provenance":[],"machine_shape":"hm","gpuType":"T4"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU","kaggle":{"accelerator":"none","dataSources":[{"sourceId":87230,"databundleVersionId":9911932,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Mount Google Drive\n# This is just for Google Colab\nfrom google.colab import drive\ndrive.mount('/content/drive')\n\n# Set the directory at MyDrive\n# This is just when running on Google colab\nimport os\nos.chdir('/content/drive/MyDrive/')","metadata":{"id":"CqGeC1HjUmmP","outputId":"9ddf2009-9b67-4950-b7fc-6498eaae98a6"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nfrom matplotlib import pyplot as plt\nimport cv2\nfrom sklearn.model_selection import train_test_split\nimport copy\nimport tensorflow as tf\nfrom tensorflow.keras.models import Model, load_model\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint\nfrom scipy import interpolate\nfrom scipy.interpolate import NearestNDInterpolator\nfrom tensorflow.keras.models import Sequential\n\n\n\n# Load the 3D array with the Bscans\n# np.shape(labels) = (number of samples, 230, 230)\n# Notice that in this script we try to fill missing values so the initial \"clean\" BScan is the ground truth that we try to recover\nlabels = np.load('Bscan_proc.npy')\n\n# The idea here is to \"corrupt\" the clean Bscans with missing values, and then try to train a model to recover the missing values\n# Initialise the training data with zeros\ntrain = np.zeros(np.shape(labels))\n\n\nfor i, c in enumerate(labels):\n    # Go through all the Bscans and normalise them from -1 to 1\n    c = c/np.max(np.abs(c))\n    labels[i] = c\n\n    # Fill the training data only with the values of the A-scans for every 10 traces.\n    train[i,:,0:-1:10] = c[:,0:-1:10]\n    # plt.imshow(train[i])\n    # plt.show()\n    # plt.imshow(labels[i])\n    # plt.show()\n\n\n# labels: the clean Bscans that we try to recover i.e. ground truth\n# train : the \"corrupted\" Bscans with the missing values\n","metadata":{"id":"50m57TAmUv5z"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport tensorflow as tf\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, Concatenate, Cropping2D, ZeroPadding2D\nfrom tensorflow.keras.losses import MeanSquaredError\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint\nfrom sklearn.model_selection import train_test_split\n\n\n# Define U-net architecture\ndef unet_model(input_size=(230, 230, 1)):\n    # This is a vanilla typical U-net that takes an image as input with 230x230 dimensions and gives another image as an output with 230x230 dimensions\n    # The hyperparameters are not tuned.\n\n    inputs = Input(input_size)\n    # Encoder\n    c1 = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)\n    c1 = Conv2D(64, (3, 3), activation='relu', padding='same')(c1)\n    p1 = MaxPooling2D((2, 2))(c1)\n\n    c2 = Conv2D(128, (3, 3), activation='relu', padding='same')(p1)\n    c2 = Conv2D(128, (3, 3), activation='relu', padding='same')(c2)\n    p2 = MaxPooling2D((2, 2))(c2)\n\n    c3 = Conv2D(256, (3, 3), activation='relu', padding='same')(p2)\n    c3 = Conv2D(256, (3, 3), activation='relu', padding='same')(c3)\n    p3 = MaxPooling2D((2, 2))(c3)\n\n    c4 = Conv2D(512, (3, 3), activation='relu', padding='same')(p3)\n    c4 = Conv2D(512, (3, 3), activation='relu', padding='same')(c4)\n    p4 = MaxPooling2D((2, 2))(c4)\n\n    c5 = Conv2D(1024, (3, 3), activation='relu', padding='same')(p4)\n    c5 = Conv2D(1024, (3, 3), activation='relu', padding='same')(c5)\n    # Decoder\n    u6 = UpSampling2D((2, 2))(c5)\n    u6 = Concatenate()([u6, c4])\n    c6 = Conv2D(512, (3, 3), activation='relu', padding='same')(u6)\n    c6 = Conv2D(512, (3, 3), activation='relu', padding='same')(c6)\n\n    u7 = UpSampling2D((2, 2))(c6)\n\n    u7 = ZeroPadding2D(padding = ((1,0),(1,0)))(u7)\n    u7 = Concatenate()([u7, c3])\n    c7 = Conv2D(256, (3, 3), activation='relu', padding='same')(u7)\n    c7 = Conv2D(256, (3, 3), activation='relu', padding='same')(c7)\n    u8 = UpSampling2D((2, 2))(c7)\n\n    u8 = ZeroPadding2D(padding = ((1,0),(1,0)))(u8)\n\n    u8 = Concatenate()([u8, c2])\n\n    c8 = Conv2D(128, (3, 3), activation='relu', padding='same')(u8)\n    c8 = Conv2D(128, (3, 3), activation='relu', padding='same')(c8)\n    u9 = UpSampling2D((2, 2))(c8)\n\n\n\n    u9 = Concatenate()([u9, c1])\n    c9 = Conv2D(64, (3, 3), activation='relu', padding='same')(u9)\n    c9 = Conv2D(64, (3, 3), activation='relu', padding='same')(c9)\n\n    outputs = Conv2D(1, (1, 1), activation='linear')(c9)\n    model = Model(inputs=[inputs], outputs=[outputs])\n    return model","metadata":{"id":"nQvVljqYVTEE"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Initiate the model\nmodel = unet_model()\n# Print hte summary of the Unet\nmodel.summary()","metadata":{"id":"U9nTE8_Eo_V4","outputId":"2a2f0dfd-d544-47fe-ae25-b4b44acf01a0"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# We add an extra dimension to train and labels to be consistent with tensorflow image inputs\nTrain = train[...,np.newaxis]\nLabels = labels[...,np.newaxis]\n\n# Divide our data to train and test. 10% of the data are used for training purposes\nx_train, x_test, y_train, y_test = train_test_split(Train, Labels, test_size = 0.1)\n\n# Mean Absolute Error gives good results as well\nmodel.compile(optimizer='adam', loss='mean_squared_error', metrics=['mae'])\n\n# The model with the best performance in the validation set will be saved in Fill_Missing.hp5\np1 = EarlyStopping(monitor=\"val_loss\", mode= \"min\", verbose=1, patience=30)\np2 = ModelCheckpoint('Fill_Missing.hp5', monitor=\"val_loss\", mode=\"min\", save_best_only=True )\n\n# The learning rate of adam can be reduced after some iterations.\nmodel.optimizer.learning_rate = 1e-4\n\n# Train the model using 10% validation for 100 epohcs with a 16 batch size\nmodel.fit(x_train, y_train, validation_split=0.1, epochs=100, batch_size=16, callbacks=[p1, p2])\n\n","metadata":{"id":"sECAj2JnWRcu","outputId":"45092555-7f71-4d47-d736-302bfbe0649f"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n# Once the model is trained we can call it in saved_model and use it to make predictions\nsaved_model = load_model('Fill_Missing.hp5')\n\n# x_test are unknown data that were not included in the training\npredictions = saved_model.predict(x_test)\n\n# Plot the results\nfor i, c in enumerate(predictions):\n    plt.subplot(1,3,1)\n    plt.imshow(x_test[i],cmap = 'bone')\n    plt.axis('off')\n    plt.title('Original')\n    plt.subplot(1,3,2)\n    plt.imshow(c, cmap='bone')\n    plt.axis('off')\n    plt.title('ML')\n    plt.title('Interpolated')\n    plt.subplot(1,3,3)\n    plt.imshow(y_test[i], cmap = 'bone')\n    plt.axis('off')\n    plt.title('Ground Truth')\n\n    plt.show()\n\n","metadata":{"id":"FcUlGcbCXIza","outputId":"829c5543-6cb2-46b2-bd43-ba7f934c60ae"},"execution_count":null,"outputs":[]}]}