--- a
+++ b/dataloader/data.py
@@ -0,0 +1,113 @@
+import os
+import cv2
+from glob import glob
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from albumentations import HorizontalFlip, VerticalFlip, Rotate
+
+
+def create_dir(path):
+    """
+    Create a directory.
+    """
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def load_data(path, split=0.2):
+    """
+    Load the images and masks
+    """
+    images = sorted(glob(f"{path}/*/image/*.png"))
+    masks = sorted(glob(f"{path}/*/mask/*.png"))
+
+    """ Split the data """
+    split_size = int(len(images) * split)
+    train_x, valid_x = train_test_split(images, test_size=split_size, random_state=42)
+    train_y, valid_y = train_test_split(masks, test_size=split_size, random_state=42)
+
+    return (train_x, train_y), (valid_x, valid_y)
+
+
+def augment_data(images, masks, save_path, augment=True):
+    """
+    Performing data augmentation.
+    """
+    IMG_HEIGHT = 512
+    IMG_WIDTH = 512
+
+    for idx, (x, y) in tqdm(enumerate(zip(images, masks)), total=len(images)):
+        """Extracting the directory and image name"""
+        directory_name = x.split("/")[-3]
+        name = directory_name + "_" + x.split("/")[-1].split(".")[0]
+
+        """ Read the image and mask """
+        x = cv2.imread(x, cv2.IMREAD_COLOR)
+        y = cv2.imread(y, cv2.IMREAD_COLOR)
+
+        if augment == True:
+            aug = HorizontalFlip(p=1.0)  # Applying Horizontal Flip 100%
+            augmented = aug(image=x, mask=y)
+            x1 = augmented["image"]
+            y1 = augmented["mask"]
+
+            aug = VerticalFlip(p=1)  # Applying Vertical Flip 100%
+            augmented = aug(image=x, mask=y)
+            x2 = augmented["image"]
+            y2 = augmented["mask"]
+
+            aug = Rotate(limit=45, p=1.0)  # Applying Rotation till 45 degress
+            augmented = aug(image=x, mask=y)
+            x3 = augmented["image"]
+            y3 = augmented["mask"]
+
+            X = [x, x1, x2, x3]
+            Y = [y, y1, y2, y3]
+
+        else:
+            X = [x]
+            Y = [y]
+
+        idx = 0
+        for i, m in zip(X, Y):
+            i = cv2.resize(i, (IMG_WIDTH, IMG_HEIGHT))
+            m = cv2.resize(m, (IMG_WIDTH, IMG_HEIGHT))
+            m = m / 255.0
+            m = (m > 0.5) * 255
+
+            if len(X) == 1:
+                tmp_image_name = f"{name}.jpg"
+                tmp_mask_name = f"{name}.jpg"
+            else:
+                tmp_image_name = f"{name}_{idx}.jpg"
+                tmp_mask_name = f"{name}_{idx}.jpg"
+
+            image_path = os.path.join(save_path, "image/", tmp_image_name)
+            mask_path = os.path.join(save_path, "mask/", tmp_mask_name)
+
+            cv2.imwrite(image_path, i)
+            cv2.imwrite(mask_path, m)
+
+            idx += 1
+
+
+if __name__ == "__main__":
+    """
+    Load the dataset
+    """
+    dataset_path = os.path.join("data", "train")
+    (train_x, train_y), (valid_x, valid_y) = load_data(dataset_path, split=0.2)
+    print("Train: ", len(train_x))
+    print("Valid: ", len(valid_x))
+
+    """ Create the directories """
+    create_dir("new_data/train/image/")
+    create_dir("new_data/train/mask/")
+    create_dir("new_data/valid/image/")
+    create_dir("new_data/valid/mask/")
+
+    """ Augment the data """
+    augment_data(
+        train_x, train_y, "new_data/train/", augment=True
+    )  # Applying Data Augmentation for training data
+    augment_data(valid_x, valid_y, "new_data/valid/", augment=False)