From b774408b73ff8f7ddf9bf2a45d20abc257d93ad1 Mon Sep 17 00:00:00 2001 From: jsta Date: Wed, 6 Jul 2022 11:24:57 -0600 Subject: [PATCH] make dwm installable --- .gitignore | 7 +- deepwatermap/__init__.py | 0 .../deepwatermap.py | 148 +++++----- inference.py => deepwatermap/inference.py | 144 ++++----- metrics.py => deepwatermap/metrics.py | 80 ++--- trainer.py => deepwatermap/trainer.py | 274 +++++++++--------- environment.yml | 8 + setup.py | 26 ++ 8 files changed, 365 insertions(+), 322 deletions(-) create mode 100644 deepwatermap/__init__.py rename deepwatermap.py => deepwatermap/deepwatermap.py (97%) rename inference.py => deepwatermap/inference.py (85%) rename metrics.py => deepwatermap/metrics.py (97%) rename trainer.py => deepwatermap/trainer.py (97%) create mode 100644 environment.yml create mode 100644 setup.py diff --git a/.gitignore b/.gitignore index e928de1..8cc510c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,9 @@ # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] -*$py.class \ No newline at end of file +*$py.class +build +*.egg-info +checkpoints + +sample_data/*.tif \ No newline at end of file diff --git a/deepwatermap/__init__.py b/deepwatermap/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/deepwatermap.py b/deepwatermap/deepwatermap.py similarity index 97% rename from deepwatermap.py rename to deepwatermap/deepwatermap.py index edf39cd..7cf1977 100644 --- a/deepwatermap.py +++ b/deepwatermap/deepwatermap.py @@ -1,75 +1,75 @@ -''' Implementation of DeepWaterMapV2. - -The model architecture is explained in: -L.F. Isikdogan, A.C. Bovik, and P. Passalacqua, -"Seeing Through the Clouds with DeepWaterMap," IEEE GRSL, 2019. -''' - -import tensorflow as tf - -def model(min_width=4): - inputs = tf.keras.layers.Input(shape=[None, None, 6]) - - def conv_block(x, num_filters, kernel_size, stride=1, use_relu=True): - x = tf.keras.layers.Conv2D( - filters=num_filters, - kernel_size=kernel_size, - kernel_initializer='he_uniform', - strides=stride, - padding='same', - use_bias=False)(x) - x = tf.keras.layers.BatchNormalization()(x) - if use_relu: - x = tf.keras.layers.Activation('relu')(x) - return x - - def downscaling_unit(x): - num_filters = int(x.get_shape()[-1]) * 4 - x_1 = conv_block(x, num_filters, kernel_size=5, stride=2) - x_2 = conv_block(x_1, num_filters, kernel_size=3, stride=1) - x = tf.keras.layers.Add()([x_1, x_2]) - return x - - def upscaling_unit(x): - num_filters = int(x.get_shape()[-1]) // 4 - x = tf.keras.layers.Lambda(lambda x: tf.nn.depth_to_space(x, 2))(x) - x_1 = conv_block(x, num_filters, kernel_size=3) - x_2 = conv_block(x_1, num_filters, kernel_size=3) - x = tf.keras.layers.Add()([x_1, x_2]) - return x - - def bottleneck_unit(x): - num_filters = int(x.get_shape()[-1]) - x_1 = conv_block(x, num_filters, kernel_size=3) - x_2 = conv_block(x_1, num_filters, kernel_size=3) - x = tf.keras.layers.Add()([x_1, x_2]) - return x - - # model flow - skip_connections = [] - num_filters = min_width - - # first layer - x = conv_block(inputs, num_filters, kernel_size=1, use_relu=False) - skip_connections.append(x) - - # encoder - for i in range(4): - x = downscaling_unit(x) - skip_connections.append(x) - - # bottleneck - x = bottleneck_unit(x) - - # decoder - for i in range(4): - x = tf.keras.layers.Add()([x, skip_connections.pop()]) - x = upscaling_unit(x) - - # last layer - x = tf.keras.layers.Add()([x, skip_connections.pop()]) - x = conv_block(x, 1, kernel_size=1, use_relu=False) - x = tf.keras.layers.Activation('sigmoid')(x) - - model = tf.keras.Model(inputs=inputs, outputs=x) +''' Implementation of DeepWaterMapV2. + +The model architecture is explained in: +L.F. Isikdogan, A.C. Bovik, and P. Passalacqua, +"Seeing Through the Clouds with DeepWaterMap," IEEE GRSL, 2019. +''' + +import tensorflow as tf + +def model(min_width=4): + inputs = tf.keras.layers.Input(shape=[None, None, 6]) + + def conv_block(x, num_filters, kernel_size, stride=1, use_relu=True): + x = tf.keras.layers.Conv2D( + filters=num_filters, + kernel_size=kernel_size, + kernel_initializer='he_uniform', + strides=stride, + padding='same', + use_bias=False)(x) + x = tf.keras.layers.BatchNormalization()(x) + if use_relu: + x = tf.keras.layers.Activation('relu')(x) + return x + + def downscaling_unit(x): + num_filters = int(x.get_shape()[-1]) * 4 + x_1 = conv_block(x, num_filters, kernel_size=5, stride=2) + x_2 = conv_block(x_1, num_filters, kernel_size=3, stride=1) + x = tf.keras.layers.Add()([x_1, x_2]) + return x + + def upscaling_unit(x): + num_filters = int(x.get_shape()[-1]) // 4 + x = tf.keras.layers.Lambda(lambda x: tf.nn.depth_to_space(x, 2))(x) + x_1 = conv_block(x, num_filters, kernel_size=3) + x_2 = conv_block(x_1, num_filters, kernel_size=3) + x = tf.keras.layers.Add()([x_1, x_2]) + return x + + def bottleneck_unit(x): + num_filters = int(x.get_shape()[-1]) + x_1 = conv_block(x, num_filters, kernel_size=3) + x_2 = conv_block(x_1, num_filters, kernel_size=3) + x = tf.keras.layers.Add()([x_1, x_2]) + return x + + # model flow + skip_connections = [] + num_filters = min_width + + # first layer + x = conv_block(inputs, num_filters, kernel_size=1, use_relu=False) + skip_connections.append(x) + + # encoder + for i in range(4): + x = downscaling_unit(x) + skip_connections.append(x) + + # bottleneck + x = bottleneck_unit(x) + + # decoder + for i in range(4): + x = tf.keras.layers.Add()([x, skip_connections.pop()]) + x = upscaling_unit(x) + + # last layer + x = tf.keras.layers.Add()([x, skip_connections.pop()]) + x = conv_block(x, 1, kernel_size=1, use_relu=False) + x = tf.keras.layers.Activation('sigmoid')(x) + + model = tf.keras.Model(inputs=inputs, outputs=x) return model \ No newline at end of file diff --git a/inference.py b/deepwatermap/inference.py similarity index 85% rename from inference.py rename to deepwatermap/inference.py index 8765b28..e15ab69 100644 --- a/inference.py +++ b/deepwatermap/inference.py @@ -1,70 +1,74 @@ -''' Runs inference on a given GeoTIFF image. - -example: -$ python inference.py --checkpoint_path checkpoints/cp.135.ckpt \ - --image_path sample_data/sentinel2_example.tif --save_path water_map.png -''' - -# Uncomment this to run inference on CPU if your GPU runs out of memory -# import os -# os.environ['CUDA_VISIBLE_DEVICES'] = '-1' - -import argparse -import deepwatermap -import tifffile as tiff -import numpy as np -import cv2 - -def find_padding(v, divisor=32): - v_divisible = max(divisor, int(divisor * np.ceil( v / divisor ))) - total_pad = v_divisible - v - pad_1 = total_pad // 2 - pad_2 = total_pad - pad_1 - return pad_1, pad_2 - -def main(checkpoint_path, image_path, save_path): - # load the model - model = deepwatermap.model() - model.load_weights(checkpoint_path) - - # load and preprocess the input image - image = tiff.imread(image_path) - pad_r = find_padding(image.shape[0]) - pad_c = find_padding(image.shape[1]) - image = np.pad(image, ((pad_r[0], pad_r[1]), (pad_c[0], pad_c[1]), (0, 0)), 'reflect') - - # solve no-pad index issue after inference - if pad_r[1] == 0: - pad_r = (pad_r[0], 1) - if pad_c[1] == 0: - pad_c = (pad_c[0], 1) - - image = image.astype(np.float32) - - # remove nans (and infinity) - replace with 0s - image = np.nan_to_num(image, copy=False, nan=0.0, posinf=0.0, neginf=0.0) - - image = image - np.min(image) - image = image / np.maximum(np.max(image), 1) - - # run inference - image = np.expand_dims(image, axis=0) - dwm = model.predict(image) - dwm = np.squeeze(dwm) - dwm = dwm[pad_r[0]:-pad_r[1], pad_c[0]:-pad_c[1]] - - # soft threshold - dwm = 1./(1+np.exp(-(16*(dwm-0.5)))) - dwm = np.clip(dwm, 0, 1) - - # save the output water map - cv2.imwrite(save_path, dwm * 255) - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--checkpoint_path', type=str, - help="Path to the dir where the checkpoints are stored") - parser.add_argument('--image_path', type=str, help="Path to the input GeoTIFF image") - parser.add_argument('--save_path', type=str, help="Path where the output map will be saved") - args = parser.parse_args() - main(args.checkpoint_path, args.image_path, args.save_path) +''' Runs inference on a given GeoTIFF image. + +example: +$ python inference.py --checkpoint_path checkpoints/cp.135.ckpt --image_path sample_data/image2.tif --save_path water_map.png +''' + +# Uncomment this to run inference on CPU if your GPU runs out of memory +# import os +# os.environ['CUDA_VISIBLE_DEVICES'] = '-1' + +import argparse +from deepwatermap import deepwatermap +import tifffile as tiff +import numpy as np +import cv2 + +def find_padding(v, divisor=32): + v_divisible = max(divisor, int(divisor * np.ceil( v / divisor ))) + total_pad = v_divisible - v + pad_1 = total_pad // 2 + pad_2 = total_pad - pad_1 + return pad_1, pad_2 + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--checkpoint_path', type=str, + help="Path to the dir where the checkpoints are stored") + parser.add_argument('--image_path', type=str, help="Path to the input GeoTIFF image") + parser.add_argument('--save_path', type=str, help="Path where the output map will be saved") + args = vars(parser.parse_args()) + + checkpoint_path=args["checkpoint_path"] + image_path=args["image_path"] + save_path=args["save_path"] + + # load the model + model = deepwatermap.model() + model.load_weights(checkpoint_path) + + # load and preprocess the input image + image = tiff.imread(image_path) + pad_r = find_padding(image.shape[0]) + pad_c = find_padding(image.shape[1]) + image = np.pad(image, ((pad_r[0], pad_r[1]), (pad_c[0], pad_c[1]), (0, 0)), 'reflect') + + # solve no-pad index issue after inference + if pad_r[1] == 0: + pad_r = (pad_r[0], 1) + if pad_c[1] == 0: + pad_c = (pad_c[0], 1) + + image = image.astype(np.float32) + + # remove nans (and infinity) - replace with 0s + image = np.nan_to_num(image, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + + image = image - np.min(image) + image = image / np.maximum(np.max(image), 1) + + # run inference + image = np.expand_dims(image, axis=0) + dwm = model.predict(image) + dwm = np.squeeze(dwm) + dwm = dwm[pad_r[0]:-pad_r[1], pad_c[0]:-pad_c[1]] + + # soft threshold + dwm = 1./(1+np.exp(-(16*(dwm-0.5)))) + dwm = np.clip(dwm, 0, 1) + + # save the output water map + cv2.imwrite(save_path, dwm * 255) + +if __name__ == '__main__': + main() diff --git a/metrics.py b/deepwatermap/metrics.py similarity index 97% rename from metrics.py rename to deepwatermap/metrics.py index 6099f40..769c408 100644 --- a/metrics.py +++ b/deepwatermap/metrics.py @@ -1,41 +1,41 @@ -''' This script defines custom metrics and loss functions. - -The Adaptive Max-Pool Loss acts as a weighting function that multiplies a -loss value with the maximum loss values within an NxN neighborhood. -An earlier version of this loss function described in: - -F. Isikdogan, A.C. Bovik, and P. Passalacqua, -"Learning a River Network Extractor using an Adaptive Loss Function," -IEEE Geoscience and Remote Sensing Letters, 2018. -''' - -import tensorflow as tf -from tensorflow.keras import backend as K -import numpy as np - -def running_recall(y_true, y_pred): - TP = K.sum(K.round(K.clip(y_true * y_pred, 0, 1))) - TP_FN = K.sum(K.round(K.clip(y_true, 0, 1))) - recall = TP / (TP_FN + K.epsilon()) - return recall - -def running_precision(y_true, y_pred): - TP = K.sum(K.round(K.clip(y_true * y_pred, 0, 1))) - TP_FP = K.sum(K.round(K.clip(y_pred, 0, 1))) - precision = TP / (TP_FP + K.epsilon()) - return precision - -def running_f1(y_true, y_pred): - precision = running_precision(y_true, y_pred) - recall = running_recall(y_true, y_pred) - return 2 * ((precision * recall) / (precision + recall + K.epsilon())) - -def adaptive_maxpool_loss(y_true, y_pred, alpha=0.25): - y_pred = K.clip(y_pred, K.epsilon(), 1. - K.epsilon()) - positive = -y_true * K.log(y_pred) * alpha - negative = -(1. - y_true) * K.log(1. - y_pred) * (1-alpha) - pointwise_loss = positive + negative - max_loss = tf.keras.layers.MaxPool2D(pool_size=8, strides=1, padding='same')(pointwise_loss) - x = pointwise_loss * max_loss - x = K.mean(x, axis=-1) +''' This script defines custom metrics and loss functions. + +The Adaptive Max-Pool Loss acts as a weighting function that multiplies a +loss value with the maximum loss values within an NxN neighborhood. +An earlier version of this loss function described in: + +F. Isikdogan, A.C. Bovik, and P. Passalacqua, +"Learning a River Network Extractor using an Adaptive Loss Function," +IEEE Geoscience and Remote Sensing Letters, 2018. +''' + +import tensorflow as tf +from tensorflow.keras import backend as K +import numpy as np + +def running_recall(y_true, y_pred): + TP = K.sum(K.round(K.clip(y_true * y_pred, 0, 1))) + TP_FN = K.sum(K.round(K.clip(y_true, 0, 1))) + recall = TP / (TP_FN + K.epsilon()) + return recall + +def running_precision(y_true, y_pred): + TP = K.sum(K.round(K.clip(y_true * y_pred, 0, 1))) + TP_FP = K.sum(K.round(K.clip(y_pred, 0, 1))) + precision = TP / (TP_FP + K.epsilon()) + return precision + +def running_f1(y_true, y_pred): + precision = running_precision(y_true, y_pred) + recall = running_recall(y_true, y_pred) + return 2 * ((precision * recall) / (precision + recall + K.epsilon())) + +def adaptive_maxpool_loss(y_true, y_pred, alpha=0.25): + y_pred = K.clip(y_pred, K.epsilon(), 1. - K.epsilon()) + positive = -y_true * K.log(y_pred) * alpha + negative = -(1. - y_true) * K.log(1. - y_pred) * (1-alpha) + pointwise_loss = positive + negative + max_loss = tf.keras.layers.MaxPool2D(pool_size=8, strides=1, padding='same')(pointwise_loss) + x = pointwise_loss * max_loss + x = K.mean(x, axis=-1) return x \ No newline at end of file diff --git a/trainer.py b/deepwatermap/trainer.py similarity index 97% rename from trainer.py rename to deepwatermap/trainer.py index ada6570..12a29ca 100644 --- a/trainer.py +++ b/deepwatermap/trainer.py @@ -1,137 +1,137 @@ -''' Trains a DeepWaterMap model. We provide a copy of the trained checkpoints. -You should not need this script unless you want to re-train the model. -''' - -import os, glob -import argparse -import tensorflow as tf -import deepwatermap -from metrics import running_precision, running_recall, running_f1 -from metrics import adaptive_maxpool_loss - -class TFModelTrainer: - def __init__(self, checkpoint_dir, data_path): - self.checkpoint_dir = checkpoint_dir - - # set training parameters - self.image_size = (512, 512) - self.learning_rate = 0.1 - self.num_epoch = 150 - self.batch_size = 24 - - # create the data generators - train_filenames = glob.glob(os.path.join(data_path, 'train_*.tfrecord')) - val_filenames = glob.glob(os.path.join(data_path, 'test_*.tfrecord')) - - self.dataset_train = self._data_layer(train_filenames) - self.dataset_val = self._data_layer(val_filenames) - - self.dataset_train_size = 137682 - self.dataset_val_size = 5000 - self.steps_per_epoch = self.dataset_train_size // self.batch_size - self.validation_steps = self.dataset_val_size // self.batch_size - - def _data_layer(self, filenames, num_threads=24): - dataset = tf.data.TFRecordDataset(filenames) - dataset = dataset.map(self._parse_tfrecord, num_parallel_calls=num_threads) - dataset = dataset.repeat() - dataset = dataset.batch(self.batch_size, drop_remainder=True) - dataset = dataset.prefetch(buffer_size=4) - return dataset - - def _parse_tfrecord(self, example_proto): - keys_to_features = {'B2': tf.io.FixedLenFeature([], tf.string), - 'B3': tf.io.FixedLenFeature([], tf.string), - 'B4': tf.io.FixedLenFeature([], tf.string), - 'B5': tf.io.FixedLenFeature([], tf.string), - 'B6': tf.io.FixedLenFeature([], tf.string), - 'B7': tf.io.FixedLenFeature([], tf.string), - 'L': tf.io.FixedLenFeature([], tf.string)} - F = tf.io.parse_single_example(example_proto, keys_to_features) - data = F['B2'], F['B3'], F['B4'], F['B5'], F['B6'], F['B7'], F['L'] - image, label = self._decode_images(data) - return image, label - - def _decode_images(self, data_strings): - bands = [[]] * len(data_strings) - for i in range(len(data_strings)): - bands[i] = tf.image.decode_png(data_strings[i]) - data = tf.concat(bands, -1) - data = tf.image.random_crop(data, size=[self.image_size[0], self.image_size[1], len(data_strings)]) - data = tf.cast(data, tf.float32) - image = data[..., :-1] / 255 - label = data[..., -1, None] / 3 - self._preprocess_images(image) - return image, label - - def _preprocess_images(self, image): - image = self._random_channel_mixing(image) - image = self._gaussian_noise(image) - image = self._normalize_image(image) - return image - - def _random_channel_mixing(self, image): - ccm = tf.eye(6)[None, :, :, None] - r = tf.random.uniform([3], maxval=0.25) + [0, 1, 0] - filter = r[None, :, None, None] - ccm = tf.nn.depthwise_conv2d(ccm, filter, strides=[1,1,1,1], padding='SAME', data_format='NHWC') - ccm = tf.squeeze(ccm) - image = tf.tensordot(image, ccm, (-1, 0)) - return image - - def _gaussian_noise(self, image): - r = tf.random.uniform((), maxval=0.04) - image = image + tf.random.normal([self.image_size[0], self.image_size[1], 6], stddev=r) - return image - - def _normalize_image(self, image): - image = tf.cast(image, tf.float32) - image = image - tf.reduce_min(image) - image = image / tf.maximum(tf.reduce_max(image), 1) - return image - - def _optimizer(self): - optimizer = tf.keras.optimizers.SGD(lr=self.learning_rate, momentum=0.9) - return optimizer - - def train(self): - # Callbacks - cp_callback = tf.keras.callbacks.ModelCheckpoint(os.path.join(self.checkpoint_dir, 'cp.{epoch:03d}.ckpt'), - save_weights_only=True) - tb_callback = tf.keras.callbacks.TensorBoard(log_dir=self.checkpoint_dir) - lr_callback = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1) - - # Model - model = deepwatermap.model() - - initial_epoch = 0 - ckpt = tf.train.get_checkpoint_state(self.checkpoint_dir) - if ckpt and ckpt.model_checkpoint_path: - model.load_weights(ckpt.model_checkpoint_path) - print("Loaded weights from", ckpt.model_checkpoint_path) - initial_epoch = int(ckpt.model_checkpoint_path.split('.')[-2]) - - model.compile(optimizer=self._optimizer(), - loss=adaptive_maxpool_loss, - metrics=[tf.keras.metrics.binary_accuracy, - running_precision, running_recall, running_f1]) - model.fit(self.dataset_train, - validation_data=self.dataset_val, - epochs=self.num_epoch, - initial_epoch=initial_epoch, - steps_per_epoch=self.steps_per_epoch, - validation_steps=self.validation_steps, - callbacks=[cp_callback, tb_callback, lr_callback]) - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--checkpoint_path', type=str, default='./checkpoints/', - help="Path to the dir where the checkpoints are saved") - parser.add_argument('--data_path', type=str, - help="Path to the tfrecord files") - args = parser.parse_args() - trainer = TFModelTrainer(args.checkpoint_path, args.data_path) - trainer.train() - -if __name__ == '__main__': - main() +''' Trains a DeepWaterMap model. We provide a copy of the trained checkpoints. +You should not need this script unless you want to re-train the model. +''' + +import os, glob +import argparse +import tensorflow as tf +import deepwatermap +from metrics import running_precision, running_recall, running_f1 +from metrics import adaptive_maxpool_loss + +class TFModelTrainer: + def __init__(self, checkpoint_dir, data_path): + self.checkpoint_dir = checkpoint_dir + + # set training parameters + self.image_size = (512, 512) + self.learning_rate = 0.1 + self.num_epoch = 150 + self.batch_size = 24 + + # create the data generators + train_filenames = glob.glob(os.path.join(data_path, 'train_*.tfrecord')) + val_filenames = glob.glob(os.path.join(data_path, 'test_*.tfrecord')) + + self.dataset_train = self._data_layer(train_filenames) + self.dataset_val = self._data_layer(val_filenames) + + self.dataset_train_size = 137682 + self.dataset_val_size = 5000 + self.steps_per_epoch = self.dataset_train_size // self.batch_size + self.validation_steps = self.dataset_val_size // self.batch_size + + def _data_layer(self, filenames, num_threads=24): + dataset = tf.data.TFRecordDataset(filenames) + dataset = dataset.map(self._parse_tfrecord, num_parallel_calls=num_threads) + dataset = dataset.repeat() + dataset = dataset.batch(self.batch_size, drop_remainder=True) + dataset = dataset.prefetch(buffer_size=4) + return dataset + + def _parse_tfrecord(self, example_proto): + keys_to_features = {'B2': tf.io.FixedLenFeature([], tf.string), + 'B3': tf.io.FixedLenFeature([], tf.string), + 'B4': tf.io.FixedLenFeature([], tf.string), + 'B5': tf.io.FixedLenFeature([], tf.string), + 'B6': tf.io.FixedLenFeature([], tf.string), + 'B7': tf.io.FixedLenFeature([], tf.string), + 'L': tf.io.FixedLenFeature([], tf.string)} + F = tf.io.parse_single_example(example_proto, keys_to_features) + data = F['B2'], F['B3'], F['B4'], F['B5'], F['B6'], F['B7'], F['L'] + image, label = self._decode_images(data) + return image, label + + def _decode_images(self, data_strings): + bands = [[]] * len(data_strings) + for i in range(len(data_strings)): + bands[i] = tf.image.decode_png(data_strings[i]) + data = tf.concat(bands, -1) + data = tf.image.random_crop(data, size=[self.image_size[0], self.image_size[1], len(data_strings)]) + data = tf.cast(data, tf.float32) + image = data[..., :-1] / 255 + label = data[..., -1, None] / 3 + self._preprocess_images(image) + return image, label + + def _preprocess_images(self, image): + image = self._random_channel_mixing(image) + image = self._gaussian_noise(image) + image = self._normalize_image(image) + return image + + def _random_channel_mixing(self, image): + ccm = tf.eye(6)[None, :, :, None] + r = tf.random.uniform([3], maxval=0.25) + [0, 1, 0] + filter = r[None, :, None, None] + ccm = tf.nn.depthwise_conv2d(ccm, filter, strides=[1,1,1,1], padding='SAME', data_format='NHWC') + ccm = tf.squeeze(ccm) + image = tf.tensordot(image, ccm, (-1, 0)) + return image + + def _gaussian_noise(self, image): + r = tf.random.uniform((), maxval=0.04) + image = image + tf.random.normal([self.image_size[0], self.image_size[1], 6], stddev=r) + return image + + def _normalize_image(self, image): + image = tf.cast(image, tf.float32) + image = image - tf.reduce_min(image) + image = image / tf.maximum(tf.reduce_max(image), 1) + return image + + def _optimizer(self): + optimizer = tf.keras.optimizers.SGD(lr=self.learning_rate, momentum=0.9) + return optimizer + + def train(self): + # Callbacks + cp_callback = tf.keras.callbacks.ModelCheckpoint(os.path.join(self.checkpoint_dir, 'cp.{epoch:03d}.ckpt'), + save_weights_only=True) + tb_callback = tf.keras.callbacks.TensorBoard(log_dir=self.checkpoint_dir) + lr_callback = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1) + + # Model + model = deepwatermap.model() + + initial_epoch = 0 + ckpt = tf.train.get_checkpoint_state(self.checkpoint_dir) + if ckpt and ckpt.model_checkpoint_path: + model.load_weights(ckpt.model_checkpoint_path) + print("Loaded weights from", ckpt.model_checkpoint_path) + initial_epoch = int(ckpt.model_checkpoint_path.split('.')[-2]) + + model.compile(optimizer=self._optimizer(), + loss=adaptive_maxpool_loss, + metrics=[tf.keras.metrics.binary_accuracy, + running_precision, running_recall, running_f1]) + model.fit(self.dataset_train, + validation_data=self.dataset_val, + epochs=self.num_epoch, + initial_epoch=initial_epoch, + steps_per_epoch=self.steps_per_epoch, + validation_steps=self.validation_steps, + callbacks=[cp_callback, tb_callback, lr_callback]) + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--checkpoint_path', type=str, default='./checkpoints/', + help="Path to the dir where the checkpoints are saved") + parser.add_argument('--data_path', type=str, + help="Path to the tfrecord files") + args = parser.parse_args() + trainer = TFModelTrainer(args.checkpoint_path, args.data_path) + trainer.train() + +if __name__ == '__main__': + main() diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..647c5a1 --- /dev/null +++ b/environment.yml @@ -0,0 +1,8 @@ +name: deepwatermap +channels: + - conda-forge +dependencies: + - tensorflow==1.12.0 + - tifffile + - opencv + - imagecodecs \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..8b81742 --- /dev/null +++ b/setup.py @@ -0,0 +1,26 @@ +import setuptools + +with open("README.md", "r", encoding="utf-8") as fh: + long_description = fh.read() + +setuptools.setup( + name="deepwatermap", + version="0.0.2", + author="Example Author", + author_email="author@example.com", + description="deepwatermap", + long_description=long_description, + long_description_content_type="text/markdown", + entry_points={ + "console_scripts": [ + "dwm_inference=deepwatermap.inference:main" + ] + }, + python_requires=">=3.6", + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: BSD License", + "Operating System :: OS Independent", + ], + packages=setuptools.find_packages(), +)