FID Generate Images.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from paint import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_dir = 'data/img_align_celeba/'\n",
    "MAX_IMAGES = 2000\n",
    "PARENT_DIR = 'fid'\n",
    "\n",
    "renderer_fn = 'renderer_constrained.pkl'\n",
    "\n",
    "actor_fns = [\n",
    "             'pretrained_models/gan/actor.pkl',\n",
    "             'pretrained_models/l2/actor.pkl',\n",
    "             'pretrained_models/l1/actor.pkl',\n",
    "             'pretrained_models/cm/actor.pkl',\n",
    "             'pretrained_models/cml1/actor.pkl'\n",
    "            ]\n",
    "labels = ['GAN', 'L2', 'L1*', 'CM', 'CM + L1*']\n",
    "dir_labels = ['gan', 'l2', 'l1', 'cm', 'cml1']\n",
    "inds_to_save = [20, 50, 100, 300]\n",
    "\n",
    "if not os.path.exists(PARENT_DIR): os.mkdir(PARENT_DIR)\n",
    "\n",
    "for ind_to_save in inds_to_save + ['final']:\n",
    "    if not os.path.exists(os.path.join(PARENT_DIR, str(ind_to_save))):\n",
    "        os.mkdir(os.path.join(PARENT_DIR, str(ind_to_save)))\n",
    "    for dir_label in dir_labels:\n",
    "        if not os.path.exists(os.path.join(PARENT_DIR, str(ind_to_save), dir_label)):\n",
    "            os.mkdir(os.path.join(PARENT_DIR, str(ind_to_save), dir_label))\n",
    "if not os.path.exists(os.path.join(PARENT_DIR, 'real')):\n",
    "    os.mkdir(os.path.join(PARENT_DIR, 'real'))\n",
    "\n",
    "# Get the files for the test\n",
    "file_list = []\n",
    "file_cnt = 0\n",
    "for filename in os.listdir(img_dir):\n",
    "    file_list.append(os.path.join(img_dir, filename))\n",
    "    file_cnt += 1\n",
    "    if file_cnt >= MAX_IMAGES: break\n",
    "    \n",
    "for fn in file_list:\n",
    "    cv2.imwrite(os.path.join(PARENT_DIR, 'real', os.path.basename(fn) + '.jpg'), cv2.imread(fn, cv2.IMREAD_COLOR))\n",
    "    for actor_ind in range(len(actor_fns)):\n",
    "        actor_fn = actor_fns[actor_ind]\n",
    "        \n",
    "        _, _, all_canvases, final_result = \\\n",
    "                paint(actor_fn, renderer_fn, img=fn, \n",
    "                      max_step=600, white_canvas=True, \n",
    "                      discrete_colors=False, div=1)\n",
    "        for ind_to_save in inds_to_save:\n",
    "            cv2.imwrite(os.path.join(PARENT_DIR, str(ind_to_save), dir_labels[actor_ind], os.path.basename(fn) + '.jpg'), all_canvases[ind_to_save][:,:,::-1])\n",
    "        cv2.imwrite(os.path.join(PARENT_DIR, 'final', dir_labels[actor_ind], os.path.basename(fn) + '.jpg'), final_result[:,:,::-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}