From 3efb58a1fd363d410e8487543d33343e2f356cb5 Mon Sep 17 00:00:00 2001
From: nlitkowski <norbert.litkowski@gmail.com>
Date: Wed, 9 Feb 2022 23:57:10 +0100
Subject: [PATCH] Script for preprocessing images

---
 .gitignore        |  2 ++
 transform_data.py | 63 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 transform_data.py

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..9bd8019
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+**/images*
+**/data/train_features*
\ No newline at end of file
diff --git a/transform_data.py b/transform_data.py
new file mode 100644
index 0000000..ebcc3e7
--- /dev/null
+++ b/transform_data.py
@@ -0,0 +1,63 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+
+import cv2 as cv
+import rasterio as rio
+
+RASTERS_DIR = './data/data/train_features'
+IMAGES_DIR = './images/'
+FC_DIR = "fc"
+
+
+def load_img(path, expand_dim=False):
+    img = cv.imread(path)
+    img = img / 255.0
+    if expand_dim:
+        img = np.expand_dims(img, axis=0)
+
+    return img
+
+
+def create_folder(name, path):
+    n = os.path.join(path, name)
+    if not os.path.exists(n):
+        os.makedirs(n)
+    return n
+
+
+def scale(band):
+    return band/np.max(band)
+
+
+def convert_raster_to_image(rasters_dir,
+                            rgb_path=None,
+                            false_color_path=None):
+
+    b2 = scale(rio.open(rasters_dir+'/B02.tif').read().reshape(512, 512, 1))
+    b3 = scale(rio.open(rasters_dir+'/B03.tif').read().reshape(512, 512, 1))
+    b4 = scale(rio.open(rasters_dir+'/B04.tif').read().reshape(512, 512, 1))
+    b8 = scale(rio.open(rasters_dir+'/B08.tif').read().reshape(512, 512, 1))
+
+    file_name = rasters_dir.split(os.sep)[-1]
+
+    rgb = np.dstack([b4, b3, b2])
+
+    plt.imsave(fname=rgb_path + f'/{file_name}.jpeg',
+               arr=rgb)
+
+    fc = np.dstack([b8, b3, b2])
+
+    plt.imsave(fname=false_color_path + f'/{file_name}.jpeg',
+               arr=fc)
+
+
+def transform_photo(raster_dir):
+    dp = create_folder(raster_dir, IMAGES_DIR)
+    fc = create_folder(FC_DIR, dp)
+    convert_raster_to_image(os.path.join(RASTERS_DIR, raster_dir), dp, fc)
+
+
+if __name__ == "__main__":
+    for d in os.listdir(RASTERS_DIR):
+        transform_photo(d)