Convert mask to jpg / progress bar

src/consts.py

@@ -3,6 +3,7 @@ FEATURES ='../data/train_features'
 LABELS  = '../data/train_labels'
 JPG_IMAGES = '../images'
 FC_DIR = "fc"
+MASK_DIR = "mask"
 METADATA = pd.read_csv('../data/On_Cloud_N_Cloud_Cover_Detection_Challenge_-_train_metadata.csv.csv')
 
 BATCH = 8

src/helpers.py

@@ -6,6 +6,27 @@ import numpy as np
 import rasterio as rio
 import matplotlib.pyplot as plt
 
+def progress_bar (iteration, total, prefix = '', suffix = '', decimals = 1, length = 100, fill = '█', printEnd = "\r"):
+    """
+    Call in a loop to create terminal progress bar
+    @params:
+        iteration   - Required  : current iteration (Int)
+        total       - Required  : total iterations (Int)
+        prefix      - Optional  : prefix string (Str)
+        suffix      - Optional  : suffix string (Str)
+        decimals    - Optional  : positive number of decimals in percent complete (Int)
+        length      - Optional  : character length of bar (Int)
+        fill        - Optional  : bar fill character (Str)
+        printEnd    - Optional  : end character (e.g. "\r", "\r\n") (Str)
+    """
+    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
+    filledLength = int(length * iteration // total)
+    bar = fill * filledLength + '-' * (length - filledLength)
+    print(f'\r{prefix} |{bar}| {percent}% {suffix}', end = printEnd)
+    # Print New Line on Complete
+    if iteration == total: 
+        print()
+
 def load_img(path,expand_dim=False):
     img = cv2.imread(path)
 #   img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
@@ -24,23 +45,29 @@ def create_folder(name, path):
 def scale(band):
     return band/np.max(band)
 
-def convert_tif_to_jpg(rasters_dir,
-                            rgb_path=None,
-                            false_color_path=None):
+def convert_tif_to_jpg(features, labels,
+                        rgb_path=None,
+                        false_color_path=None,
+                        mask_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]
+    b2 = scale(rio.open(features+'/B02.tif').read().reshape(512, 512, 1))
+    b3 = scale(rio.open(features+'/B03.tif').read().reshape(512, 512, 1))
+    b4 = scale(rio.open(features+'/B04.tif').read().reshape(512, 512, 1))
+    b8 = scale(rio.open(features+'/B08.tif').read().reshape(512, 512, 1))
 
+    #zwykła konwersja
+    file_name = features.split(os.sep)[-1]
     rgb = np.dstack([b4, b3, b2])
 
     plt.imsave(fname=rgb_path + f'/{file_name}.jpeg',
                arr=rgb)
-
+    # konwersja false color
     fc = np.dstack([b8, b3, b2])
-
     plt.imsave(fname=false_color_path + f'/{file_name}.jpeg',
                arr=fc)
+
+    #konwersja maski
+    mask= rio.open(labels).read().reshape(512,512,1)
+    mask *= 255
+    cv2.imwrite(filename= mask_path + f'/{file_name}.jpeg',  
+                img=mask.astype(np.uint8))

src/setup.py

@@ -12,18 +12,21 @@ import cv2
 import warnings 
 #deep learning
 
-from consts import SEED, JPG_IMAGES, FC_DIR, FEATURES
-from helpers import  create_folder, convert_tif_to_jpg
+from consts import SEED, JPG_IMAGES, FC_DIR, FEATURES, MASK_DIR, LABELS
+from helpers import  create_folder, convert_tif_to_jpg, progress_bar
 
 warnings.filterwarnings('ignore')
 
 def transform_photo(tif_dir):
     dp = create_folder(tif_dir, JPG_IMAGES)
     fc = create_folder(FC_DIR, dp)
-    convert_tif_to_jpg(os.path.join(FEATURES, tif_dir), dp, fc)
+    mk = create_folder(MASK_DIR, dp)
+    convert_tif_to_jpg(os.path.join(FEATURES, tif_dir), os.path.join(LABELS, tif_dir + ".tif"),  dp, fc, mk)
 
 
 if __name__ == "__main__":
-    for d in os.listdir(FEATURES):
-        print(str(d), end='\t')
+    dirs = os.listdir(FEATURES)
+    progress_bar(0, len(dirs), prefix = 'Converting TIF to JPG:', suffix = 'Complete', length = 50)
+    for i, d in enumerate(dirs):
+        progress_bar(i, len(dirs), prefix = 'Converting TIF to JPG:', suffix = 'Complete', length = 50)
         transform_photo(d)