snapshot...

modules/GSOC/E2_ESRGAN/lib/dataset.py

@@ -166,22 +166,35 @@ class OpdNpyDataset:
         self.dataset = dataset
 
     def read_numpy_file_s(self, f_idxs):
-        data_s = []
+        opd_s = []
+        refl_s = []
         for fi in f_idxs:
             fname = self.filenames[fi]
             data = np.load(fname)
-            data = data[0, ]
-            data = scale(data, 'cld_opd_dcomp', mean_std_dct)
-            data = data.astype(np.float32)
-            data_s.append(data)
-        hr_image = np.concatenate(data_s)
-        hr_image = tf.expand_dims(hr_image, axis=3)
+
+            refl = data[0, ]
+            refl = scale(refl, 'refl_0_65um_nom', mean_std_dct)
+            refl = refl.astype(np.float32)
+            refl_s.append(refl)
+
+            opd = data[1, ]
+            opd = scale(opd, 'cld_opd_dcomp', mean_std_dct)
+            opd = opd.astype(np.float32)
+            opd_s.append(opd)
+
+        opd = np.concatenate(opd_s)
+        refl = np.concatenate(refl_s)
+
+        hr_image = np.stack([refl, opd], axis=3)
         hr_image = tf.image.crop_to_bounding_box(hr_image, 0, 0, self.hr_size, self.hr_size)
+
         low_resolution = tf.image.resize(hr_image, [self.lr_size, self.lr_size], method='bicubic')
         low_resolution = tf.math.multiply(low_resolution, 255.0)
-        hr_image = tf.math.multiply(hr_image, 255.0)
         low_resolution = tf.clip_by_value(low_resolution, 0, 255)
+
+        hr_image = tf.math.multiply(hr_image[:, :, :, 1], 255.0)
         high_resolution = tf.clip_by_value(hr_image, 0, 255)
+        high_resolution = tf.expand_dims(high_resolution, axis=3)
 
         low_resolution, high_resolution = augment_image()(low_resolution, high_resolution)