snapshot...

modules/deeplearning/espcn_l1b_l2.py

 import glob
 import tensorflow as tf
-from util.setup import logdir, modeldir, cachepath, now, ancillary_path
-from util.util import EarlyStop, normalize, denormalize, resample
+from util.setup import logdir, modeldir, now, ancillary_path
+from util.util import EarlyStop, normalize, denormalize, resample, resample_2d_linear, get_grid_values_all, \
+    resample_2d_linear_one
 import os, datetime
 import numpy as np
 import pickle
@@ -32,6 +33,8 @@ NOISE_TRAINING = False
 NOISE_STDDEV = 0.10
 DO_AUGMENT = True
 
+IMG_DEPTH = 1
+
 # setup scaling parameters dictionary
 mean_std_dct = {}
 mean_std_file = ancillary_path+'mean_std_lo_hi_l2.pkl'
@@ -47,26 +50,29 @@ f.close()
 mean_std_dct.update(mean_std_dct_l1b)
 mean_std_dct.update(mean_std_dct_l2)
 
-# data_params = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'temp_13_3um_nom', 'temp_3_75um_nom',
-#                'temp_6_7um_nom', 'temp_6_2um_nom', 'temp_7_3um_nom', 'temp_8_5um_nom', 'temp_9_7um_nom']
-data_params = ['temp_10_4um_nom', 'temp_11_0um_nom', 'temp_12_0um_nom', 'temp_13_3um_nom', 'temp_3_9um_nom',
-               'temp_6_7um_nom']
-# data_params = ['refl_0_65um_nom', 'temp_11_0um_nom', 'cld_temp_acha', 'cld_press_acha', 'cloud_fraction']
-label_params = ['refl_0_65um_nom', 'temp_11_0um_nom', 'cld_temp_acha', 'cld_press_acha', 'cloud_fraction', 'cld_opd_acha']
+DO_ZERO_OUT = False
 
+# label_param = 'cloud_fraction'
+label_param = 'cld_opd_dcomp'
+# label_param = 'cloud_probability'
 
-DO_ZERO_OUT = False
+params = ['temp_11_0um_nom', 'temp_12_0um_nom', 'refl_0_65um_nom', label_param]
+data_params = ['temp_11_0um_nom', 'refl_0_65um_nom']
+
+label_idx = params.index(label_param)
+
+print('data_params: ', data_params)
+print('label_param: ', label_param)
 
-label_idx = 1
-label_param = label_params[label_idx]
 
 x_134 = np.arange(134)
 y_134 = np.arange(134)
-#x_134_2 = x_134[3:131:2]
-#y_134_2 = y_134[3:131:2]
 x_134_2 = x_134[2:133:2]
 y_134_2 = y_134[2:133:2]
 
+slc_x = slice(3, 131)
+slc_y = slice(3, 131)
+
 
 def build_residual_conv2d_block(conv, num_filters, block_name, activation=tf.nn.leaky_relu, padding='SAME', scale=None):
     # kernel_initializer = 'glorot_uniform'
@@ -179,12 +185,9 @@ class ESPCN:
         self.test_data_nda = None
         self.test_label_nda = None
 
-        self.n_chans = len(data_params)
+        self.n_chans = len(data_params) + 1
 
         self.X_img = tf.keras.Input(shape=(None, None, self.n_chans))
-        # self.X_img = tf.keras.Input(shape=(36, 36, self.n_chans))
-        # self.X_img = tf.keras.Input(shape=(34, 34, self.n_chans))
-        # self.X_img = tf.keras.Input(shape=(66, 66, self.n_chans))
 
         self.inputs.append(self.X_img)
 
@@ -192,25 +195,16 @@ class ESPCN:
 
     def get_in_mem_data_batch(self, idxs, is_training):
         if is_training:
-            label_files = self.train_label_files
-            data_files = self.train_data_files
+            files = self.train_data_files
         else:
-            label_files = self.test_label_files
-            data_files = self.test_data_files
+            files = self.test_data_files
 
-        label_s = []
         data_s = []
         for k in idxs:
-            f = label_files[k]
-            nda = np.load(f)
-            label_s.append(nda)
-
-            f = data_files[k]
+            f = files[k]
             nda = np.load(f)
             data_s.append(nda)
-
-        label = np.concatenate(label_s)
-        data = np.concatenate(data_s)
+        input_data = np.concatenate(data_s)
 
         add_noise = None
         noise_scale = None
@@ -220,22 +214,30 @@ class ESPCN:
 
         data_norm = []
         for k, param in enumerate(data_params):
-            tmp = data[:, k, :, :]
-            tmp = resample(x_134, y_134, tmp, x_134_2, y_134_2)
+            tmp = input_data[:, k, :, :]
             tmp = normalize(tmp, param, mean_std_dct, add_noise=add_noise, noise_scale=noise_scale)
