image_denoising.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. from keras.layers import Input, Conv3D, MaxPooling3D, UpSampling3D
  2. from keras.models import Model
  3. from keras import backend as K
  4. from keras.callbacks import TensorBoard
  5. import pandas as pd
  6. import numpy as np
  7. from sklearn.utils import shuffle
  8. from modules.utils import config as cfg
  9. import cv2
  10. import argparse
  11. def generate_model(input_shape=(3, 200, 200)):
  12. print(input_shape)
  13. input_img = Input(shape=input_shape) # adapt this if using `channels_first` image data format
  14. x = Conv3D(32, (1, 3, 3), activation='relu', padding='same')(input_img)
  15. x = MaxPooling3D((1, 2, 2), padding='same')(x)
  16. x = Conv3D(32, (1, 3, 3), activation='relu', padding='same')(x)
  17. encoded = MaxPooling3D((1, 2, 2), padding='same')(x)
  18. x = Conv3D(32, (1, 3, 3), activation='relu', padding='same')(encoded)
  19. x = UpSampling3D((1, 2, 2))(x)
  20. x = Conv3D(32, (1, 3, 3), activation='relu', padding='same')(x)
  21. x = UpSampling3D((1, 2, 2))(x)
  22. decoded = Conv3D(1, (1, 3, 3), activation='sigmoid', padding='same')(x)
  23. autoencoder = Model(input_img, decoded)
  24. # TODO : check if
  25. autoencoder.compile(optimizer='adadelta', loss='binary_crossentropy')
  26. return autoencoder
  27. def main():
  28. # load params
  29. parser = argparse.ArgumentParser(description="Train Keras model and save it into .json file")
  30. parser.add_argument('--data', type=str, help='dataset filename prefix (without .train and .test)', required=True)
  31. parser.add_argument('--output', type=str, help='output file name desired for model (without .json extension)', required=True)
  32. parser.add_argument('--batch_size', type=int, help='batch size used as model input', default=cfg.keras_batch)
  33. parser.add_argument('--epochs', type=int, help='number of epochs used for training model', default=cfg.keras_epochs)
  34. args = parser.parse_args()
  35. p_data_file = args.data
  36. p_output = args.output
  37. p_batch_size = args.batch_size
  38. p_epochs = args.epochs
  39. # load data from `p_data_file`
  40. ########################
  41. # 1. Get and prepare data
  42. ########################
  43. print("Preparing data...")
  44. dataset_train = pd.read_csv(p_data_file + '.train', header=None, sep=";")
  45. dataset_test = pd.read_csv(p_data_file + '.test', header=None, sep=";")
  46. print("Train set size : ", len(dataset_train))
  47. print("Test set size : ", len(dataset_test))
  48. # default first shuffle of data
  49. dataset_train = shuffle(dataset_train)
  50. dataset_test = shuffle(dataset_test)
  51. print("Reading all images data...")
  52. # getting number of chanel
  53. n_channels = len(dataset_train[1][1].split('::'))
  54. print("Number of channels : ", n_channels)
  55. img_width, img_height = cfg.keras_img_size
  56. # specify the number of dimensions
  57. if K.image_data_format() == 'channels_first':
  58. if n_channels > 1:
  59. input_shape = (1, n_channels, img_width, img_height)
  60. else:
  61. input_shape = (n_channels, img_width, img_height)
  62. else:
  63. if n_channels > 1:
  64. input_shape = (1, img_width, img_height, n_channels)
  65. else:
  66. input_shape = (img_width, img_height, n_channels)
  67. # `:` is the separator used for getting each img path
  68. if n_channels > 1:
  69. dataset_train[1] = dataset_train[1].apply(lambda x: [cv2.imread(path, cv2.IMREAD_GRAYSCALE) for path in x.split('::')])
  70. dataset_test[1] = dataset_test[1].apply(lambda x: [cv2.imread(path, cv2.IMREAD_GRAYSCALE) for path in x.split('::')])
  71. else:
  72. dataset_train[1] = dataset_train[1].apply(lambda x: cv2.imread(x, cv2.IMREAD_GRAYSCALE))
  73. dataset_test[1] = dataset_test[1].apply(lambda x: cv2.imread(x, cv2.IMREAD_GRAYSCALE))
  74. x_dataset_train = dataset_train[1].apply(lambda x: np.array(x).reshape(input_shape))
  75. x_dataset_test = dataset_test[1].apply(lambda x: np.array(x).reshape(input_shape))
  76. y_dataset_train = dataset_train[0].apply(lambda x: cv2.imread(x).reshape(input_shape))
  77. y_dataset_test = dataset_test[0].apply(lambda x: cv2.imread(x).reshape(input_shape))
  78. # format correct data
  79. x_data_train = np.array([item for item in x_dataset_train.values])
  80. #x_data_train = np.array(x_dataset_train.values)
  81. x_data_test = np.array([item for item in x_dataset_test.values])
  82. #x_data_test = np.array(x_dataset_test.values)
  83. y_data_train = np.array([item for item in y_dataset_train.values])
  84. #y_data_train = np.array(y_dataset_train.values)
  85. y_data_test = np.array([item for item in y_dataset_test.values])
  86. #y_data_test = np.array(y_dataset_test.values)
  87. # load model
  88. autoencoder = generate_model(input_shape)
  89. # tensorboard --logdir=/tmp/autoencoder
  90. autoencoder.fit(x_data_train, y_data_train,
  91. epochs=100,
  92. batch_size=32,
  93. shuffle=True,
  94. validation_data=(x_data_test, y_data_test),
  95. callbacks=[TensorBoard(log_dir='/tmp/autoencoder', histogram_freq=0, write_graph=False)])
  96. # save model
  97. if __name__ == "__main__":
  98. main()