generator.py

from pathlib import Path

import albumentations as A
import cv2
import numpy as np
from tensorflow.keras.utils import Sequence

transforms = A.Compose(
    [
        A.ShiftScaleRotate(
            shift_limit=0.03125,
            scale_limit=0.20,
            rotate_limit=20,
            border_mode=cv2.BORDER_CONSTANT,
            value=0,
            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, mode):
        self.df = df
        self.indices = np.arange(len(df))
        self.batch_size = cfg.train.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
        ]
        imgs = []
        genders = []
        ages = []

        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"])

        imgs = np.asarray(imgs)
        genders = np.asarray(genders)
        ages = np.asarray(ages)

        return imgs, (genders, ages)

    def __len__(self):
        return len(self.df) // self.batch_size

    def on_epoch_end(self):
        np.random.shuffle(self.indices)