aug

requirements.txt

@@ -10,3 +10,4 @@ h5py
 tables
 hydra-core
 scikit-learn
+albumentations

src/generator.py

@@ -1,16 +1,30 @@
 from pathlib import Path
 import numpy as np
 import cv2
-from tensorflow.keras.utils import Sequence, to_categorical
+import albumentations as A
+from tensorflow.keras.utils import Sequence
+
+
+transforms = A.Compose([
+    A.OneOf([
+        A.ShiftScaleRotate(shift_limit=0.03125, scale_limit=0.20, rotate_limit=20, border_mode=cv2.BORDER_CONSTANT,
+                           value=0, p=1.0),
+        A.IAAAffine(scale=(0.8, 1.2), translate_percent=(-0.03125, 0.03125), rotate=(-10, 10), shear=(-40, 40),
+                    mode="constant", p=1.0)
+    ]),
+    A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
+    A.HorizontalFlip(p=0.5)
+])
 
 
 class ImageSequence(Sequence):
-    def __init__(self, cfg, df):
+    def __init__(self, cfg, df, mode):
         self.df = df
         self.indices = np.arange(len(df))
         self.batch_size = cfg.model.batch_size
         self.img_dir = Path(__file__).resolve().parents[1].joinpath("data", f"{cfg.data.db}_crop")
         self.img_size = cfg.model.img_size
+        self.mode = mode
 
     def __getitem__(self, idx):
         sample_indices = self.indices[idx * self.batch_size:(idx + 1) * self.batch_size]
@@ -21,6 +35,10 @@ def __getitem__(self, idx):
         for _, row in self.df.iloc[sample_indices].iterrows():
             img = cv2.imread(str(self.img_dir.joinpath(row["img_paths"])))
             img = cv2.resize(img, (self.img_size, self.img_size))
+
+            if self.mode == "train":
+                img = transforms(image=img)["image"]
+
             imgs.append(img)
             genders.append(row["genders"])
             ages.append(row["ages"])

train.py

@@ -23,8 +23,8 @@ def main(cfg):
     csv_path = Path(to_absolute_path(__file__)).parent.joinpath("meta", f"{cfg.data.db}.csv")
     df = pd.read_csv(str(csv_path))
     train, val = train_test_split(df, random_state=42, test_size=0.1)
-    train_gen = ImageSequence(cfg, train)
-    val_gen = ImageSequence(cfg, val)
+    train_gen = ImageSequence(cfg, train, "train")
+    val_gen = ImageSequence(cfg, val, "val")
 
     strategy = tf.distribute.MirroredStrategy()