+            tmp = resample_2d_linear(x_134, y_134, tmp, x_134_2, y_134_2)
             data_norm.append(tmp)
-        data = np.stack(data_norm, axis=3)
-        data = data.astype(np.float32)
 
-        # label = label[:, label_idx, :, :]
-        label = label[:, label_idx, 3:131:2, 3:131:2]
-        # label = label[:, label_idx, 3:131, 3:131]
-        label = np.expand_dims(label, axis=3)
-        label = label.astype(np.float32)
+        tmp = input_data[:, label_idx, :, ]
+        if label_param != 'cloud_fraction':
+            tmp = normalize(tmp, label_param, mean_std_dct, add_noise=add_noise, noise_scale=noise_scale)
+        else:
+            tmp = np.where(np.isnan(tmp), 0, tmp)
+        tmp = resample_2d_linear(x_134, y_134, tmp, x_134_2, y_134_2)
+        data_norm.append(tmp)
+
+        data = np.stack(data_norm, axis=3)
 
+        label = input_data[:, label_idx, slc_y, slc_x]
         if label_param != 'cloud_fraction':
             label = normalize(label, label_param, mean_std_dct)
-            # label = label = scale(label, label_param, mean_std_dct)
+        else:
+            label = np.where(np.isnan(label), 0, label)
+        label = np.expand_dims(label, axis=3)
+
+        data = data.astype(np.float32)
+        label = label.astype(np.float32)
 
         if is_training and DO_AUGMENT:
             data_ud = np.flip(data, axis=1)
@@ -312,16 +314,14 @@ class ESPCN:
         dataset = dataset.map(self.data_function_evaluate, num_parallel_calls=8)
         self.eval_dataset = dataset
 
-    def setup_pipeline(self, train_data_files, train_label_files, test_data_files, test_label_files, num_train_samples):
+    def setup_pipeline(self, train_data_files, test_data_files, num_train_samples):
 
         self.train_data_files = train_data_files
         self.test_data_files = test_data_files
-        self.train_label_files = train_label_files
-        self.test_label_files = test_label_files
 
         trn_idxs = np.arange(len(train_data_files))
         np.random.shuffle(trn_idxs)
-        tst_idxs = np.arange(len(train_data_files))
+        tst_idxs = np.arange(len(test_data_files))
 
         self.get_train_dataset(trn_idxs)
         self.get_test_dataset(tst_idxs)
@@ -336,6 +336,7 @@ class ESPCN:
         print('num test samples: ', tst_idxs.shape[0])
         print('setup_pipeline: Done')
 
+
     def setup_test_pipeline(self, test_data_files):
         self.test_data_files = test_data_files
         tst_idxs = np.arange(len(test_data_files))
@@ -359,7 +360,7 @@ class ESPCN:
         kernel_initializer = 'he_uniform'
         momentum = 0.99
 
-        num_filters = 32
+        num_filters = 64
 
         input_2d = self.inputs[0]
         print('input: ', input_2d.shape)
@@ -382,18 +383,19 @@ class ESPCN:
 
         conv_b = build_residual_conv2d_block(conv_b, num_filters, 'Residual_Block_5', scale=scale)
 
-        conv_b = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, kernel_initializer=kernel_initializer)(conv_b)
+        # conv_b = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, kernel_initializer=kernel_initializer)(conv_b)
 
-        conv = conv + conv_b
+        # conv = conv + conv_b
+        conv = conv_b
         print(conv.shape)
 
-        # conv = tf.keras.layers.Conv2D(num_filters * (factor ** 2), 3, padding='same')(conv)
-        # print(conv.shape)
+        conv = tf.keras.layers.Conv2D(IMG_DEPTH * (factor ** 2), 3, padding='same')(conv)
+        print(conv.shape)
 
         conv = tf.nn.depth_to_space(conv, factor)
         print(conv.shape)
 
-        self.logits = tf.keras.layers.Conv2D(1, kernel_size=3, strides=1, padding=padding, name='regression')(conv)
+        self.logits = tf.keras.layers.Conv2D(IMG_DEPTH, kernel_size=1, strides=1, padding=padding, activation=activation, name='regression')(conv)
 
         print(self.logits.shape)
 
@@ -591,7 +593,6 @@ class ESPCN:
 
         ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model)
         ckpt_manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=3)
-
         ckpt.restore(ckpt_manager.latest_checkpoint)
 
         self.reset_test_metrics()
@@ -604,42 +605,35 @@ class ESPCN:
 
         print('loss, acc: ', self.test_loss.result().numpy(), self.test_accuracy.result().numpy())
 
-    def do_evaluate(self, nda_lr, param, ckpt_dir):
+    def do_evaluate(self, data, ckpt_dir):
 
         ckpt = tf.train.Checkpoint(step=tf.Variable(1), model=self.model)
         ckpt_manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=3)
 
         ckpt.restore(ckpt_manager.latest_checkpoint)
 
