Merge pull request yu4u#53 from yu4u/feature/age_estimation

Feature/age estimation

age_estimation/README.md

@@ -0,0 +1,68 @@
+# Age Estimation
+This sub-project focuses on improving the accuracy of age estimation.
+
+
+## Dataset Preparation
+
+### APPA-REAL Dataset
+Follow the instructions [here](../appa-real/README.md) to download and extract the dataset.
+
+### UTK Face Dataset
+Firstly, download images from [the website of the UTKFace dataset](https://susanqq.github.io/UTKFace/).
+`part1.tar.gz`, `part2.tar.gz`, and `part3.tar.gz` can be downloaded from `In-the-wild Faces` in Datasets section.
+Then, extract the archives:
+
+```sh
+tar zxf part1.tar.gz
+tar zxf part2.tar.gz
+tar zxf part3.tar.gz
+```
+
+Finally, run the following script to create the training data:
+
+```
+python3 utkface/create_db_utkface_with_margin.py --input [PATH_TO_DATASET_DIR] --output [OUTPUT_DIR]
+```
+
+`[PATH_TO_DATASET_DIR]` should be a directory that includes `part1`, `part2`, and `part3` directories.
+The cropped face images with margin will be created in `[OUTPUT_DIR]`.
+
+
+### Training
+
+```bash
+python3 train.py --appa_dir [PATH_to_appa-real-release] --utk_dir [PATH_TO_UTK_CROPPED_FACE_DIR] --nb_epochs 100
+```
+
+Options:
+
+```bash
+usage: train.py [-h] --appa_dir APPA_DIR [--utk_dir UTK_DIR]
+                [--output_dir OUTPUT_DIR] [--batch_size BATCH_SIZE]
+                [--nb_epochs NB_EPOCHS] [--lr LR] [--model_name MODEL_NAME]
+
+This script trains the CNN model for age estimation.
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --appa_dir APPA_DIR   path to the APPA-REAL dataset (default: None)
+  --utk_dir UTK_DIR     path to the UTK face dataset (default: None)
+  --output_dir OUTPUT_DIR
+                        checkpoint dir (default: checkpoints)
+  --batch_size BATCH_SIZE
+                        batch size (default: 32)
+  --nb_epochs NB_EPOCHS
+                        number of epochs (default: 30)
+  --lr LR               learning rate (default: 0.1)
+  --model_name MODEL_NAME
+                        model name: 'ResNet50' or 'InceptionResNetV2'
+                        (default: ResNet50)
+```
+
+### Result
+
+Currently the best MAE (against apparent age) is 5.250.
+
+<img src="result/result.png" width="480px">
+
+weights: https://github.com/yu4u/age-gender-estimation/releases/download/v0.5/age_only_weights.029-4.027-5.250.hdf5

age_estimation/generator.py

@@ -0,0 +1,131 @@
+import better_exceptions
+import random
+from pathlib import Path
+from PIL import Image
+import numpy as np
+import pandas as pd
+import cv2
+from keras.utils import Sequence, to_categorical
+import Augmentor
+
+
+def get_transform_func():
+    p = Augmentor.Pipeline()
+    p.flip_left_right(probability=0.5)
+    p.rotate(probability=1, max_left_rotation=5, max_right_rotation=5)
+    p.zoom_random(probability=0.5, percentage_area=0.95)
+    p.random_distortion(probability=0.5, grid_width=2, grid_height=2, magnitude=8)
+    p.random_color(probability=1, min_factor=0.8, max_factor=1.2)
+    p.random_contrast(probability=1, min_factor=0.8, max_factor=1.2)
+    p.random_brightness(probability=1, min_factor=0.8, max_factor=1.2)
+    p.random_erasing(probability=0.5, rectangle_area=0.2)
+
+    def transform_image(image):
+        image = [Image.fromarray(image)]
+        for operation in p.operations:
+            r = round(random.uniform(0, 1), 1)
+            if r <= operation.probability:
+                image = operation.perform_operation(image)
+        return image[0]
+    return transform_image
+
+
+class FaceGenerator(Sequence):
+    def __init__(self, appa_dir, utk_dir=None, batch_size=32, image_size=224):
+        self.image_path_and_age = []
+        self._load_appa(appa_dir)
+
+        if utk_dir:
+            self._load_utk(utk_dir)
+
+        self.image_num = len(self.image_path_and_age)
+        self.batch_size = batch_size
+        self.image_size = image_size
+        self.indices = np.random.permutation(self.image_num)
+        self.transform_image = get_transform_func()
+
+    def __len__(self):
+        return self.image_num // self.batch_size
+
+    def __getitem__(self, idx):
+        batch_size = self.batch_size
+        image_size = self.image_size
+        x = np.zeros((batch_size, image_size, image_size, 3), dtype=np.uint8)
+        y = np.zeros((batch_size, 1), dtype=np.int32)
+
+        sample_indices = self.indices[idx * batch_size:(idx + 1) * batch_size]
+
+        for i, sample_id in enumerate(sample_indices):
+            image_path, age = self.image_path_and_age[sample_id]
+            image = cv2.imread(str(image_path))
+            x[i] = self.transform_image(cv2.resize(image, (image_size, image_size)))
+            y[i] = age
+
+        return x, to_categorical(y, 101)
+
+    def on_epoch_end(self):
+        self.indices = np.random.permutation(self.image_num)
+
+    def _load_appa(self, appa_dir):
+        appa_root = Path(appa_dir)
+        train_image_dir = appa_root.joinpath("train")
+        gt_train_path = appa_root.joinpath("gt_avg_train.csv")
+        df = pd.read_csv(str(gt_train_path))
+
+        for i, row in df.iterrows():
+            age = min(100, int(row.apparent_age_avg))
+            # age = int(row.real_age)
+            image_path = train_image_dir.joinpath(row.file_name + "_face.jpg")
+
+            if image_path.is_file():
+                self.image_path_and_age.append([str(image_path), age])
+
+    def _load_utk(self, utk_dir):
+        image_dir = Path(utk_dir)
+
+        for image_path in image_dir.glob("*.jpg"):
+            image_name = image_path.name  # [age]_[gender]_[race]_[date&time].jpg
+            age = min(100, int(image_name.split("_")[0]))
+
+            if image_path.is_file():
+                self.image_path_and_age.append([str(image_path), age])
+
+
+class ValGenerator(Sequence):
+    def __init__(self, appa_dir, batch_size=32, image_size=224):
+        self.image_path_and_age = []
+        self._load_appa(appa_dir)
+        self.image_num = len(self.image_path_and_age)
+        self.batch_size = batch_size
+        self.image_size = image_size
+
+    def __len__(self):
+        return self.image_num // self.batch_size
+
+    def __getitem__(self, idx):
+        batch_size = self.batch_size
+        image_size = self.image_size
+        x = np.zeros((batch_size, image_size, image_size, 3), dtype=np.uint8)
+        y = np.zeros((batch_size, 1), dtype=np.int32)
+
+        for i in range(batch_size):
+            image_path, age = self.image_path_and_age[idx * batch_size + i]
+            image = cv2.imread(str(image_path))
+            x[i] = cv2.resize(image, (image_size, image_size))
+            y[i] = age
+
+        return x, to_categorical(y, 101)
+
+    def _load_appa(self, appa_dir):
+        appa_root = Path(appa_dir)
+        val_image_dir = appa_root.joinpath("valid")
+        gt_val_path = appa_root.joinpath("gt_avg_valid.csv")
+        df = pd.read_csv(str(gt_val_path))
+
+        for i, row in df.iterrows():
+            age = min(100, int(row.apparent_age_avg))
+            # age = int(row.real_age)
+            image_path = val_image_dir.joinpath(row.file_name + "_face.jpg")
+
+            if image_path.is_file():
+                self.image_path_and_age.append([str(image_path), age])

age_estimation/model.py

@@ -0,0 +1,37 @@
+import better_exceptions
+from keras.applications import ResNet50, InceptionResNetV2
+from keras.layers import Dense
+from keras.models import Model
+from keras import backend as K
+
+
+def age_mae(y_true, y_pred):
+    true_age = K.sum(y_true * K.arange(0, 101, dtype="float32"), axis=-1)
+    pred_age = K.sum(y_pred * K.arange(0, 101, dtype="float32"), axis=-1)
+    mae = K.mean(K.abs(true_age - pred_age))
+    return mae
+
+
+def get_model(model_name="ResNet50"):
+    base_model = None
+
+    if model_name == "ResNet50":
+        base_model = ResNet50(include_top=False, weights='imagenet', input_shape=(224, 224, 3), pooling="avg")
+    elif model_name == "InceptionResNetV2":
+        base_model = InceptionResNetV2(include_top=False, weights='imagenet', input_shape=(299, 299, 3), pooling="avg")
+
+    prediction = Dense(units=101, kernel_initializer="he_normal", use_bias=False, activation="softmax",
+                       name="pred_age")(base_model.output_layers[0].output)
+
+    model = Model(inputs=base_model.input, outputs=prediction)
+
+    return model
+
+
+def main():
+    model = get_model("InceptionResNetV2")
+    model.summary()
+
+
+if __name__ == '__main__':
+    main()

age_estimation/train.py

@@ -0,0 +1,86 @@
+import argparse
+from pathlib import Path
+import numpy as np
+from keras.callbacks import LearningRateScheduler, ModelCheckpoint
+from keras.optimizers import SGD
+from generator import FaceGenerator, ValGenerator
+from model import get_model, age_mae
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="This script trains the CNN model for age estimation.",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--appa_dir", type=str, required=True,
+                        help="path to the APPA-REAL dataset")
+    parser.add_argument("--utk_dir", type=str, default=None,
+                        help="path to the UTK face dataset")
+    parser.add_argument("--output_dir", type=str, default="checkpoints",
+                        help="checkpoint dir")
+    parser.add_argument("--batch_size", type=int, default=32,
+                        help="batch size")
+    parser.add_argument("--nb_epochs", type=int, default=30,
+                        help="number of epochs")
+    parser.add_argument("--lr", type=float, default=0.1,
+                        help="learning rate")
+    parser.add_argument("--model_name", type=str, default="ResNet50",
+                        help="model name: 'ResNet50' or 'InceptionResNetV2'")
+    args = parser.parse_args()
+    return args
+
+
+class Schedule:
+    def __init__(self, nb_epochs, initial_lr):
+        self.epochs = nb_epochs
+        self.initial_lr = initial_lr
+
+    def __call__(self, epoch_idx):
+        if epoch_idx < self.epochs * 0.25:
+            return self.initial_lr
+        elif epoch_idx < self.epochs * 0.50:
+            return self.initial_lr * 0.2
+        elif epoch_idx < self.epochs * 0.75:
+            return self.initial_lr * 0.04
+        return self.initial_lr * 0.008
+
+
+def main():
+    args = get_args()
+    appa_dir = args.appa_dir
+    utk_dir = args.utk_dir
+    model_name = args.model_name
+    batch_size = args.batch_size
+    nb_epochs = args.nb_epochs
+    lr = args.lr
+
+    if model_name == "ResNet50":
+        image_size = 224
+    elif model_name == "InceptionResNetV2":
+        image_size = 299
+
+    train_gen = FaceGenerator(appa_dir, utk_dir=utk_dir, batch_size=batch_size, image_size=image_size)
+    val_gen = ValGenerator(appa_dir, batch_size=batch_size, image_size=image_size)
+    model = get_model(model_name=model_name)
+    sgd = SGD(lr=0.1, momentum=0.9, nesterov=True)
+    model.compile(optimizer=sgd, loss="categorical_crossentropy", metrics=[age_mae])
+    model.summary()
+    output_dir = Path(__file__).resolve().parent.joinpath(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    callbacks = [LearningRateScheduler(schedule=Schedule(nb_epochs, initial_lr=lr)),
+                 ModelCheckpoint(str(output_dir) + "/weights.{epoch:03d}-{val_loss:.3f}-{val_age_mae:.3f}.hdf5",
+                                 monitor="val_age_mae",
+                                 verbose=1,
+                                 save_best_only=True,
+                                 mode="min")
+                 ]
+
+    hist = model.fit_generator(generator=train_gen,
+                               epochs=nb_epochs,
+                               validation_data=val_gen,
+                               verbose=1,
+                               callbacks=callbacks)
+
+    np.savez(str(output_dir.joinpath("history.npz")), history=hist.history)
+
+
+if __name__ == '__main__':
+    main()

utkface/create_db_utkface_with_margin.py

@@ -0,0 +1,67 @@
+import argparse
+from pathlib import Path
+from tqdm import tqdm
+import cv2
+import dlib
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="This script detect faces using dlib and save detected faces",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--input", "-i", type=str, required=True,
+                        help="path to input directory that includes part1, part2, part3 sub-directories")
+    parser.add_argument("--output", "-o", type=str, required=True,
+                        help="path to output directory")
+    parser.add_argument("--margin", type=float, default=0.4,
+                        help="margin around face regions")
+    args = parser.parse_args()
+    return args
+
+
+# robust image cropping from
+# https://stackoverflow.com/questions/15589517/how-to-crop-an-image-in-opencv-using-python
+def imcrop(img, x1, y1, x2, y2):
+    if x1 < 0 or y1 < 0 or x2 > img.shape[1] or y2 > img.shape[0]:
+        img, x1, x2, y1, y2 = pad_img_to_fit_bbox(img, x1, x2, y1, y2)
+    return img[y1:y2, x1:x2, :]
+
+
+def pad_img_to_fit_bbox(img, x1, x2, y1, y2):
+    img = cv2.copyMakeBorder(img, - min(0, y1), max(y2 - img.shape[0], 0),
+                             -min(0, x1), max(x2 - img.shape[1], 0), cv2.BORDER_REPLICATE)
+    y2 += -min(0, y1)
+    y1 += -min(0, y1)
+    x2 += -min(0, x1)
+    x1 += -min(0, x1)
+    return img, x1, x2, y1, y2
+
+
+def main():
+    args = get_args()
+    root_dir = Path(args.input)
+    output_dir = Path(args.output)
+    margin = args.margin
+    detector = dlib.get_frontal_face_detector()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for image_path in tqdm(root_dir.glob("*/*.jpg")):
+        img = cv2.imread(str(image_path))
+        input_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        detected = detector(input_img, 1)
+
+        if len(detected) != 1:
+            continue
+
+        d = detected[0]
+        x1, y1, x2, y2, w, h = d.left(), d.top(), d.right() + 1, d.bottom() + 1, d.width(), d.height()
+        xw1 = int(x1 - margin * w)
+        yw1 = int(y1 - margin * h)
+        xw2 = int(x2 + margin * w)
+        yw2 = int(y2 + margin * h)
+        image_name = image_path.name
+        cropped_img = imcrop(img, xw1, yw1, xw2, yw2)
+        cv2.imwrite(str(output_dir.joinpath(image_name)), cropped_img)
+
+
+if __name__ == '__main__':
+    main()