[ { "path": ".gitignore", "content": "*.swp\n__pycache__\n*.pbtext\ncompletions\n*.pyc\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2017 TeckYian Lim\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "Semantic Image Inpainting With Deep Generative Models and image restoration with GANS\n=====================================================\n[[Inpainting CVPR2017]](http://www.isle.illinois.edu/~yeh17/projects/semantic_inpaint/index.html)\n[[Image Restore ICASSP2018]](https://teckyianlim.me/image-restore/image-restore.html)\n\nTensorflow implementation for semantic image inpainting:\n\n![](http://www.isle.illinois.edu/~yeh17/projects/semantic_inpaint/img/process.png)\n\nSemantic Image Inpainting With Deep Generative Models\n\n[Raymond A. Yeh*](http://www.isle.illinois.edu/~yeh17/),\n[Teck Yian Lim*](http://tlim11.web.engr.illinois.edu/),\n[Chen Chen](http://cchen156.web.engr.illinois.edu/),\n[Alexander G. Schwing](http://www.alexander-schwing.de/),\n[Mark Hasegawa-Johnson](http://www.ifp.illinois.edu/~hasegawa/),\n[Minh N. Do](http://minhdo.ece.illinois.edu/)\n\nIn CVPR 2017\n\n\\* indicating equal contributions.\n\nOverview\n--------\nImplementation of proposed cost function and backpropogation to input. \n\nIn this code release, we load a pretrained DCGAN model, and apply our proposed\nobjective function for the task of image completion. \n\nNotes\n-----\nTo reproduce the CVPR2017 work, run the inpainting example.\n\nDependencies\n------------\n - Tensorflow >= 1.0\n - scipy + PIL/pillow (image io)\n - pyamg (for Poisson blending)\n\nTested to work with both Python 2.7 and Python 3.5\n\n\nFiles\n-----\n - src/model.py - Main implementation\n - src/inpaint.py - command line application\n - src/external - external code used. Citations in code\n - graphs/dcgan-100.pb - frozen pretrained DCGAN with 100-dimension latent space\n \nWeights\n-------\n\nGit doesn't work nicely with large binary files. Please download our weights from \n[here](https://www.dropbox.com/s/3uo97fzu4jfi2ms/dcgan-100.pb?dl=0), trained on the \n[CelebA dataset](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html).\n\nAlternatively, train your own GAN using your dataset. Conversion from checkpoint to \nTensorflow ProtoBuf format can be done with \n[this script](https://gist.github.com/moodoki/e37a85fb0258b045c005ca3db9cbc7f6)\n\n\nRunning\n-------\n\n\nGenerate multiple candidates for completion:\n```\npython src/inpaint.py --model_file graphs/dcgan-100.pb \\\n --maskType center --in_image testimages/face1.png \\\n --nIter 1000 --blend\n```\n\nGenerate completions for multiple input images:\n```\npython src/inpaint.py --model_file graphs/dcgan-100.pb \\\n --maskType center --inDir testimages \\\n --nIter 1000 --blend\n```\n\n\nCitation\n--------\n\n~~~\n@inproceedings{\n yeh2017semantic,\n title={Semantic Image Inpainting with Deep Generative Models},\n author={Yeh$^\\ast$, Raymond A. and Chen$^\\ast$, Chen and Lim, Teck Yian and Schwing Alexander G. and Hasegawa-Johnson, Mark and Do, Minh N.},\n booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},\n year={2017},\n note = {$^\\ast$ equal contribution},\n}\n~~~\n\n" }, { "path": "graphs/.gitignore", "content": "*\n!.gitignore\n" }, { "path": "requirements.txt", "content": "Pillow\ntensorflow\npyamg\n" }, { "path": "runall.sh", "content": "batches=`seq 0 15`\n\nfor b in $batches\ndo\n echo test_batches/quantized/batch_$b\n python src/quantize.py --model_file graphs/dcgan-100.pb --inDir test_batches/quantized/batch_$b --nIter 1000 --blend --outDir test_out\ndone\n\nfor b in $batches\ndo\n echo test_batches/gaussian_noise/batch_$b\n python src/denoising.py --model_file graphs/dcgan-100.pb --inDir test_batches/gaussian_noise/batch_$b --nIter 1000 --blend --outDir test_out\ndone\n\nfor b in $batches\ndo\n echo test_batches/testset/batch_$b\n python src/colorize.py --model_file graphs/dcgan-100.pb --inDir test_batches/testset/batch_$b --nIter 1000 --blend --outDir test_out\ndone\n\nfor b in $batches\ndo\n echo test_batches/sr_linear/batch_$b\n python src/superres.py --model_file graphs/dcgan-100.pb --inDir test_batches/sr_linear/batch_$b --nIter 1000 --blend --outDir test_out\ndone\n\nfor b in $batches\ndo\n echo test_batches/sr_nn/batch_$b\n python src/superres.py --model_file graphs/dcgan-100.pb --inDir test_batches/sr_nn/batch_$b --nIter 1000 --blend --outDir test_out/sr_nn\ndone\n\n\n" }, { "path": "src/colorize.py", "content": "import tensorflow as tf\nimport scipy.misc\nimport argparse\nimport os\nimport numpy as np\nfrom glob import glob\nfrom helper import loadimage_gray, saveimages\n\nfrom model_colorize import ModelColorize\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--model_file', type=str, help=\"Pretrained GAN model\")\nparser.add_argument('--lr', type=float, default=0.01)\nparser.add_argument('--momentum', type=float, default=0.9)\nparser.add_argument('--nIter', type=int, default=1000)\nparser.add_argument('--imgSize', type=int, default=64)\nparser.add_argument('--batch_size', type=int, default=64)\nparser.add_argument('--lambda_p', type=float, default=0.1)\nparser.add_argument('--checkpointDir', type=str, default='checkpoint')\nparser.add_argument('--outDir', type=str, default='colorization')\nparser.add_argument('--blend', action='store_true', default=False,\n help=\"Blend predicted image to original image\")\nparser.add_argument('--in_image', type=str, default=None,\n help='Input Image (ignored if inDir is specified')\nparser.add_argument('--inDir', type=str, default=None,\n help='Path to input images')\nparser.add_argument('--imgExt', type=str, default='png',\n help='input images file extension')\nparser.add_argument('-c', action='store_true', help='corrupt image on the fly')\n\nargs = parser.parse_args()\n\ndef main():\n m = ModelColorize(args.model_file, args)\n\n # Generate some samples from the model as a test\n imout = m.sample()\n #saveimages(imout)\n\n if args.inDir is not None:\n imgfilenames = glob( args.inDir + '/*.' + args.imgExt )\n print('{} images found'.format(len(imgfilenames)))\n in_img = np.array([loadimage_gray(f) for f in imgfilenames])\n elif args.in_img is not None:\n in_img = loadimage_gray(args.in_image)\n else:\n print('Input image needs to be specified')\n exit(1)\n\n saveimages(in_img, 'colorize_in', imgfilenames, args.outDir)\n\n inpaint_out, g_out = m.colorize(in_img, args.blend)\n saveimages(g_out, 'colorize_gen', imgfilenames, args.outDir)\n saveimages(inpaint_out, 'colorize', imgfilenames, args.outDir)\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "src/corrupt_image.py", "content": "from scipy.misc import imread, imsave, imresize\nimport argparse\nimport numpy as np\n\nparser = argparse.ArgumentParser()\nparser.add_argument('-i', '--input', type=str, required=True)\nparser.add_argument('-o', '--output', type=str, required=True)\nparser.add_argument('-t', '--type', type=str, \n choices=(['gaussian_noise',\n 'quantize',\n 'sr_linear',\n 'sr_nn',\n 'sr_black'\n ]),\n default='gaussian_noise')\n\ndef tofloat(img):\n img = img.astype(np.float32)\n return img/np.max(img)\n\ndef gaussian_noise(img, std=0.1):\n i = tofloat(img)\n i = i + std*np.random.normal(size=img.shape)\n return i\n\ndef quantize_2(img, levels=4):\n i = tofloat(img)\n i = np.floor(i*levels)/levels\n return i\n\ndef quantize(img, quantize_factor=55):\n img = img.astype(float)\n img /= 255.0\n img *= (255.0 / quantize_factor)\n images = img\n images = np.floor(images)\n images /= (255.0 / quantize_factor)\n images *= 255.0\n images = images.astype(int)\n return images\n\n\ndef resize(img, rate=4., interp='bilinear'):\n return imresize(img, float(rate), interp=interp) \n\ndef sr_linear(img, rate=4.):\n s = resize(img, 1./rate, interp='bilinear')\n return resize(s, rate, interp='bilinear')\n\ndef sr_nn(img, rate=4.):\n s = resize(img, 1./rate, interp='bilinear')\n return resize(s, rate, interp='nearest')\n\ndef sr_black(img, rate=4):\n o = np.zeros_like(img)\n s = resize(img, 1./rate, interp='bilinear')\n o[::rate, ::rate, :] = s\n return o\n\ndef main(args):\n i = imread(args.input)\n o = eval(args.type)(i)\n imsave(args.output, o)\n\nif __name__ == '__main__':\n args = parser.parse_args()\n main(args)\n \n" }, { "path": "src/denoising.py", "content": "import tensorflow as tf\nimport scipy.misc\nimport argparse\nimport os\nimport numpy as np\nfrom glob import glob\nfrom helper import loadimage, saveimages\n\nfrom model_denoising import ModelDenoising\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--model_file', type=str, help=\"Pretrained GAN model\")\nparser.add_argument('--lr', type=float, default=0.0001)\nparser.add_argument('--momentum', type=float, default=0.9)\nparser.add_argument('--nIter', type=int, default=1000)\nparser.add_argument('--imgSize', type=int, default=64)\nparser.add_argument('--batch_size', type=int, default=64)\nparser.add_argument('--lambda_p', type=float, default=0.03)\nparser.add_argument('--checkpointDir', type=str, default='checkpoint')\nparser.add_argument('--outDir', type=str, default='denoising')\nparser.add_argument('--blend', action='store_true', default=False,\n help=\"Blend predicted image to original image\")\nparser.add_argument('--in_image', type=str, default=None,\n help='Input Image (ignored if inDir is specified')\nparser.add_argument('--inDir', type=str, default=None,\n help='Path to input images')\nparser.add_argument('--imgExt', type=str, default='png',\n help='input images file extension')\nparser.add_argument('-c', action='store_true', help='corrupt image on the fly')\n\nargs = parser.parse_args()\n\n\ndef corrupt_image(img, down_sample_factor=4):\n noisy = img + 0.9*img.std()*np.random.random(img.shape)\n return noisy\n\ndef main():\n m = ModelDenoising(args.model_file, args)\n m.sigma=0.1\n\n # Generate some samples from the model as a test\n #imout = m.sample()\n #saveimages(imout)\n\n if args.inDir is not None:\n imgfilenames = glob( args.inDir + '/*.' + args.imgExt )\n print('{} images found'.format(len(imgfilenames)))\n in_img = np.array([loadimage(f) for f in imgfilenames])\n if args.c:\n in_corrupt_img = np.array([corrupt_image(loadimage(f), dd) for f in imgfilenames])\n else:\n in_corrupt_img = np.copy(in_img)\n\n\n elif args.in_image is not None:\n in_img = in_corrupt_img #loadimage(args.in_image)\n else:\n print('Input image needs to be specified')\n exit(1)\n #saveimages(in_corrupt_img, prefix='input')\n inpaint_out, g_out = m.restore_image(in_img)\n saveimages(g_out, 'denoise_gen', imgfilenames, args.outDir)\n saveimages(inpaint_out, 'denoise', imgfilenames, args.outDir)\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "src/external/__init__.py", "content": "" }, { "path": "src/external/poissonblending.py", "content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n#\n#\n# Original code from https://github.com/parosky/poissonblending \n\nimport numpy as np\nimport scipy.sparse\nimport PIL.Image\nimport pyamg\n\n# pre-process the mask array so that uint64 types from opencv.imread can be adapted\ndef prepare_mask(mask):\n if type(mask[0][0]) is np.ndarray:\n result = np.ndarray((mask.shape[0], mask.shape[1]), dtype=np.uint8)\n for i in range(mask.shape[0]):\n for j in range(mask.shape[1]):\n if sum(mask[i][j]) > 0:\n result[i][j] = 1\n else:\n result[i][j] = 0\n mask = result\n return mask\n\ndef blend(img_target, img_source, img_mask, offset=(0, 0)):\n # compute regions to be blended\n region_source = (\n max(-offset[0], 0),\n max(-offset[1], 0),\n min(img_target.shape[0]-offset[0], img_source.shape[0]),\n min(img_target.shape[1]-offset[1], img_source.shape[1]))\n region_target = (\n max(offset[0], 0),\n max(offset[1], 0),\n min(img_target.shape[0], img_source.shape[0]+offset[0]),\n min(img_target.shape[1], img_source.shape[1]+offset[1]))\n region_size = (region_source[2]-region_source[0], region_source[3]-region_source[1])\n\n # clip and normalize mask image\n img_mask = img_mask[region_source[0]:region_source[2], region_source[1]:region_source[3]]\n img_mask = prepare_mask(img_mask)\n img_mask[img_mask==0] = False\n img_mask[img_mask!=False] = True\n\n # create coefficient matrix\n A = scipy.sparse.identity(np.prod(region_size), format='lil')\n for y in range(region_size[0]):\n for x in range(region_size[1]):\n if img_mask[y,x]:\n index = x+y*region_size[1]\n A[index, index] = 4\n if index+1 < np.prod(region_size):\n A[index, index+1] = -1\n if index-1 >= 0:\n A[index, index-1] = -1\n if index+region_size[1] < np.prod(region_size):\n A[index, index+region_size[1]] = -1\n if index-region_size[1] >= 0:\n A[index, index-region_size[1]] = -1\n A = A.tocsr()\n \n # create poisson matrix for b\n P = pyamg.gallery.poisson(img_mask.shape)\n\n # for each layer (ex. RGB)\n for num_layer in range(img_target.shape[2]):\n # get subimages\n t = img_target[region_target[0]:region_target[2],region_target[1]:region_target[3],num_layer]\n s = img_source[region_source[0]:region_source[2], region_source[1]:region_source[3],num_layer]\n t = t.flatten()\n s = s.flatten()\n\n # create b\n b = P * s\n for y in range(region_size[0]):\n for x in range(region_size[1]):\n if not img_mask[y,x]:\n index = x+y*region_size[1]\n b[index] = t[index]\n\n # solve Ax = b\n x = pyamg.solve(A,b,verb=False,tol=1e-10)\n\n # assign x to target image\n x = np.reshape(x, region_size)\n x[x>255] = 255\n x[x<0] = 0\n x = np.array(x, img_target.dtype)\n img_target[region_target[0]:region_target[2],region_target[1]:region_target[3],num_layer] = x\n\n return img_target\n\n\ndef test():\n img_mask = np.asarray(PIL.Image.open('./testimages/test1_mask.png'))\n img_mask.flags.writeable = True\n img_source = np.asarray(PIL.Image.open('./testimages/test1_src.png'))\n img_source.flags.writeable = True\n img_target = np.asarray(PIL.Image.open('./testimages/test1_target.png'))\n img_target.flags.writeable = True\n img_ret = blend(img_target, img_source, img_mask, offset=(40,-30))\n img_ret = PIL.Image.fromarray(np.uint8(img_ret))\n img_ret.save('./testimages/test1_ret.png')\n\n\nif __name__ == '__main__':\n test()\n" }, { "path": "src/helper.py", "content": "import scipy\nimport os\nimport numpy as np\n\ndef loadimage(filename):\n img = scipy.misc.imread(filename).astype(np.float)\n return img\n\ndef loadimage_gray(filename):\n img = scipy.misc.imread(filename, mode='L').astype(np.float)[:,:,np.newaxis]\n return img\n\ndef saveimages(outimages, prefix='samples', filenames=None, outdir='out'):\n numimages = len(outimages)\n print(\"Array shape {}\".format(outimages.shape))\n\n if not os.path.exists(outdir):\n os.mkdir(outdir)\n\n for i in range(numimages):\n if filenames is None:\n filename = '{}_{}.png'.format(prefix, i)\n else:\n filename = '{}_{}'.format(prefix, os.path.basename(filenames[i]))\n filename = os.path.join(outdir, filename)\n scipy.misc.imsave(filename, np.squeeze(outimages[i, :, :, :]))\n\n\n" }, { "path": "src/imgsnr.py", "content": "from scipy.misc import imread\nimport numpy as np\nimport argparse\n\nparser = argparse.ArgumentParser()\nparser.add_argument('-i', '--input', type=str, help=\"original\")\nparser.add_argument('-o', '--output', type=str, help=\"processed\")\n\ndef psnr(a, b):\n a = normalize(a)\n b = normalize(b)\n e = a-b\n n = np.mean(np.multiply(e, e))\n return 10*np.log(1/n)/np.log(10)\n\ndef normalize(i):\n i = i.astype(np.float32)\n return i/np.max(i)\n\ndef main(args):\n a = imread(args.input)\n b = imread(args.output)\n print(args.output, psnr(a, b))\n\nif __name__ == '__main__':\n args = parser.parse_args()\n main(args)\n\n \n" }, { "path": "src/inpaint.py", "content": "import tensorflow as tf\nimport scipy.misc\nimport argparse\nimport os\nimport numpy as np\nfrom glob import glob\n\nfrom model_inpaint import ModelInpaint\nfrom helper import loadimage, saveimages\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--model_file', type=str, help=\"Pretrained GAN model\")\nparser.add_argument('--lr', type=float, default=0.0003)\nparser.add_argument('--momentum', type=float, default=0.9)\nparser.add_argument('--nIter', type=int, default=3000)\nparser.add_argument('--imgSize', type=int, default=64)\nparser.add_argument('--batch_size', type=int, default=64)\nparser.add_argument('--lambda_p', type=float, default=0.2)\nparser.add_argument('--checkpointDir', type=str, default='checkpoint')\nparser.add_argument('--outDir', type=str, default='completions')\nparser.add_argument('--blend', action='store_true', default=False,\n help=\"Blend predicted image to original image\")\nparser.add_argument('--maskType', type=str,\n choices=['random', 'center', 'left', 'file'],\n default='center')\nparser.add_argument('--maskFile', type=str,\n default=None,\n help='Input binary mask for file mask type')\nparser.add_argument('--maskThresh', type=int,\n default=128,\n help='Threshold in case input mask is not binary')\nparser.add_argument('--in_image', type=str, default=None,\n help='Input Image (ignored if inDir is specified')\nparser.add_argument('--inDir', type=str, default=None,\n help='Path to input images')\nparser.add_argument('--imgExt', type=str, default='png',\n help='input images file extension')\n\nargs = parser.parse_args()\n\n\n#def loadimage(filename):\n# img = scipy.misc.imread(filename, mode='RGB').astype(np.float)\n# return img\n#\n#\n#def saveimages(outimages, prefix='samples'):\n# numimages = len(outimages)\n#\n# if not os.path.exists(args.outDir):\n# os.mkdir(args.outDir)\n#\n# for i in range(numimages):\n# filename = '{}_{}.png'.format(prefix, i)\n# filename = os.path.join(args.outDir, filename)\n# scipy.misc.imsave(filename, outimages[i, :, :, :])\n\n\ndef gen_mask(maskType):\n image_shape = [args.imgSize, args.imgSize]\n if maskType == 'random':\n fraction_masked = 0.2\n mask = np.ones(image_shape)\n mask[np.random.random(image_shape[:2]) < fraction_masked] = 0.0\n elif maskType == 'center':\n scale = 0.25\n assert(scale <= 0.5)\n mask = np.ones(image_shape)\n sz = args.imgSize\n l = int(args.imgSize*scale)\n u = int(args.imgSize*(1.0-scale))\n mask[l:u, l:u] = 0.0\n elif maskType == 'left':\n mask = np.ones(image_shape)\n c = args.imgSize // 2\n mask[:, :c] = 0.0\n elif maskType == 'file':\n mask = loadmask(args.maskfile, args.maskthresh)\n else:\n assert(False)\n return mask\n\n\ndef loadmask(filename, thresh=128):\n immask = scipy.misc.imread(filename, mode='L')\n image_shape = [args.imgSize, args.imgSize]\n mask = np.ones(image_shape)\n mask[immask < 128] = 0\n mask[immaks >= 128] = 1\n return mask\n\n\ndef main():\n m = ModelInpaint(args.model_file, args)\n\n # Generate some samples from the model as a test\n #imout = m.sample()\n #saveimages(imout)\n\n mask = gen_mask(args.maskType)\n if args.inDir is not None:\n imgfilenames = glob( args.inDir + '/*.' + args.imgExt )\n print('{} images found'.format(len(imgfilenames)))\n in_img = np.array([loadimage(f) for f in imgfilenames])\n elif args.in_img is not None:\n in_img = loadimage(args.in_image)\n else:\n print('Input image needs to be specified')\n exit(1)\n\n #saveimages(in_img*mask[:,:,np.newaxis], 'inpaint_in', imgfilenames, args.outDir)\n inpaint_out, g_out = m.restore_image(in_img, mask, args.blend)\n scipy.misc.imsave(os.path.join(args.outDir, 'mask.png'), mask)\n\n saveimages(g_out, 'inpaint_gen', imgfilenames, args.outDir)\n saveimages(inpaint_out, 'inpaint', imgfilenames, args.outDir)\n\nif __name__ == '__main__':\n main()\n" }, { "path": "src/inpaint_test.py", "content": "import tensorflow as tf\nimport scipy.misc\nimport argparse\nimport os\nimport numpy as np\nfrom glob import glob\n\nfrom model_inpaint_test import ModelInpaintTest as ModelInpaint\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--model_file', type=str, help=\"Pretrained GAN model\")\nparser.add_argument('--lr', type=float, default=0.01)\nparser.add_argument('--momentum', type=float, default=0.9)\nparser.add_argument('--nIter', type=int, default=1000)\nparser.add_argument('--imgSize', type=int, default=64)\nparser.add_argument('--batch_size', type=int, default=64)\nparser.add_argument('--lambda_p', type=float, default=0.1)\nparser.add_argument('--checkpointDir', type=str, default='checkpoint')\nparser.add_argument('--outDir', type=str, default='completions')\nparser.add_argument('--blend', action='store_true', default=False,\n help=\"Blend predicted image to original image\")\nparser.add_argument('--maskType', type=str,\n choices=['random', 'center', 'left', 'file'],\n default='center')\nparser.add_argument('--maskFile', type=str,\n default=None,\n help='Input binary mask for file mask type')\nparser.add_argument('--maskThresh', type=int,\n default=128,\n help='Threshold in case input mask is not binary')\nparser.add_argument('--in_image', type=str, default=None,\n help='Input Image (ignored if inDir is specified')\nparser.add_argument('--inDir', type=str, default=None,\n help='Path to input images')\nparser.add_argument('--imgExt', type=str, default='png',\n help='input images file extension')\n\nargs = parser.parse_args()\n\n\ndef loadimage(filename):\n img = scipy.misc.imread(filename, mode='RGB').astype(np.float)\n return img\n\n\ndef saveimages(outimages, prefix='samples'):\n numimages = len(outimages)\n\n if not os.path.exists(args.outDir):\n os.mkdir(args.outDir)\n\n for i in range(numimages):\n filename = '{}_{}.png'.format(prefix, i)\n filename = os.path.join(args.outDir, filename)\n scipy.misc.imsave(filename, outimages[i, :, :, :])\n\n\ndef gen_mask(maskType):\n image_shape = [args.imgSize, args.imgSize]\n if maskType == 'random':\n fraction_masked = 0.2\n mask = np.ones(image_shape)\n mask[np.random.random(image_shape[:2]) < fraction_masked] = 0.0\n elif maskType == 'center':\n scale = 0.25\n assert(scale <= 0.5)\n mask = np.ones(image_shape)\n sz = args.imgSize\n l = int(args.imgSize*scale)\n u = int(args.imgSize*(1.0-scale))\n mask[l:u, l:u] = 0.0\n elif maskType == 'left':\n mask = np.ones(image_shape)\n c = args.imgSize // 2\n mask[:, :c] = 0.0\n elif maskType == 'file':\n mask = loadmask(args.maskfile, args.maskthresh)\n else:\n assert(False)\n return mask\n\n\ndef loadmask(filename, thresh=128):\n immask = scipy.misc.imread(filename, mode='L')\n image_shape = [args.imgSize, args.imgSize]\n mask = np.ones(image_shape)\n mask[immask < 128] = 0\n mask[immaks >= 128] = 1\n return mask\n\n\ndef main():\n m = ModelInpaint(args.model_file, args)\n\n # Generate some samples from the model as a test\n imout = m.sample()\n saveimages(imout)\n\n mask = gen_mask(args.maskType)\n if args.inDir is not None:\n imgfilenames = glob( args.inDir + '/*.' + args.imgExt )\n print('{} images found'.format(len(imgfilenames)))\n in_img = np.array([loadimage(f) for f in imgfilenames])\n elif args.in_img is not None:\n in_img = loadimage(args.in_image)\n else:\n print('Input image needs to be specified')\n exit(1)\n\n inpaint_out, g_out = m.restore_image(in_img, mask, args.blend)\n scipy.misc.imsave(os.path.join(args.outDir, 'mask.png'), mask)\n saveimages(g_out, 'gen')\n saveimages(inpaint_out, 'inpaint')\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "src/model.py", "content": "import tensorflow as tf\nimport numpy as np\nimport external.poissonblending as blending\nfrom scipy.signal import convolve2d\n\n\nclass ModelInpaint():\n def __init__(self, modelfilename, config,\n model_name='dcgan',\n gen_input='z:0', gen_output='Tanh:0', gen_loss='Mean_2:0',\n disc_input='real_images:0', disc_output='Sigmoid:0',\n z_dim=100, batch_size=64):\n \"\"\"\n Model for Semantic image inpainting.\n Loads frozen weights of a GAN and create the graph according to the\n loss function as described in paper\n\n Arguments:\n modelfilename - tensorflow .pb file with weights to be loaded\n config - training parameters: lambda_p, nIter\n gen_input - node name for generator input\n gen_output - node name for generator output\n disc_input - node name for discriminator input\n disc_output - node name for discriminator output\n z_dim - latent space dimension of GAN\n batch_size - training batch size\n \"\"\"\n\n self.config = config\n\n self.batch_size = batch_size\n self.z_dim = z_dim\n self.graph, self.graph_def = ModelInpaint.loadpb(modelfilename,\n model_name)\n\n self.gi = self.graph.get_tensor_by_name(model_name+'/'+gen_input)\n self.go = self.graph.get_tensor_by_name(model_name+'/'+gen_output)\n self.gl = self.graph.get_tensor_by_name(model_name+'/'+gen_loss)\n self.di = self.graph.get_tensor_by_name(model_name+'/'+disc_input)\n self.do = self.graph.get_tensor_by_name(model_name+'/'+disc_output)\n\n self.image_shape = self.go.shape[1:].as_list()\n\n self.l = config.lambda_p\n\n self.sess = tf.Session(graph=self.graph)\n\n self.init_z()\n\n def init_z(self):\n \"\"\"Initializes latent variable z\"\"\"\n self.z = np.random.randn(self.batch_size, self.z_dim)\n\n def sample(self, z=None):\n \"\"\"GAN sampler. Useful for checking if the GAN was loaded correctly\"\"\"\n if z is None:\n z = self.z\n sample_out = self.sess.run(self.go, feed_dict={self.gi: z})\n return sample_out\n\n def preprocess(self, images, imask, useWeightedMask = True, nsize=7):\n \"\"\"Default preprocessing pipeline\n Prepare the data to be fed to the network. Weighted mask is computed\n and images and masks are duplicated to fill the batch.\n\n Arguments:\n image - input image\n mask - input mask\n\n Returns:\n None\n \"\"\"\n images = ModelInpaint.imtransform(images)\n if useWeightedMask:\n mask = ModelInpaint.createWeightedMask(imask, nsize)\n else:\n mask = imask\n mask = ModelInpaint.create3ChannelMask(mask)\n \n bin_mask = ModelInpaint.binarizeMask(imask, dtype='uint8')\n self.bin_mask = ModelInpaint.create3ChannelMask(bin_mask)\n\n self.masks_data = np.repeat(mask[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n\n #Generate multiple candidates for completion if single image is given\n if len(images.shape) is 3:\n self.images_data = np.repeat(images[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n elif len(images.shape) is 4:\n #Ensure batch is filled\n num_images = images.shape[0]\n self.images_data = np.repeat(images[np.newaxis, 0, :, :, :],\n self.batch_size,\n axis=0)\n ncpy = min(num_images, self.batch_size)\n self.images_data[:ncpy, :, :, :] = images[:ncpy, :, :, :].copy()\n\n def postprocess(self, g_out, blend = True):\n \"\"\"Default post processing pipeline\n Applies poisson blending using binary mask. (default)\n\n Arguments:\n g_out - generator output\n blend - Use poisson blending (True) or alpha blending (False)\n \"\"\"\n images_out = ModelInpaint.iminvtransform(g_out)\n images_in = ModelInpaint.iminvtransform(self.images_data)\n\n if blend:\n for i in range(len(g_out)):\n images_out[i] = ModelInpaint.poissonblending(\n images_in[i], images_out[i], self.bin_mask\n )\n else:\n images_out = np.multiply(images_out, 1-self.masks_data) \\\n + np.multiply(images_in, self.masks_data)\n\n return images_out\n\n def build_inpaint_graph(self):\n \"\"\"Builds the context and prior loss objective\"\"\"\n with self.graph.as_default():\n self.masks = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='mask')\n self.images = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='images')\n self.context_loss = tf.reduce_sum(\n tf.contrib.layers.flatten(\n tf.abs(tf.multiply(self.masks, self.go) -\n tf.multiply(self.masks, self.images))), 1\n )\n\n self.perceptual_loss = self.gl\n self.inpaint_loss = self.context_loss + self.l*self.perceptual_loss\n self.inpaint_grad = tf.gradients(self.inpaint_loss, self.gi)\n\n def inpaint(self, image, mask, blend=True):\n \"\"\"Perform inpainting with the given image and mask with the standard\n pipeline as described in paper. To skip steps or try other pre/post\n processing, the methods can be called seperately.\n\n Arguments:\n image - input 3 channel image\n mask - input binary mask, single channel. Nonzeros values are \n treated as 1\n blend - Flag to apply Poisson blending on output, Default = True\n\n Returns:\n post processed image (merged/blneded), raw generator output\n \"\"\"\n self.build_inpaint_graph()\n self.preprocess(image, mask)\n\n imout = self.backprop_to_input()\n\n return self.postprocess(imout, blend), imout\n\n def backprop_to_input(self, verbose=True):\n \"\"\"Main worker function. To be called after all initilization is done.\n Performs backpropagation to input using (accelerated) gradient descent\n to obtain latent space representation of target image\n\n Returns:\n generator output image\n \"\"\"\n v = 0\n for i in range(self.config.nIter):\n out_vars = [self.inpaint_loss, self.inpaint_grad, self.go]\n in_dict = {self.masks: self.masks_data,\n self.gi: self.z,\n self.images: self.images_data}\n\n loss, grad, imout = self.sess.run(out_vars, feed_dict=in_dict)\n\n v_prev = np.copy(v)\n v = self.config.momentum*v - self.config.lr*grad[0]\n self.z += (-self.config.momentum * v_prev +\n (1 + self.config.momentum) * v)\n self.z = np.clip(self.z, -1, 1)\n\n if verbose:\n print('Iteration {}: {}'.format(i, np.mean(loss)))\n\n return imout\n\n @staticmethod\n def loadpb(filename, model_name='dcgan'):\n \"\"\"Loads pretrained graph from ProtoBuf file\n\n Arguments:\n filename - path to ProtoBuf graph definition\n model_name - prefix to assign to loaded graph node names\n\n Returns:\n graph, graph_def - as per Tensorflow definitions\n \"\"\"\n with tf.gfile.GFile(filename, 'rb') as f:\n graph_def = tf.GraphDef()\n graph_def.ParseFromString(f.read())\n\n with tf.Graph().as_default() as graph:\n tf.import_graph_def(graph_def,\n input_map=None,\n return_elements=None,\n op_dict=None,\n producer_op_list=None,\n name=model_name)\n\n return graph, graph_def\n\n @staticmethod\n def imtransform(img):\n \"\"\"Helper: Rescale pixel value ranges to -1 and 1\"\"\"\n return np.array(img) / 127.5-1\n\n @staticmethod\n def iminvtransform(img):\n \"\"\"Helper: Rescale pixel value ranges to 0 and 1\"\"\"\n return (np.array(img) + 1.0) / 2.0\n\n @staticmethod\n def poissonblending(img1, img2, mask):\n \"\"\"Helper: interface to external poisson blending\"\"\"\n return blending.blend(img1, img2, 1 - mask)\n\n @staticmethod\n def createWeightedMask(mask, nsize=7):\n \"\"\"Takes binary weighted mask to create weighted mask as described in \n paper.\n\n Arguments:\n mask - binary mask input. numpy float32 array\n nsize - pixel neighbourhood size. default = 7\n \"\"\"\n ker = np.ones((nsize,nsize), dtype=np.float32)\n ker = ker/np.sum(ker)\n wmask = mask * convolve2d(mask, ker, mode='same', boundary='symm')\n return wmask\n\n @staticmethod\n def binarizeMask(mask, dtype=np.float32):\n \"\"\"Helper function, ensures mask is 0/1 or 0/255 and single channel\n If dtype specified as float32 (default), output mask will be 0, 1\n if required dtype is uint8, output mask will be 0, 255\n \"\"\"\n assert(np.dtype(dtype) == np.float32 or np.dtype(dtype) == np.uint8)\n bmask = np.array(mask, dtype=np.float32)\n bmask[bmask>0] = 1.0\n bmask[bmask<=0] = 0\n if dtype == np.uint8:\n bmask = np.array(bmask*255, dtype=np.uint8)\n return bmask\n \n @staticmethod\n def create3ChannelMask(mask):\n \"\"\"Helper function, repeats single channel mask to 3 channels\"\"\"\n assert(len(mask.shape)==2)\n return np.repeat(mask[:,:,np.newaxis], 3, axis=2)\n" }, { "path": "src/model_base.py", "content": "\"\"\"Model base class.\"\"\"\n\nimport tensorflow as tf\nimport numpy as np\nimport tools\nimport abc\n\nclass ModelBase(object):\n def __init__(self, modelfilename, config,\n model_name='dcgan',\n gen_input='z:0', gen_output='Tanh:0', gen_loss='Mean_2:0',\n disc_input='real_images:0', disc_output='Sigmoid:0',\n z_dim=100, batch_size=64):\n \"\"\"\n Model for Semantic image inpainting.\n Loads frozen weights of a GAN and create the graph according to the\n loss function as described in paper\n\n Args:\n modelfilename: tensorflow .pb file with weights to be loaded\n config: training parameters: lambda_p, nIter\n gen_input: node name for generator input\n gen_output: node name for generator output\n disc_input: node name for discriminator input\n disc_output: node name for discriminator output\n z_dim: latent space dimension of GAN\n batch_size: training batch size\n \"\"\"\n __metaclass__ = abc.ABCMeta\n self.config = config\n\n self.batch_size = batch_size\n self.z_dim = z_dim\n self.graph, self.graph_def = tools.loadpb(modelfilename,\n model_name)\n\n self.gi = self.graph.get_tensor_by_name(model_name+'/'+gen_input)\n self.go = self.graph.get_tensor_by_name(model_name+'/'+gen_output)\n self.gl = self.graph.get_tensor_by_name(model_name+'/'+gen_loss)\n self.di = self.graph.get_tensor_by_name(model_name+'/'+disc_input)\n self.do = self.graph.get_tensor_by_name(model_name+'/'+disc_output)\n\n self.image_shape = self.go.shape[1:].as_list()\n\n self.l = config.lambda_p\n\n self.sess = tf.Session(graph=self.graph)\n\n self.init_z()\n\n def init_z(self):\n \"\"\"Initializes latent variable z\"\"\"\n self.z = np.random.randn(self.batch_size, self.z_dim)\n #self.z = np.random.uniform(size=[self.batch_size, self.z_dim])\n\n def sample(self, z=None):\n \"\"\"GAN sampler. Useful for checking if the GAN was loaded correctly\"\"\"\n if z is None:\n z = self.z\n sample_out = self.sess.run(self.go, feed_dict={self.gi: z})\n return sample_out\n\n @abc.abstractmethod\n def preprocess(self, **kwargs):\n \"\"\"\n \"\"\"\n pass\n\n @abc.abstractmethod\n def postprocess(self, g_out, **kwargs):\n \"\"\"\n \"\"\"\n pass\n\n @abc.abstractmethod\n def build_context_loss(self):\n \"\"\"Builds context loss function.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def build_input_placeholders(self):\n pass\n\n @abc.abstractmethod\n def restore_image(self, image, **kwargs):\n \"\"\"\n \"\"\"\n pass\n\n def build_restore_graph(self):\n \"\"\"\n \"\"\"\n self.build_input_placeholders()\n self.build_context_loss()\n\n self.perceptual_loss = self.gl\n self.inpaint_loss = self.context_loss + self.l*self.perceptual_loss\n self.inpaint_grad = tf.gradients(self.inpaint_loss, self.gi)\n\n\n def backprop_to_input(self, verbose=True):\n \"\"\"Main worker function. To be called after all initilization is done.\n\n Performs backpropagation to input using (accelerated) gradient descent\n to obtain latent space representation of target image\n\n Returns:\n generator output image\n \"\"\"\n v = 0\n for i in range(self.config.nIter):\n out_vars = [self.inpaint_loss, self.inpaint_grad, self.go]\n if hasattr(self, 'masks_data'):\n in_dict = {self.masks: self.masks_data,\n self.gi: self.z,\n self.images: self.images_data}\n else:\n in_dict = {self.gi: self.z,\n self.images: self.images_data}\n\n\n loss, grad, imout = self.sess.run(out_vars, feed_dict=in_dict)\n\n v_prev = np.copy(v)\n v = self.config.momentum*v - self.config.lr*grad[0]\n self.z += (-self.config.momentum * v_prev +\n (1 + self.config.momentum) * v)\n self.z = np.clip(self.z, -1, 1)\n\n if verbose:\n print('Iteration {}: {}'.format(i, np.mean(loss)))\n return imout\n" }, { "path": "src/model_colorize.py", "content": "import tensorflow as tf\nimport numpy as np\nimport external.poissonblending as blending\nfrom scipy.signal import convolve2d\nfrom PIL import Image\n\n\nclass ModelColorize():\n def __init__(self, modelfilename, config,\n model_name='dcgan',\n gen_input='z:0', gen_output='Tanh:0', gen_loss='Mean_2:0',\n disc_input='real_images:0', disc_output='Sigmoid:0',\n z_dim=100, batch_size=64):\n \"\"\"\n Model for Semantic image inpainting.\n Loads frozen weights of a GAN and create the graph according to the\n loss function as described in paper\n\n Arguments:\n modelfilename - tensorflow .pb file with weights to be loaded\n config - training parameters: lambda_p, nIter\n gen_input - node name for generator input\n gen_output - node name for generator output\n disc_input - node name for discriminator input\n disc_output - node name for discriminator output\n z_dim - latent space dimension of GAN\n batch_size - training batch size\n \"\"\"\n\n self.config = config\n\n self.batch_size = batch_size\n self.z_dim = z_dim\n self.graph, self.graph_def = ModelColorize.loadpb(modelfilename,\n model_name)\n\n self.gi = self.graph.get_tensor_by_name(model_name+'/'+gen_input)\n self.go = self.graph.get_tensor_by_name(model_name+'/'+gen_output)\n self.gl = self.graph.get_tensor_by_name(model_name+'/'+gen_loss)\n self.di = self.graph.get_tensor_by_name(model_name+'/'+disc_input)\n self.do = self.graph.get_tensor_by_name(model_name+'/'+disc_output)\n\n self.image_shape = self.go.shape[1:-1].as_list() + [1]\n\n self.l = config.lambda_p\n\n self.sess = tf.Session(graph=self.graph)\n\n self.init_z()\n\n def init_z(self):\n \"\"\"Initializes latent variable z\"\"\"\n self.z = np.random.randn(self.batch_size, self.z_dim)\n #self.z = np.random.uniform(size=[self.batch_size, self.z_dim])\n\n def sample(self, z=None):\n \"\"\"GAN sampler. Useful for checking if the GAN was loaded correctly\"\"\"\n if z is None:\n z = self.z\n sample_out = self.sess.run(self.go, feed_dict={self.gi: z})\n return sample_out\n\n def preprocess(self, images):\n \"\"\"Default preprocessing pipeline\n Converts input image from single channel to 3 channel grayscale\n Prepare the data to be fed to the network. Weighted mask is computed\n and images and masks are duplicated to fill the batch.\n\n Arguments:\n image - input image\n\n Returns:\n None\n \"\"\"\n images = ModelColorize.imtransform(images)\n\n #Generate multiple candidates for completion if single image is given\n if len(images.shape) is 3:\n self.images_data = np.repeat(images[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n elif len(images.shape) is 4:\n #Ensure batch is filled\n num_images = images.shape[0]\n self.images_data = np.repeat(images[np.newaxis, 0, :, :, :],\n self.batch_size,\n axis=0)\n ncpy = min(num_images, self.batch_size)\n self.images_data[:ncpy, :, :, :] = images[:ncpy, :, :, :].copy()\n\n def postprocess(self, g_out, blend=False):\n \"\"\"Default post processing pipeline\n Currently does nothing\n\n Arguments:\n g_out - generator output\n \"\"\"\n images_out = ModelColorize.iminvtransform(g_out)\n images_in = ModelColorize.iminvtransform(self.images_data)\n\n if blend:\n for idx, (i, o) in enumerate(zip(images_in, images_out)):\n images_out[idx, :, :, :] = ModelColorize.colorblend( i, o )\n\n return images_out\n\n def build_colorization_graph(self):\n \"\"\"Builds the context and prior loss objective\"\"\"\n with self.graph.as_default():\n self.images = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='images')\n self.context_loss = tf.reduce_sum(\n tf.contrib.layers.flatten(\n tf.abs(tf.image.rgb_to_grayscale(self.go) -\n self.images)), 1\n )\n\n self.prior_loss = self.gl\n self.colorization_loss = self.context_loss + self.l*self.prior_loss\n self.colorization_grad = tf.gradients(self.colorization_loss, self.gi)\n\n def colorize(self, image, blend=False):\n \"\"\"Perform inpainting with the given image and mask with the standard\n pipeline as described in paper. To skip steps or try other pre/post\n processing, the methods can be called seperately.\n\n Arguments:\n image - input 3 channel image\n mask - input binary mask, single channel. Nonzeros values are \n treated as 1\n blend - Flag to apply Poisson blending on output, Default = True\n\n Returns:\n post processed image (merged/blneded), raw generator output\n \"\"\"\n self.build_colorization_graph()\n self.preprocess(image)\n\n imout = self.backprop_to_input()\n\n return self.postprocess(imout, blend), imout\n\n def backprop_to_input(self, verbose=True):\n \"\"\"Main worker function. To be called after all initilization is done.\n Performs backpropagation to input using (accelerated) gradient descent\n to obtain latent space representation of target image\n\n Returns:\n generator output image\n \"\"\"\n v = 0\n for i in range(self.config.nIter):\n out_vars = [self.colorization_loss, self.colorization_grad, self.go]\n in_dict = {self.gi: self.z,\n self.images: self.images_data}\n\n loss, grad, imout = self.sess.run(out_vars, feed_dict=in_dict)\n\n v_prev = np.copy(v)\n v = self.config.momentum*v - self.config.lr*grad[0]\n self.z += (-self.config.momentum * v_prev +\n (1 + self.config.momentum) * v)\n self.z = np.clip(self.z, -1, 1)\n\n if verbose:\n print('Iteration {}: {}'.format(i, np.mean(loss)))\n\n return imout\n\n @staticmethod\n def loadpb(filename, model_name='dcgan'):\n \"\"\"Loads pretrained graph from ProtoBuf file\n\n Arguments:\n filename - path to ProtoBuf graph definition\n model_name - prefix to assign to loaded graph node names\n\n Returns:\n graph, graph_def - as per Tensorflow definitions\n \"\"\"\n with tf.gfile.GFile(filename, 'rb') as f:\n graph_def = tf.GraphDef()\n graph_def.ParseFromString(f.read())\n\n with tf.Graph().as_default() as graph:\n tf.import_graph_def(graph_def,\n input_map=None,\n return_elements=None,\n op_dict=None,\n producer_op_list=None,\n name=model_name)\n\n return graph, graph_def\n\n @staticmethod\n def imtransform(img):\n \"\"\"Helper: Rescale pixel value ranges to -1 and 1\"\"\"\n return np.array(img) / 127.5-1\n\n @staticmethod\n def iminvtransform(img):\n \"\"\"Helper: Rescale pixel value ranges to 0 and 1\"\"\"\n return (np.array(img) + 1.0) / 2.0\n\n @staticmethod\n def poissonblending(img1, img2, mask):\n \"\"\"Helper: interface to external poisson blending\"\"\"\n return blending.blend(img1, img2, 1 - mask)\n\n @staticmethod\n def colorblend(img_gray, img_color):\n \"\"\"Helper to apply color from one image to another\"\"\"\n img_blended_hsv = np.zeros_like(img_color, dtype=np.uint8)\n img_blended_hsv[:,:,2:] = np.uint8(np.copy(img_gray*255))\n img_color_hsv = np.array(Image.fromarray(np.uint8(img_color*255), mode='RGB').convert('HSV'))\n img_blended_hsv[:,:,:2] = img_color_hsv[:,:,:2]\n img_out = np.array(Image.fromarray(img_blended_hsv, mode='HSV').convert('RGB'))\n\n return img_out\n" }, { "path": "src/model_denoising.py", "content": "import tensorflow as tf\nimport numpy as np\nimport tools\nfrom scipy.ndimage.filters import gaussian_filter \n\nfrom model_base import ModelBase\n\nclass ModelDenoising(ModelBase):\n def preprocess(self, images, mask=None):\n \"\"\"Default preprocessing pipeline\n Prepare the data to be fed to the network. Weighted mask is computed\n and images and masks are duplicated to fill the batch.\n\n Arguments:\n image - input image\n mask - input mask\n\n Returns:\n None\n \"\"\"\n images = tools.imtransform(images)\n #Generate multiple candidates for completion if single image is given\n if len(images.shape) is 3:\n ii = np.repeat(images[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n elif len(images.shape) is 4:\n #Ensure batch is filled\n num_images = images.shape[0]\n ii = np.repeat(images[np.newaxis, 0, :, :, :],\n self.batch_size,\n axis=0)\n ncpy = min(num_images, self.batch_size)\n ii[:ncpy, :, :, :] = images[:ncpy, :, :, :].copy()\n\n self.images_data=np.stack([gaussian_filter(i, self.sigma) for i in ii])\n\n def postprocess(self, g_out, blend = True):\n \"\"\"Default post processing pipeline\n Applies poisson blending using binary mask. (default)\n\n Arguments:\n g_out - generator output\n blend - Use poisson blending (True) or alpha blending (False)\n \"\"\"\n images_out = tools.iminvtransform(g_out)\n images_in = tools.iminvtransform(self.images_data)\n return images_out\n\n def build_input_placeholders(self):\n with self.graph.as_default():\n self.masks = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='mask')\n self.images = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='images')\n\n def perform_corruption(self, images):\n return images\n\n def build_context_loss(self):\n \"\"\"Builds the context and prior loss objective\"\"\"\n with self.graph.as_default():\n self.context_loss = tf.reduce_sum(\n tf.contrib.layers.flatten(\n tf.abs(self.perform_corruption(self.go) -\n self.perform_corruption(self.images))), 1\n )\n\n def restore_image(self, image, mask=None, blend=True):\n \"\"\"Perform inpainting with the given image and mask with the standard\n pipeline as described in paper. To skip steps or try other pre/post\n processing, the methods can be called seperately.\n\n Arguments:\n image - input 3 channel image\n mask - input binary mask, single channel. Nonzeros values are\n treated as 1\n blend - Flag to apply Poisson blending on output, Default = True\n\n Returns:\n post processed image (merged/blneded), raw generator output\n \"\"\"\n self.build_restore_graph()\n self.preprocess(image, mask)\n\n imout = self.backprop_to_input()\n\n return self.postprocess(imout, blend), imout\n" }, { "path": "src/model_inpaint.py", "content": "import tensorflow as tf\nimport numpy as np\nimport tools\n\nfrom model_base import ModelBase\n\nclass ModelInpaint(ModelBase):\n def preprocess(self, images, imask, useWeightedMask=True, nsize=15):\n \"\"\"Default preprocessing pipeline\n Prepare the data to be fed to the network. Weighted mask is computed\n and images and masks are duplicated to fill the batch.\n\n Arguments:\n image - input image\n mask - input mask\n\n Returns:\n None\n \"\"\"\n images = tools.imtransform(images)\n if useWeightedMask:\n mask = tools.createWeightedMask(imask, nsize)\n else:\n mask = imask\n mask = tools.create3ChannelMask(mask)\n\n bin_mask = tools.binarizeMask(imask, dtype='uint8')\n self.bin_mask = tools.create3ChannelMask(bin_mask)\n\n self.masks_data = np.repeat(mask[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n\n #Generate multiple candidates for completion if single image is given\n if len(images.shape) is 3:\n self.images_data = np.repeat(images[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n elif len(images.shape) is 4:\n #Ensure batch is filled\n num_images = images.shape[0]\n self.images_data = np.repeat(images[np.newaxis, 0, :, :, :],\n self.batch_size,\n axis=0)\n ncpy = min(num_images, self.batch_size)\n self.images_data[:ncpy, :, :, :] = images[:ncpy, :, :, :].copy()\n\n def postprocess(self, g_out, blend = True):\n \"\"\"Default post processing pipeline\n Applies poisson blending using binary mask. (default)\n\n Arguments:\n g_out - generator output\n blend - Use poisson blending (True) or alpha blending (False)\n \"\"\"\n images_out = tools.iminvtransform(g_out)\n images_in = tools.iminvtransform(self.images_data)\n\n if blend:\n for i in range(len(g_out)):\n images_out[i] = tools.poissonblending(\n images_in[i], images_out[i], self.bin_mask\n )\n else:\n images_out = np.multiply(images_out, 1-self.masks_data) \\\n + np.multiply(images_in, self.masks_data)\n\n return images_out\n\n def build_input_placeholders(self):\n with self.graph.as_default():\n self.masks = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='mask')\n self.images = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='images')\n\n def build_context_loss(self):\n \"\"\"Builds the context and prior loss objective\"\"\"\n with self.graph.as_default():\n self.context_loss = tf.reduce_sum(\n tf.contrib.layers.flatten(\n tf.abs(tf.multiply(self.masks, self.go) -\n tf.multiply(self.masks, self.images))), 1\n )\n\n def restore_image(self, image, mask, blend=True):\n \"\"\"Perform inpainting with the given image and mask with the standard\n pipeline as described in paper. To skip steps or try other pre/post\n processing, the methods can be called seperately.\n\n Arguments:\n image - input 3 channel image\n mask - input binary mask, single channel. Nonzeros values are\n treated as 1\n blend - Flag to apply Poisson blending on output, Default = True\n\n Returns:\n post processed image (merged/blneded), raw generator output\n \"\"\"\n self.build_restore_graph()\n self.preprocess(image, mask)\n\n imout = self.backprop_to_input()\n\n return self.postprocess(imout, blend), imout\n" }, { "path": "src/model_inpaint_test.py", "content": "import tensorflow as tf\nimport numpy as np\nimport tools\n\nfrom model_inpaint import ModelInpaint\n\ndef p_standard_gaussian(z):\n\n return tf.exp( -.5 * tf.reduce_sum( tf.square(z), axis=1 ) )\n\nclass ModelInpaintTest(ModelInpaint):\n \"\"\"Test of small modificaiton to prior loss\"\"\"\n\n def build_restore_graph(self):\n super(ModelInpaintTest, self).build_restore_graph()\n\n p = p_standard_gaussian\n self.perceptual_loss = (self.gl + tf.log(tf.clip_by_value(\n 1-tf.exp(tf.clip_by_value(self.gl, -5, 5)), 1e-4, 99999)\n ))*p(self.gi)\n\n self.inpaint_loss = self.context_loss + self.l*self.perceptual_loss\n self.inpaint_grad = tf.gradients(self.inpaint_loss, self.gi)\n\n" }, { "path": "src/model_quantize.py", "content": "import tensorflow as tf\nimport numpy as np\nimport tools\n\nfrom model_base import ModelBase\n\nclass ModelQuantize(ModelBase):\n def preprocess(self, images, mask=None):\n \"\"\"Default preprocessing pipeline\n Prepare the data to be fed to the network. Weighted mask is computed\n and images and masks are duplicated to fill the batch.\n\n Arguments:\n image - input image\n mask - input mask\n\n Returns:\n None\n \"\"\"\n images = tools.imtransform(images)\n #Generate multiple candidates for completion if single image is given\n if len(images.shape) is 3:\n self.images_data = np.repeat(images[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n elif len(images.shape) is 4:\n #Ensure batch is filled\n num_images = images.shape[0]\n self.images_data = np.repeat(images[np.newaxis, 0, :, :, :],\n self.batch_size,\n axis=0)\n ncpy = min(num_images, self.batch_size)\n self.images_data[:ncpy, :, :, :] = images[:ncpy, :, :, :].copy()\n\n def postprocess(self, g_out, blend = True):\n \"\"\"Default post processing pipeline\n Applies poisson blending using binary mask. (default)\n\n Arguments:\n g_out - generator output\n blend - Use poisson blending (True) or alpha blending (False)\n \"\"\"\n images_out = tools.iminvtransform(g_out)\n images_in = tools.iminvtransform(self.images_data)\n return images_out\n\n def build_input_placeholders(self):\n with self.graph.as_default():\n self.masks = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='mask')\n self.images = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='images')\n\n def perform_corruption(self, images):\n images = (images + 1.0) / 2.0\n images *= (255.0 / self.quantize_factor)\n images = tf.floor(images)\n images /= (255.0 / self.quantize_factor)\n images *= 2.\n images -= 1.\n return images\n \n def perform_corruption_new(self, images):\n images = (images + 1.0)/2.0\n images = tf.floor(images*self.levels)/self.levels\n images *= 2.\n images -= 1.\n return images\n\n def build_context_loss(self):\n \"\"\"Builds the context and prior loss objective\"\"\"\n with self.graph.as_default():\n self.context_loss = tf.reduce_sum(\n tf.contrib.layers.flatten(\n tf.abs(self.go -\n self.perform_corruption(self.images))), 1\n )\n\n def restore_image(self, image, mask=None, blend=True):\n \"\"\"Perform inpainting with the given image and mask with the standard\n pipeline as described in paper. To skip steps or try other pre/post\n processing, the methods can be called seperately.\n\n Arguments:\n image - input 3 channel image\n mask - input binary mask, single channel. Nonzeros values are\n treated as 1\n blend - Flag to apply Poisson blending on output, Default = True\n\n Returns:\n post processed image (merged/blneded), raw generator output\n \"\"\"\n self.build_restore_graph()\n self.preprocess(image, mask)\n\n imout = self.backprop_to_input()\n\n return self.postprocess(imout, blend), imout\n" }, { "path": "src/model_superres.py", "content": "import tensorflow as tf\nimport numpy as np\nimport tools\n\nfrom model_base import ModelBase\n\nclass ModelSuperres(ModelBase):\n def preprocess(self, images, mask=None):\n \"\"\"Default preprocessing pipeline\n Prepare the data to be fed to the network. Weighted mask is computed\n and images and masks are duplicated to fill the batch.\n\n Arguments:\n image - input image\n mask - input mask\n\n Returns:\n None\n \"\"\"\n images = tools.imtransform(images)\n #Generate multiple candidates for completion if single image is given\n if len(images.shape) is 3:\n self.images_data = np.repeat(images[np.newaxis, :, :, :],\n self.batch_size,\n axis=0)\n elif len(images.shape) is 4:\n #Ensure batch is filled\n num_images = images.shape[0]\n self.images_data = np.repeat(images[np.newaxis, 0, :, :, :],\n self.batch_size,\n axis=0)\n ncpy = min(num_images, self.batch_size)\n self.images_data[:ncpy, :, :, :] = images[:ncpy, :, :, :].copy()\n\n def postprocess(self, g_out, blend = True):\n \"\"\"Default post processing pipeline\n Applies poisson blending using binary mask. (default)\n\n Arguments:\n g_out - generator output\n blend - Use poisson blending (True) or alpha blending (False)\n \"\"\"\n images_out = tools.iminvtransform(g_out)\n images_in = tools.iminvtransform(self.images_data)\n return images_out\n\n def build_input_placeholders(self):\n with self.graph.as_default():\n self.masks = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='mask')\n self.images = tf.placeholder(tf.float32,\n [None] + self.image_shape,\n name='images')\n\n def perform_corruption(self, images):\n down_sample_factor = self.down_sample_factor\n\n ret = tf.image.resize_bilinear(\n images,\n [tf.shape(images)[1] // down_sample_factor,\n tf.shape(images)[2] // down_sample_factor,\n ],\n align_corners=True)\n\n ret = tf.image.resize_bilinear(ret,\n [tf.shape(images)[1],\n tf.shape(images)[2]\n ],\n align_corners=True)\n return ret\n\n def build_context_loss(self):\n \"\"\"Builds the context and prior loss objective\"\"\"\n with self.graph.as_default():\n self.context_loss = tf.reduce_sum(\n tf.contrib.layers.flatten(\n tf.abs(self.perform_corruption(self.go) -\n self.perform_corruption(self.images))), 1\n )\n\n def restore_image(self, image, mask=None, blend=True):\n \"\"\"Perform inpainting with the given image and mask with the standard\n pipeline as described in paper. To skip steps or try other pre/post\n processing, the methods can be called seperately.\n\n Arguments:\n image - input 3 channel image\n mask - input binary mask, single channel. Nonzeros values are\n treated as 1\n blend - Flag to apply Poisson blending on output, Default = True\n\n Returns:\n post processed image (merged/blneded), raw generator output\n \"\"\"\n self.build_restore_graph()\n self.preprocess(image, mask)\n\n imout = self.backprop_to_input()\n\n return self.postprocess(imout, blend), imout\n" }, { "path": "src/quantize.py", "content": "import tensorflow as tf\nimport scipy.misc\nimport argparse\nimport os\nimport numpy as np\nfrom glob import glob\nfrom helper import loadimage, saveimages\n\nfrom model_quantize import ModelQuantize\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--model_file', type=str, help=\"Pretrained GAN model\")\nparser.add_argument('--lr', type=float, default=0.0001)\nparser.add_argument('--momentum', type=float, default=0.9)\nparser.add_argument('--nIter', type=int, default=1000)\nparser.add_argument('--imgSize', type=int, default=64)\nparser.add_argument('--batch_size', type=int, default=64)\nparser.add_argument('--lambda_p', type=float, default=0.001)\nparser.add_argument('--checkpointDir', type=str, default='checkpoint')\nparser.add_argument('--outDir', type=str, default='quantize')\nparser.add_argument('--blend', action='store_true', default=False,\n help=\"Blend predicted image to original image\")\nparser.add_argument('--in_image', type=str, default=None,\n help='Input Image (ignored if inDir is specified')\nparser.add_argument('--inDir', type=str, default=None,\n help='Path to input images')\nparser.add_argument('--imgExt', type=str, default='png',\n help='input images file extension')\nparser.add_argument('-c', action='store_true', help='corrupt image on the fly')\n\nargs = parser.parse_args()\n\n\ndef corrupt_image(img, quantize_factor=4):\n img /= 255.0\n img *= (255.0 / quantize_factor)\n images = img\n images = np.floor(images)\n images /= (255.0 / quantize_factor)\n images *= 255.0\n images = images.astype(int)\n return images\n\ndef main():\n m = ModelQuantize(args.model_file, args)\n dd = 50\n m.quantize_factor = dd\n #m.levels = 4\n\n # Generate some samples from the model as a test\n #imout = m.sample()\n #saveimages(imout)\n\n if args.inDir is not None:\n imgfilenames = glob( args.inDir + '/*.' + args.imgExt )\n print('{} images found'.format(len(imgfilenames)))\n in_img = np.array([loadimage(f) for f in imgfilenames])\n if args.c:\n in_corrupt_img = np.array([corrupt_image(loadimage(f), dd) for f in imgfilenames])\n else:\n in_corrupt_img = np.copy(in_img)\n\n elif args.in_image is not None:\n in_img = in_corrupt_img #loadimage(args.in_image)\n else:\n print('Input image needs to be specified')\n exit(1)\n #saveimages(in_corrupt_img, prefix='input')\n inpaint_out, g_out = m.restore_image(in_img)\n saveimages(g_out, 'quantize_gen', imgfilenames, args.outDir)\n saveimages(inpaint_out, 'quantize', imgfilenames, args.outDir)\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "src/superres.py", "content": "import tensorflow as tf\nimport scipy.misc\nimport argparse\nimport os\nimport numpy as np\nfrom glob import glob\nfrom helper import loadimage, saveimages\n\nfrom model_superres import ModelSuperres\n\nparser = argparse.ArgumentParser()\nparser.add_argument('--model_file', type=str, help=\"Pretrained GAN model\")\nparser.add_argument('--lr', type=float, default=0.0001)\nparser.add_argument('--momentum', type=float, default=0.9)\nparser.add_argument('--nIter', type=int, default=1000)\nparser.add_argument('--imgSize', type=int, default=64)\nparser.add_argument('--batch_size', type=int, default=64)\nparser.add_argument('--lambda_p', type=float, default=0.01)\nparser.add_argument('--checkpointDir', type=str, default='checkpoint')\nparser.add_argument('--outDir', type=str, default='superres')\nparser.add_argument('--blend', action='store_true', default=False,\n help=\"Blend predicted image to original image\")\nparser.add_argument('--in_image', type=str, default=None,\n help='Input Image (ignored if inDir is specified')\nparser.add_argument('--inDir', type=str, default=None,\n help='Path to input images')\nparser.add_argument('--imgExt', type=str, default='png',\n help='input images file extension')\nparser.add_argument('-c', action='store_true', help='corrupt image on the fly')\n\nargs = parser.parse_args()\n\n\ndef corrupt_image(img, down_sample_factor=4):\n ret = scipy.misc.imresize(img, 1./down_sample_factor)\n ret = scipy.misc.imresize(ret, 1.*down_sample_factor)\n return ret\n\ndef main():\n m = ModelSuperres(args.model_file, args)\n dd = 4\n m.down_sample_factor = dd\n\n # Generate some samples from the model as a test\n #imout = m.sample()\n #saveimages(imout)\n\n if args.inDir is not None:\n imgfilenames = glob( args.inDir + '/*.' + args.imgExt )\n print('{} images found'.format(len(imgfilenames)))\n in_img = np.array([loadimage(f) for f in imgfilenames])\n if args.c:\n in_corrupt_img = np.array([corrupt_image(loadimage(f), dd) for f in imgfilenames])\n else:\n in_corrupt_img = np.copy(in_img)\n\n\n elif args.in_image is not None:\n in_img = loadimage(args.in_image)\n else:\n print('Input image needs to be specified')\n exit(1)\n #saveimages(in_corrupt_img, prefix='input')\n inpaint_out, g_out = m.restore_image(in_img)\n\n saveimages(g_out, 'superres_gen', imgfilenames, args.outDir)\n saveimages(inpaint_out, 'superres', imgfilenames, args.outDir)\n\nif __name__ == '__main__':\n main()\n" }, { "path": "src/tools.py", "content": "\"\"\"Helper functions\n\"\"\"\nimport tensorflow as tf\nimport numpy as np\nimport external.poissonblending as blending\nfrom scipy.signal import convolve2d\n\n###################Helpers for image inapinting #######################\ndef imtransform(img):\n\t\"\"\"Helper: Rescale pixel value ranges to -1 and 1\"\"\"\n\treturn np.array(img) / 127.5-1\n\ndef iminvtransform(img):\n\t\"\"\"Helper: Rescale pixel value ranges to 0 and 1\"\"\"\n\treturn (np.array(img) + 1.0) / 2.0\n\ndef poissonblending(img1, img2, mask):\n\t\"\"\"Helper: interface to external poisson blending\"\"\"\n\treturn blending.blend(img1, img2, 1 - mask)\n\ndef createWeightedMask(mask, nsize=7):\n\t\"\"\"Takes binary weighted mask to create weighted mask as described in paper.\n\tArgs:\n\t\tmask: binary mask input. numpy float32 array\n\t\tnsize: pixel neighbourhood size. default = 7\n\t\"\"\"\n\tker = np.ones((nsize,nsize), dtype=np.float32)\n\tker = ker/np.sum(ker)\n\twmask = mask * convolve2d(1-mask, ker, mode='same', boundary='symm')\n\treturn wmask\n\ndef binarizeMask(mask, dtype=np.float32):\n\t\"\"\"Helper function, ensures mask is 0/1 or 0/255 and single channel.\n\n\tIf dtype specified as float32 (default), output mask will be 0, 1\n\tif required dtype is uint8, output mask will be 0, 255\n\n\tArgs:\n\t\tmask:.\n\t\tdtype:.\n\t\"\"\"\n\tassert(np.dtype(dtype) == np.float32 or np.dtype(dtype) == np.uint8)\n\tbmask = np.array(mask, dtype=np.float32)\n\tbmask[bmask>0] = 1.0\n\tbmask[bmask<=0] = 0\n\tif dtype == np.uint8:\n\t\tbmask = np.array(bmask*255, dtype=np.uint8)\n\treturn bmask\n\ndef create3ChannelMask(mask):\n\t\"\"\"Helper function, repeats single channel mask to 3 channels\"\"\"\n\tassert(len(mask.shape)==2)\n\treturn np.repeat(mask[:,:,np.newaxis], 3, axis=2)\n\ndef loadpb(filename, model_name='dcgan'):\n \"\"\"Loads pretrained graph from ProtoBuf file\n Args:\n filename: path to ProtoBuf graph definition.\n model_name: prefix to assign to loaded graph node names.\n Returns:\n graph, graph_def: as per Tensorflow definitions.\n \"\"\"\n with tf.gfile.GFile(filename, 'rb') as f:\n graph_def = tf.GraphDef()\n graph_def.ParseFromString(f.read())\n\n with tf.Graph().as_default() as graph:\n tf.import_graph_def(graph_def,\n input_map=None,\n return_elements=None,\n op_dict=None,\n producer_op_list=None,\n name=model_name)\n\n return graph, graph_def\n" } ]