-        data = normalize(nda_lr, param, mean_std_dct)
-        #data = np.expand_dims(data, axis=0)
-        #data = np.expand_dims(data, axis=3)
-
         self.reset_test_metrics()
 
         pred = self.model([data], training=False)
         self.test_probs = pred
         pred = pred.numpy()
 
-        return denormalize(pred, param, mean_std_dct)
+        return pred
 
-    def run(self, directory):
-        train_data_files = glob.glob(directory+'data_train*.npy')
-        valid_data_files = glob.glob(directory+'data_valid*.npy')
-        train_label_files = glob.glob(directory+'label_train*.npy')
-        valid_label_files = glob.glob(directory+'label_valid*.npy')
+    def run(self, directory, ckpt_dir=None, num_data_samples=50000):
+        train_data_files = glob.glob(directory+'data_train_*.npy')
+        valid_data_files = glob.glob(directory+'data_valid_*.npy')
+        train_data_files = train_data_files[::2]
+        valid_data_files = valid_data_files[::2]
 
-        train_data_files.sort()
-        valid_data_files.sort()
-
-        self.setup_pipeline(train_data_files, train_label_files, valid_data_files, valid_label_files, 100000)
+        self.setup_pipeline(train_data_files, valid_data_files, num_data_samples)
         self.build_model()
         self.build_training()
         self.build_evaluation()
-        self.do_training()
+        self.do_training(ckpt_dir=ckpt_dir)
 
     def run_restore(self, directory, ckpt_dir):
-        valid_data_files = glob.glob(directory + 'data_valid*.npy')
+        valid_data_files = glob.glob(directory + 'data_*.npy')
         self.num_data_samples = 1000
         self.setup_test_pipeline(valid_data_files)
         self.build_model()
@@ -647,12 +641,59 @@ class ESPCN:
         self.build_evaluation()
         self.restore(ckpt_dir)
 
-    def run_evaluate(self, nda_lr, param, ckpt_dir):
+    def run_evaluate(self, data, ckpt_dir):
         self.num_data_samples = 80000
         self.build_model()
         self.build_training()
         self.build_evaluation()
-        return self.do_evaluate(nda_lr, param, ckpt_dir)
+        return self.do_evaluate(data, ckpt_dir)
+
+
+def run_restore_static(directory, ckpt_dir):
+    nn = ESPCN()
+    nn.run_restore(directory, ckpt_dir)
+
+
+def run_evaluate_static(in_file, out_file, ckpt_dir):
+    N = 8
+    sub_y, sub_x = (N+1) * 128, (N+1) * 128
+    y_0, x_0, = 2500 - int(sub_y/2), 2500 - int(sub_x/2)
+
+    slc_y_2, slc_x_2 = slice(1, 128*N + 6, 2), slice(1, 128*N + 6, 2)
+    y_2, x_2 = np.arange((128*N)/2 + 3), np.arange((128*N)/2 + 3)
+    t, s = np.arange(1, (128*N)/2 + 2, 0.5), np.arange(1, (128*N)/2 + 2, 0.5)
+
+    h5f = h5py.File(in_file, 'r')
+    grd_a = get_grid_values_all(h5f, 'temp_11_0um_nom')
+    grd_a = grd_a[y_0:y_0+sub_y, x_0:x_0+sub_x]
+    grd_a = grd_a[slc_y_2, slc_x_2]
+    grd_a = normalize(grd_a, 'temp_11_0um_nom', mean_std_dct)
+    grd_a = resample_2d_linear_one(x_2, y_2, grd_a, t, s)
+
+    grd_b = get_grid_values_all(h5f, 'refl_0_65um_nom')
+    grd_b = grd_b[y_0:y_0+sub_y, x_0:x_0+sub_x]
+    grd_b = grd_b[slc_y_2, slc_x_2]
+    grd_b = normalize(grd_b, 'refl_0_65um_nom', mean_std_dct)
+    grd_b = resample_2d_linear_one(x_2, y_2, grd_b, t, s)
+
+    grd_c = get_grid_values_all(h5f, label_param)
+    grd_c = grd_c[y_0:y_0+sub_y, x_0:x_0+sub_x]
+    grd_c = grd_c[slc_y_2, slc_x_2]
+    if label_param != 'cloud_fraction':
+        grd_c = normalize(grd_c, label_param, mean_std_dct)
+    grd_c = resample_2d_linear_one(x_2, y_2, grd_c, t, s)
+
+    data = np.stack([grd_a, grd_b, grd_c], axis=2)
+    data = np.expand_dims(data, axis=0)
+
+    nn = ESPCN()
+    out_sr = nn.run_evaluate(data, ckpt_dir)
+    if label_param != 'cloud_fraction':
+        out_sr = denormalize(out_sr, label_param, mean_std_dct)
+    if out_file is not None:
+        np.save(out_file, out_sr)
+    else:
+        return out_sr
 
 
 if __name__ == "__main__":