Made DataLoader object

The DataLoader object loads all files with a specific extension in the given directory. All the loaded files are stored in ascending order to keep track of the data's year. Issue #9.

data_loader/data_loader.py

@@ -0,0 +1,37 @@
+import numpy as np
+import os
+import sys
+from osgeo import gdal
+
+class DataLoader():
+
+    def __init__(self, data_dir):
+        self.data_dir = data_dir
+        self.files = []
+        self.arrays = []
+
+    def load_directory(self, ext):
+        for file in os.listdir(self.data_dir):
+            if file.endswith(ext):
+                # Stores the file without extension
+                self.files.append(os.path.splitext(file)[0])
+
+                print(os.path.join(self.data_dir, file))
+
+        # Turning all the string-values into integers
+        self.files = [int(file) for file in self.files]
+
+        # Sorts the file in ascending order based on the year
+        self.files = sorted(self.files, key=int)
+
+        # Turns the integers back into string values with extensions
+        self.files = [str(file) + ext for file in self.files]
+
+    def create_np_arrays(self):
+        for file in self.files:
+            pop_data = gdal.Open(os.path.join(self.data_dir, file))
+            array = np.array(pop_data.GetRasterBand(1).ReadAsArray())
+            # Null-values (neg-values) are replaced with zeros
+            array[array < 0] = 0
+            self.arrays.append(array)
+

mains/main.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from data_loader.data_generator import DataGenerator, PrepData, PrepTrainTest
+from data_loader.data_loader import DataLoader
 from models.pop_model import PopModel
 from trainers.pop_trainer import PopTrainer
 from utils.config import process_config
@@ -28,25 +29,12 @@ def main():
     else:
         config = process_config(os.path.join(config_dir, 'example.json'))
 
-    # Loads the geotiff
-    pop_data_10 = gdal.Open(os.path.join(data_dir, 'TGO10v4.tif'))
-    pop_data_14 = gdal.Open(os.path.join(data_dir, 'TGO14adjv1.tif'))
-    pop_data_15 = gdal.Open(os.path.join(data_dir, 'TGO15adjv4.tif'))
+    data_loader = DataLoader(data_dir)
+    data_loader.load_directory('.tif')
+    data_loader.create_np_arrays()
 
-    # Store the values of the geotif into a np.array
-    pop_arr_10 = np.array(pop_data_10.GetRasterBand(1).ReadAsArray())
-    pop_arr_10 = np.delete(pop_arr_10, -1, axis=1)  # Shape not the same as pop data from 14 and 15
-
-    pop_arr_14 = np.array(pop_data_14.GetRasterBand(1).ReadAsArray())
-    pop_arr_15 = np.array(pop_data_15.GetRasterBand(1).ReadAsArray())
-
-    # Null-values (neg-values) are replaced with zeros
-    pop_arr_10[pop_arr_10 < 0] = 0
-    pop_arr_14[pop_arr_14 < 0] = 0
-    pop_arr_15[pop_arr_15 < 0] = 0
-
-    preptt = PrepTrainTest(pop_arr_10, pop_arr_14, config.batch_size, config.chunk_height, config.chunk_width)
-    prepd = PrepData(pop_arr_10, pop_arr_14, config.batch_size, config.chunk_height, config.chunk_width)
+    preptt = PrepTrainTest(data_loader.arrays[0], data_loader.arrays[1], config.batch_size, config.chunk_height, config.chunk_width)
+    prepd = PrepData(data_loader.arrays[0], data_loader.arrays[1], config.batch_size, config.chunk_height, config.chunk_width)
 
 
     # create the experiments dirs

mains/use_model.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from data_loader.data_generator import DataGenerator, PrepData, PrepTrainTest
+from data_loader.data_loader import DataLoader
 from models.pop_model import PopModel
 from trainers.pop_trainer import PopTrainer
 from utils.config import process_config
@@ -28,26 +29,12 @@ def main():
     else:
         config = process_config(os.path.join(config_dir, 'example.json'))
 
+    data_loader = DataLoader(data_dir)
+    data_loader.load_directory('.tif')
+    data_loader.create_np_arrays()
 
-    # Loads the geotiff
-    pop_data_10 = gdal.Open(os.path.join(data_dir, 'TGO10v4.tif'))
-    pop_data_14 = gdal.Open(os.path.join(data_dir, 'TGO14adjv1.tif'))
-    pop_data_15 = gdal.Open(os.path.join(data_dir, 'TGO15adjv4.tif'))
-
-    # Store the values of the geotif into a np.array
-    pop_arr_10 = np.array(pop_data_10.GetRasterBand(1).ReadAsArray())
-    pop_arr_10 = np.delete(pop_arr_10, -1, axis=1)  # Shape not the same as pop data from 14 and 15
-
-    pop_arr_14 = np.array(pop_data_14.GetRasterBand(1).ReadAsArray())
-    pop_arr_15 = np.array(pop_data_15.GetRasterBand(1).ReadAsArray())
-
-    # Null-values (neg-values) are replaced with zeros
-    pop_arr_10[pop_arr_10 < 0] = 0
-    pop_arr_14[pop_arr_14 < 0] = 0
-    pop_arr_15[pop_arr_15 < 0] = 0
-
-    preptt = PrepTrainTest(pop_arr_10, pop_arr_14, config.batch_size)
-    prepd = PrepData(pop_arr_10, pop_arr_14, config.batch_size)
+    preptt = PrepTrainTest(data_loader.arrays[0], data_loader.arrays[1], config.batch_size, config.chunk_height, config.chunk_width)
+    prepd = PrepData(data_loader.arrays[0], data_loader.arrays[1], config.batch_size, config.chunk_height, config.chunk_width)
 
 
     # create the experiments dirs

tests/test.py

@@ -0,0 +1,44 @@
+import numpy as np
+import numpy as np
+import os
+import sys
+from osgeo import gdal
+import osr
+
+base_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(base_dir)
+data_dir = os.path.relpath('..\\data', base_dir)
+config_dir = os.path.relpath('..\\configs', base_dir)
+
+
+# Loads the geotiff
+pop_data_10 = gdal.Open(os.path.join(data_dir, '2010.tif'))
+
+# Store the values of the geotif into a np.array
+pop_arr_10 = np.array(pop_data_10.GetRasterBand(1).ReadAsArray())
+pop_arr_10 = np.delete(pop_arr_10, -1, axis=1)  # Shape not the same as pop data from 14 and 15
+
+print(pop_arr_10.shape)
+
+Projection = osr.SpatialReference()
+Projection.ImportFromWkt(pop_data_10.GetProjectionRef())
+
+geoTransform = pop_data_10.GetGeoTransform()
+
+driver = gdal.GetDriverByName('GTiff')
+
+dst_ds = driver.Create('2010.tif', xsize=pop_arr_10.shape[1], ysize=pop_arr_10.shape[0],
+                       bands=1, eType=gdal.GDT_Float32)
+
+dst_ds.SetGeoTransform((
+    geoTransform[0],  # x_min
+    geoTransform[1],  # pixel width
+    geoTransform[2],  # rotation
+    geoTransform[3],  # y_max
+    geoTransform[4],  # rotation
+    geoTransform[5]  # pixel height
+))
+
+dst_ds.SetProjection(Projection.ExportToWkt())
+dst_ds.GetRasterBand(1).WriteArray(pop_arr_10)
+dst_ds.FlushCache()  # Write to disk.