6. Train face attributes model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "\n",
    "from keras.applications.xception import Xception, preprocess_input\n",
    "from keras.layers import GlobalAveragePooling2D\n",
    "from keras.layers import Input, Dense, LeakyReLU\n",
    "from keras import backend as K\n",
    "from keras.models import Model\n",
    "from keras.callbacks import ModelCheckpoint\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "import tensorflow as tf\n",
    "\n",
    "PATH = '/workspace/dataset/'\n",
    "FACE_DEFAULT_SHAPE = (218, 178)\n",
    "BS = 32\n",
    "\n",
    "celeb_data = pd.read_csv('identity_CelebA.txt', sep=\" \", header=None)\n",
    "celeb_data.columns = [\"image\", \"label\"]\n",
    "attributes = pd.read_csv(PATH + 'list_attr_celeba.csv')\n",
    "attributes = attributes.replace(-1, 0)\n",
    "\n",
    "# 0 - train, 1 - validation, 2 - test\n",
    "train_val_test = pd.read_csv('list_eval_partition.csv', usecols=['partition']).values[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 162770 images.\n",
      "Found 19867 images.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n"
     ]
    }
   ],
   "source": [
    "# checkpoint = ModelCheckpoint('attributes.h5', monitor='val_loss', verbose=1, save_best_only=True, save_weights_only=True)\n",
    "\n",
    "attributes = attributes[['image_id','Attractive','Bald','Male','Bags_Under_Eyes','Narrow_Eyes',\n",
    "                         'Oval_Face','Pointy_Nose','Receding_Hairline','Young']]\n",
    "\n",
    "features = attributes.drop(['image_id'], axis=1).columns\n",
    "\n",
    "df_train = attributes.iloc[train_val_test == 0]\n",
    "df_valid = attributes.iloc[train_val_test == 1]\n",
    "df_test = attributes.iloc[train_val_test == 2]\n",
    "\n",
    "# necessary for flow_from_dataframe method\n",
    "df_valid = df_valid.reset_index()\n",
    "df_test = df_test.reset_index()\n",
    "\n",
    "datagen = ImageDataGenerator(horizontal_flip=True, preprocessing_function=preprocess_input)\n",
    "\n",
    "train_gen = datagen.flow_from_dataframe(df_train, directory=PATH+'img_align_celeba', x_col='image_id', \n",
    "                                        y_col=features, target_size=FACE_DEFAULT_SHAPE, color_mode='rgb',\n",
    "                                        classes=None, class_mode='other', batch_size=BS, shuffle=True)\n",
    "valid_gen = datagen.flow_from_dataframe(df_valid, directory=PATH+'img_align_celeba', x_col='image_id', \n",
    "                                        y_col=features, target_size=FACE_DEFAULT_SHAPE, color_mode='rgb',\n",
    "                                        classes=None, class_mode='other', batch_size=BS, shuffle=True)\n",
    "\n",
    "\n",
    "xception = Xception(include_top=False, weights=None, input_shape = FACE_DEFAULT_SHAPE + (3,))\n",
    "output = GlobalAveragePooling2D()(xception.output)\n",
    "base_model = Model(xception.input, output, name = 'base_xception')\n",
    "\n",
    "def get_attr_model(conv_feat_size, num_feat):\n",
    "    '''\n",
    "    Takes the output of the conv feature extractor and yields the embeddings\n",
    "    '''\n",
    "    input = Input((conv_feat_size,), name = 'input')\n",
    "    x = Dense(512)(input)\n",
    "    x = LeakyReLU(alpha=0.1)(x)\n",
    "    x = Dense(128)(x)\n",
    "    x = LeakyReLU(alpha=0.1)(x)\n",
    "    x = Dense(num_feat, activation='sigmoid')(x)\n",
    "    model = Model(input, x, name = 'attr_classification')\n",
    "    return model\n",
    "\n",
    "inp_shape = K.int_shape(base_model.input)[1:]\n",
    "conv_feat_size = K.int_shape(base_model.output)[-1]\n",
    "\n",
    "input = Input( inp_shape )\n",
    "emb_attr = get_attr_model(conv_feat_size, len(features))\n",
    "att_model = Model(input, emb_attr(base_model(input)))\n",
    "\n",
    "att_model.compile(Adam(lr=0.0002), loss = 'binary_crossentropy', metrics=['binary_accuracy'])\n",
    "att_model.fit_generator(train_gen, steps_per_epoch=len(train_gen), epochs=50, initial_epoch = 0,\n",
    "                             validation_data=valid_gen, validation_steps=len(valid_gen), \n",
    "                             use_multiprocessing=True, workers=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "att_model.save('attributes.h5')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}