Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,4 @@
+model
+logs
+data
+load_model.py

README.md

@@ -1,12 +1,85 @@
----
-title: Watermark Removal
-emoji: 🐨
-colorFrom: purple
-colorTo: purple
-sdk: gradio
-sdk_version: 5.49.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Chimzuruoke Okafor
+
+# Watermark-Removal
+
+<p>
+  <a href="https://pepy.tech/project/prompttools" target="_blank"><img src="https://static.pepy.tech/badge/prompttools" alt="Total Downloads"/></a>
+  <a href="https://github.com/hegelai/watermark-removal"><img src="https://img.shields.io/github/stars/zuruoke/watermark-removal" /></a>
+  <a href="https://twitter.com/zuruoke_okafor"><img src="https://img.shields.io/twitter/follow/Zuruoke_Okafor?style=social"></a>
+</p>
+
+![version](https://img.shields.io/badge/version-v1.0.0-green.svg?style=plastic)
+![pytorch](https://img.shields.io/badge/tensorflow-v1.15.0-green.svg?style=plastic)
+![license](https://img.shields.io/badge/license-CC_BY--NC-green.svg?style=plastic)
+
+<a href="https://coff.ee/zuruokeokafor" target="_blank">
+  <img src="https://www.buymeacoffee.com/assets/img/custom_images/orange_img.png"
+       alt="Buy Me A Coffee"
+       style="height: 41px;width:174px;box-shadow:0px 3px 2px rgba(190,190,190,0.5);" />
+</a>
+
+An open source project that uses a machine learning based image inpainting methodology to remove watermark from images which is totally indistinguishable from the ground truth version of the image.
+
+This project was inspired by the [Contextual Attention](https://arxiv.org/abs/1801.07892) (CVPR 2018) and [Gated Convolution](https://arxiv.org/abs/1806.03589) (ICCV 2019 Oral).
+
+And also a shoutout to [Chu-Tak Li](https://chutakcode.wixsite.com/website) for his [Medium article series](https://towardsdatascience.com/10-papers-you-must-read-for-deep-image-inpainting-2e41c589ced0) that really gave me a deep insight into the image inpainting papers stated above
+
+<img src="https://user-images.githubusercontent.com/51057490/140277713-c7d6e2b9-db62-4793-823a-25ed0c4e2771.png" width="45%"/> <img src="https://user-images.githubusercontent.com/51057490/140277781-5b5218bb-9044-4ec9-a349-eea93bc56d4a.png" width="45%"/> <img src="https://user-images.githubusercontent.com/51057490/140277929-3f187647-0e63-4bcb-b9f1-472f7558aae5.jpeg" width="45%"/> <img src="https://user-images.githubusercontent.com/51057490/140277957-6ddb7dec-25c8-42f1-8e39-be491d4f2248.png" width="45%"/> <img src="https://user-images.githubusercontent.com/51057490/140277983-265a1c9e-6093-4154-8252-838baca21c41.jpeg" width="45%" /> <img src="https://user-images.githubusercontent.com/51057490/140278002-56c4ae3d-6bfb-4ba3-aa02-7bd28474bfdf.png" width="45%" /> <img src="https://user-images.githubusercontent.com/51057490/140278030-d2a962ce-3722-43f1-b1bd-0ffde2aa7026.jpeg" width="45%" /> <img src="https://user-images.githubusercontent.com/51057490/140278040-10e401d7-4b7d-4d81-91fe-e9f01ef4ce7f.png" width="45%" /> <img src="https://user-images.githubusercontent.com/51057490/140278017-34862de0-86eb-40f0-b04b-7dc02fe38a77.jpeg" width="45%" /> <img src="https://user-images.githubusercontent.com/51057490/140278011-e0ae9ed0-e4ed-44ed-a9ac-28eb8456797a.png" width="45%" />
+
+## Run
+
+- use [Google colab](https://research.google.com/colaboratory/)
+
+- First of all, clone this repo
+
+      !git clone https://github.com/zuruoke/watermark-removal
+
+- Change Directory to the repo
+
+      !cd watermark-removal
+
+- Since Google Colab uses the latest Tensorflow 2x version and this project uses 1.15.0, downgrade to Tensorflow 1.15.0 version and restart the runtime, (`although the new version of Google Colab does not need you to restart the runtime`).
+
+      !pip install tensorflow==1.15.0
+
+- Install tensorflow toolkit [neuralgym](https://github.com/JiahuiYu/neuralgym).
+
+      !pip install git+https://github.com/JiahuiYu/neuralgym
+
+- Download the model dirs using this [link](https://drive.google.com/drive/folders/1xRV4EdjJuAfsX9pQme6XeoFznKXG0ptJ?usp=sharing) and put it under `model/` (rename `checkpoint.txt` to `checkpoint` because sometimes google drive automatically adds .txt after download)
+
+And you're all Set!!
+
+- Now remove the watermark on the image by runing the `main.py` file
+
+      !python main.py --image path-to-input-image --output path-to-output-image --checkpoint_dir model/ --watermark_type istock
+
+## Citing
+
+```
+@article{yu2018generative,
+  title={Generative Image Inpainting with Contextual Attention},
+  author={Yu, Jiahui and Lin, Zhe and Yang, Jimei and Shen, Xiaohui and Lu, Xin and Huang, Thomas S},
+  journal={arXiv preprint arXiv:1801.07892},
+  year={2018}
+}
+
+@article{yu2018free,
+  title={Free-Form Image Inpainting with Gated Convolution},
+  author={Yu, Jiahui and Lin, Zhe and Yang, Jimei and Shen, Xiaohui and Lu, Xin and Huang, Thomas S},
+  journal={arXiv preprint arXiv:1806.03589},
+  year={2018}
+}
+```
+
+<p align="center">
+<a href="https://star-history.com/#zuruoke/watermark-removal">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/svg?repos=zuruoke/watermark-removal&type=Date&theme=dark" />
+    <source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/svg?repos=zuruoke/watermark-removal&type=Date" />
+    <img alt="Star History Chart" src="https://api.star-history.com/svg?repos=zuruoke/watermark-removal&type=Date" />
+  </picture>
+</a>
+</p>
+
+## © Chimzuruoke Okafor

batch_test.py

@@ -0,0 +1,92 @@
+import time
+import os
+import argparse
+
+import cv2
+import numpy as np
+import tensorflow as tf
+import neuralgym as ng
+
+from inpaint_model import InpaintCAModel
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--flist', default='', type=str,
+    help='The filenames of image to be processed: input, mask, output.')
+parser.add_argument(
+    '--image_height', default=-1, type=int,
+    help='The height of images should be defined, otherwise batch mode is not'
+    ' supported.')
+parser.add_argument(
+    '--image_width', default=-1, type=int,
+    help='The width of images should be defined, otherwise batch mode is not'
+    ' supported.')
+parser.add_argument(
+    '--checkpoint_dir', default='', type=str,
+    help='The directory of tensorflow checkpoint.')
+
+
+if __name__ == "__main__":
+    FLAGS = ng.Config('inpaint.yml')
+    ng.get_gpus(1)
+    # os.environ['CUDA_VISIBLE_DEVICES'] =''
+    args = parser.parse_args()
+
+    sess_config = tf.ConfigProto()
+    sess_config.gpu_options.allow_growth = True
+    sess = tf.Session(config=sess_config)
+
+    model = InpaintCAModel()
+    input_image_ph = tf.placeholder(
+        tf.float32, shape=(1, args.image_height, args.image_width*2, 3))
+    output = model.build_server_graph(FLAGS, input_image_ph)
+    output = (output + 1.) * 127.5
+    output = tf.reverse(output, [-1])
+    output = tf.saturate_cast(output, tf.uint8)
+    vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
+    assign_ops = []
+    for var in vars_list:
+        vname = var.name
+        from_name = vname
+        var_value = tf.contrib.framework.load_variable(
+            args.checkpoint_dir, from_name)
+        assign_ops.append(tf.assign(var, var_value))
+    sess.run(assign_ops)
+    print('Model loaded.')
+
+    with open(args.flist, 'r') as f:
+        lines = f.read().splitlines()
+    t = time.time()
+    for line in lines:
+    # for i in range(100):
+        image, mask, out = line.split()
+        base = os.path.basename(mask)
+
+        image = cv2.imread(image)
+        mask = cv2.imread(mask)
+        image = cv2.resize(image, (args.image_width, args.image_height))
+        mask = cv2.resize(mask, (args.image_width, args.image_height))
+        # cv2.imwrite(out, image*(1-mask/255.) + mask)
+        # # continue
+        # image = np.zeros((128, 256, 3))
+        # mask = np.zeros((128, 256, 3))
+
+        assert image.shape == mask.shape
+
+        h, w, _ = image.shape
+        grid = 4
+        image = image[:h//grid*grid, :w//grid*grid, :]
+        mask = mask[:h//grid*grid, :w//grid*grid, :]
+        print('Shape of image: {}'.format(image.shape))
+
+        image = np.expand_dims(image, 0)
+        mask = np.expand_dims(mask, 0)
+        input_image = np.concatenate([image, mask], axis=2)
+
+        # load pretrained model
+        result = sess.run(output, feed_dict={input_image_ph: input_image})
+        print('Processed: {}'.format(out))
+        cv2.imwrite(out, result[0][:, :, ::-1])
+
+    print('Time total: {}'.format(time.time() - t))

guided_batch_test.py

@@ -0,0 +1,95 @@
+import time
+import os
+import argparse
+
+import cv2
+import numpy as np
+import tensorflow as tf
+import neuralgym as ng
+
+from inpaint_model import InpaintCAModel
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--flist', default='', type=str,
+    help='The filenames of image to be processed: input, mask, output.')
+parser.add_argument(
+    '--image_height', default=-1, type=int,
+    help='The height of images should be defined, otherwise batch mode is not'
+    ' supported.')
+parser.add_argument(
+    '--image_width', default=-1, type=int,
+    help='The width of images should be defined, otherwise batch mode is not'
+    ' supported.')
+parser.add_argument(
+    '--checkpoint_dir', default='', type=str,
+    help='The directory of tensorflow checkpoint.')
+
+
+if __name__ == "__main__":
+    ng.get_gpus(1)
+    # os.environ['CUDA_VISIBLE_DEVICES'] =''
+    args = parser.parse_args()
+
+    sess_config = tf.ConfigProto()
+    sess_config.gpu_options.allow_growth = True
+    sess = tf.Session(config=sess_config)
+
+    model = InpaintCAModel()
+    input_image_ph = tf.placeholder(
+        tf.float32, shape=(1, args.image_height, args.image_width*3, 3))
+    output = model.build_server_graph(input_image_ph)
+    output = (output + 1.) * 127.5
+    output = tf.reverse(output, [-1])
+    output = tf.saturate_cast(output, tf.uint8)
+    vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
+    assign_ops = []
+    for var in vars_list:
+        vname = var.name
+        from_name = vname
+        var_value = tf.contrib.framework.load_variable(
+            args.checkpoint_dir, from_name)
+        assign_ops.append(tf.assign(var, var_value))
+    sess.run(assign_ops)
+    print('Model loaded.')
+
+    with open(args.flist, 'r') as f:
+        lines = f.read().splitlines()
+    t = time.time()
+    for line in lines:
+    # for i in range(100):
+        image, mask, out = line.split()
+        base = os.path.basename(mask)
+
+        guidance = cv2.imread(image[:-4] + '_edge.jpg')
+        image = cv2.imread(image)
+        mask = cv2.imread(mask)
+        image = cv2.resize(image, (args.image_width, args.image_height))
+        guidance = cv2.resize(guidance, (args.image_width, args.image_height))
+        mask = cv2.resize(mask, (args.image_width, args.image_height))
+        # cv2.imwrite(out, image*(1-mask/255.) + mask)
+        # # continue
+        # image = np.zeros((128, 256, 3))
+        # mask = np.zeros((128, 256, 3))
+
+        assert image.shape == mask.shape
+
+        h, w, _ = image.shape
+        grid = 4
+        image = image[:h//grid*grid, :w//grid*grid, :]
+        mask = mask[:h//grid*grid, :w//grid*grid, :]
+        guidance = guidance[:h//grid*grid, :w//grid*grid, :]
+        print('Shape of image: {}'.format(image.shape))
+
+        image = np.expand_dims(image, 0)
+        guidance = np.expand_dims(guidance, 0)
+        mask = np.expand_dims(mask, 0)
+        input_image = np.concatenate([image, guidance, mask], axis=2)
+
+        # load pretrained model
+        result = sess.run(output, feed_dict={input_image_ph: input_image})
+        print('Processed: {}'.format(out))
+        cv2.imwrite(out, result[0][:, :, ::-1])
+
+    print('Time total: {}'.format(time.time() - t))

inpaint.yml

@@ -0,0 +1,70 @@
+# =========================== Basic Settings ===========================
+# machine info
+num_gpus_per_job: 1  # number of gpus each job need
+num_cpus_per_job: 4  # number of gpus each job need
+num_hosts_per_job: 1
+memory_per_job: 32  # number of gpus each job need
+gpu_type: 'nvidia-tesla-p100'
+
+# parameters
+name: places2_gated_conv_v100  # any name
+model_restore: ''  # logs/places2_gated_conv
+dataset: 'celebahq'  # 'tmnist', 'dtd', 'places2', 'celeba', 'imagenet', 'cityscapes'
+random_crop: False  # Set to false when dataset is 'celebahq', meaning only resize the images to img_shapes, instead of crop img_shapes from a larger raw image. This is useful when you train on images with different resolutions like places2. In these cases, please set random_crop to true.
+val: False  # true if you want to view validation results in tensorboard
+log_dir: logs/full_model_celeba_hq_256
+
+gan: 'sngan'
+gan_loss_alpha: 1
+gan_with_mask: True
+discounted_mask: True
+random_seed: False
+padding: 'SAME'
+
+# training
+train_spe: 4000
+max_iters: 100000000
+viz_max_out: 10
+val_psteps: 2000
+
+# data
+data_flist:
+  # https://github.com/jiahuiyu/progressive_growing_of_gans_tf
+  celebahq: [
+    'data/celeba_hq/train_shuffled.flist',
+    'data/celeba_hq/validation_static_view.flist'
+  ]
+  # http://mmlab.ie.cuhk.edu.hk/projects/celeba.html, please to use random_crop: True
+  celeba: [
+    'data/celeba/train_shuffled.flist',
+    'data/celeba/validation_static_view.flist'
+  ]
+  # http://places2.csail.mit.edu/, please download the high-resolution dataset and use random_crop: True
+  places2: [
+    'data/places2/train_shuffled.flist',
+    'data/places2/validation_static_view.flist'
+  ]
+  # http://www.image-net.org/, please use random_crop: True
+  imagenet: [
+    'data/imagenet/train_shuffled.flist',
+    'data/imagenet/validation_static_view.flist',
+  ]
+
+static_view_size: 30
+img_shapes: [256, 256, 3]
+height: 128
+width: 128
+max_delta_height: 32
+max_delta_width: 32
+batch_size: 16
+vertical_margin: 0
+horizontal_margin: 0
+
+# loss
+ae_loss: True
+l1_loss: True
+l1_loss_alpha: 1.
+
+# to tune
+guided: False
+edge_threshold: 0.6

inpaint_model.py

@@ -0,0 +1,297 @@
+""" common model for DCGAN """
+import logging
+
+import cv2
+import neuralgym as ng
+import tensorflow as tf
+from tensorflow.contrib.framework.python.ops import arg_scope
+
+from neuralgym.models import Model
+from neuralgym.ops.summary_ops import scalar_summary, images_summary
+from neuralgym.ops.summary_ops import gradients_summary
+from neuralgym.ops.layers import flatten, resize
+from neuralgym.ops.gan_ops import gan_hinge_loss
+from neuralgym.ops.gan_ops import random_interpolates
+
+from inpaint_ops import gen_conv, gen_deconv, dis_conv
+from inpaint_ops import random_bbox, bbox2mask, local_patch, brush_stroke_mask
+from inpaint_ops import resize_mask_like, contextual_attention
+
+
+logger = logging.getLogger()
+
+
+class InpaintCAModel(Model):
+    def __init__(self):
+        super().__init__('InpaintCAModel')
+
+    def build_inpaint_net(self, x, mask, reuse=False,
+                          training=True, padding='SAME', name='inpaint_net'):
+        """Inpaint network.
+
+        Args:
+            x: incomplete image, [-1, 1]
+            mask: mask region {0, 1}
+        Returns:
+            [-1, 1] as predicted image
+        """
+        xin = x
+        offset_flow = None
+        ones_x = tf.ones_like(x)[:, :, :, 0:1]
+        x = tf.concat([x, ones_x, ones_x*mask], axis=3)
+
+        # two stage network
+        cnum = 48
+        with tf.variable_scope(name, reuse=reuse), \
+                arg_scope([gen_conv, gen_deconv],
+                          training=training, padding=padding):
+            # stage1
+            x = gen_conv(x, cnum, 5, 1, name='conv1')
+            x = gen_conv(x, 2*cnum, 3, 2, name='conv2_downsample')
+            x = gen_conv(x, 2*cnum, 3, 1, name='conv3')
+            x = gen_conv(x, 4*cnum, 3, 2, name='conv4_downsample')
+            x = gen_conv(x, 4*cnum, 3, 1, name='conv5')
+            x = gen_conv(x, 4*cnum, 3, 1, name='conv6')
+            mask_s = resize_mask_like(mask, x)
+            x = gen_conv(x, 4*cnum, 3, rate=2, name='conv7_atrous')
+            x = gen_conv(x, 4*cnum, 3, rate=4, name='conv8_atrous')
+            x = gen_conv(x, 4*cnum, 3, rate=8, name='conv9_atrous')
+            x = gen_conv(x, 4*cnum, 3, rate=16, name='conv10_atrous')
+            x = gen_conv(x, 4*cnum, 3, 1, name='conv11')
+            x = gen_conv(x, 4*cnum, 3, 1, name='conv12')
+            x = gen_deconv(x, 2*cnum, name='conv13_upsample')
+            x = gen_conv(x, 2*cnum, 3, 1, name='conv14')
+            x = gen_deconv(x, cnum, name='conv15_upsample')
+            x = gen_conv(x, cnum//2, 3, 1, name='conv16')
+            x = gen_conv(x, 3, 3, 1, activation=None, name='conv17')
+            x = tf.nn.tanh(x)
+            x_stage1 = x
+
+            # stage2, paste result as input
+            x = x*mask + xin[:, :, :, 0:3]*(1.-mask)
+            x.set_shape(xin[:, :, :, 0:3].get_shape().as_list())
+            # conv branch
+            # xnow = tf.concat([x, ones_x, ones_x*mask], axis=3)
+            xnow = x
+            x = gen_conv(xnow, cnum, 5, 1, name='xconv1')
+            x = gen_conv(x, cnum, 3, 2, name='xconv2_downsample')
+            x = gen_conv(x, 2*cnum, 3, 1, name='xconv3')
+            x = gen_conv(x, 2*cnum, 3, 2, name='xconv4_downsample')
+            x = gen_conv(x, 4*cnum, 3, 1, name='xconv5')
+            x = gen_conv(x, 4*cnum, 3, 1, name='xconv6')
+            x = gen_conv(x, 4*cnum, 3, rate=2, name='xconv7_atrous')
+            x = gen_conv(x, 4*cnum, 3, rate=4, name='xconv8_atrous')
+            x = gen_conv(x, 4*cnum, 3, rate=8, name='xconv9_atrous')
+            x = gen_conv(x, 4*cnum, 3, rate=16, name='xconv10_atrous')
+            x_hallu = x
+            # attention branch
+            x = gen_conv(xnow, cnum, 5, 1, name='pmconv1')
+            x = gen_conv(x, cnum, 3, 2, name='pmconv2_downsample')
+            x = gen_conv(x, 2*cnum, 3, 1, name='pmconv3')
+            x = gen_conv(x, 4*cnum, 3, 2, name='pmconv4_downsample')
+            x = gen_conv(x, 4*cnum, 3, 1, name='pmconv5')
+            x = gen_conv(x, 4*cnum, 3, 1, name='pmconv6',
+                                activation=tf.nn.relu)
+            x, offset_flow = contextual_attention(x, x, mask_s, 3, 1, rate=2)
+            x = gen_conv(x, 4*cnum, 3, 1, name='pmconv9')
+            x = gen_conv(x, 4*cnum, 3, 1, name='pmconv10')
+            pm = x
+            x = tf.concat([x_hallu, pm], axis=3)
+
+            x = gen_conv(x, 4*cnum, 3, 1, name='allconv11')
+            x = gen_conv(x, 4*cnum, 3, 1, name='allconv12')
+            x = gen_deconv(x, 2*cnum, name='allconv13_upsample')
+            x = gen_conv(x, 2*cnum, 3, 1, name='allconv14')
+            x = gen_deconv(x, cnum, name='allconv15_upsample')
+            x = gen_conv(x, cnum//2, 3, 1, name='allconv16')
+            x = gen_conv(x, 3, 3, 1, activation=None, name='allconv17')
+            x = tf.nn.tanh(x)
+            x_stage2 = x
+        return x_stage1, x_stage2, offset_flow
+
+    def build_sn_patch_gan_discriminator(self, x, reuse=False, training=True):
+        with tf.variable_scope('sn_patch_gan', reuse=reuse):
+            cnum = 64
+            x = dis_conv(x, cnum, name='conv1', training=training)
+            x = dis_conv(x, cnum*2, name='conv2', training=training)
+            x = dis_conv(x, cnum*4, name='conv3', training=training)
+            x = dis_conv(x, cnum*4, name='conv4', training=training)
+            x = dis_conv(x, cnum*4, name='conv5', training=training)
+            x = dis_conv(x, cnum*4, name='conv6', training=training)
+            x = flatten(x, name='flatten')
+            return x
+
+    def build_gan_discriminator(
+            self, batch, reuse=False, training=True):
+        with tf.variable_scope('discriminator', reuse=reuse):
+            d = self.build_sn_patch_gan_discriminator(
+                batch, reuse=reuse, training=training)
+            return d
+
+    def build_graph_with_losses(
+            self, FLAGS, batch_data, training=True, summary=False,
+            reuse=False):
+        if FLAGS.guided:
+            batch_data, edge = batch_data
+            edge = edge[:, :, :, 0:1] / 255.
+            edge = tf.cast(edge > FLAGS.edge_threshold, tf.float32)
+        batch_pos = batch_data / 127.5 - 1.
+        # generate mask, 1 represents masked point
+        bbox = random_bbox(FLAGS)
+        regular_mask = bbox2mask(FLAGS, bbox, name='mask_c')
+        irregular_mask = brush_stroke_mask(FLAGS, name='mask_c')
+        mask = tf.cast(
+            tf.logical_or(
+                tf.cast(irregular_mask, tf.bool),
+                tf.cast(regular_mask, tf.bool),
+            ),
+            tf.float32
+        )
+
+        batch_incomplete = batch_pos*(1.-mask)
+        if FLAGS.guided:
+            edge = edge * mask
+            xin = tf.concat([batch_incomplete, edge], axis=3)
+        else:
+            xin = batch_incomplete
+        x1, x2, offset_flow = self.build_inpaint_net(
+            xin, mask, reuse=reuse, training=training,
+            padding=FLAGS.padding)
+        batch_predicted = x2
+        losses = {}
+        # apply mask and complete image
+        batch_complete = batch_predicted*mask + batch_incomplete*(1.-mask)
+        # local patches
+        losses['ae_loss'] = FLAGS.l1_loss_alpha * tf.reduce_mean(tf.abs(batch_pos - x1))
+        losses['ae_loss'] += FLAGS.l1_loss_alpha * tf.reduce_mean(tf.abs(batch_pos - x2))
+        if summary:
+            scalar_summary('losses/ae_loss', losses['ae_loss'])
+            if FLAGS.guided:
+                viz_img = [
+                    batch_pos,
+                    batch_incomplete + edge,
+                    batch_complete]
+            else:
+                viz_img = [batch_pos, batch_incomplete, batch_complete]
+            if offset_flow is not None:
+                viz_img.append(
+                    resize(offset_flow, scale=4,
+                           func=tf.image.resize_bilinear))
+            images_summary(
+                tf.concat(viz_img, axis=2),
+                'raw_incomplete_predicted_complete', FLAGS.viz_max_out)
+
+        # gan
+        batch_pos_neg = tf.concat([batch_pos, batch_complete], axis=0)
+        if FLAGS.gan_with_mask:
+            batch_pos_neg = tf.concat([batch_pos_neg, tf.tile(mask, [FLAGS.batch_size*2, 1, 1, 1])], axis=3)
+        if FLAGS.guided:
+            # conditional GANs
+            batch_pos_neg = tf.concat([batch_pos_neg, tf.tile(edge, [2, 1, 1, 1])], axis=3)
+        # wgan with gradient penalty
+        if FLAGS.gan == 'sngan':
+            pos_neg = self.build_gan_discriminator(batch_pos_neg, training=training, reuse=reuse)
+            pos, neg = tf.split(pos_neg, 2)
+            g_loss, d_loss = gan_hinge_loss(pos, neg)
+            losses['g_loss'] = g_loss
+            losses['d_loss'] = d_loss
+        else:
+            raise NotImplementedError('{} not implemented.'.format(FLAGS.gan))
+        if summary:
+            # summary the magnitude of gradients from different losses w.r.t. predicted image
+            gradients_summary(losses['g_loss'], batch_predicted, name='g_loss')
+            gradients_summary(losses['g_loss'], x2, name='g_loss_to_x2')
+            # gradients_summary(losses['ae_loss'], x1, name='ae_loss_to_x1')
+            gradients_summary(losses['ae_loss'], x2, name='ae_loss_to_x2')
+        losses['g_loss'] = FLAGS.gan_loss_alpha * losses['g_loss']
+        if FLAGS.ae_loss:
+            losses['g_loss'] += losses['ae_loss']
+        g_vars = tf.get_collection(
+            tf.GraphKeys.TRAINABLE_VARIABLES, 'inpaint_net')
+        d_vars = tf.get_collection(
+            tf.GraphKeys.TRAINABLE_VARIABLES, 'discriminator')
+        return g_vars, d_vars, losses
+
+    def build_infer_graph(self, FLAGS, batch_data, bbox=None, name='val'):
+        """
+        """
+        if FLAGS.guided:
+            batch_data, edge = batch_data
+            edge = edge[:, :, :, 0:1] / 255.
+            edge = tf.cast(edge > FLAGS.edge_threshold, tf.float32)
+        regular_mask = bbox2mask(FLAGS, bbox, name='mask_c')
+        irregular_mask = brush_stroke_mask(FLAGS, name='mask_c')
+        mask = tf.cast(
+            tf.logical_or(
+                tf.cast(irregular_mask, tf.bool),
+                tf.cast(regular_mask, tf.bool),
+            ),
+            tf.float32
+        )
+
+        batch_pos = batch_data / 127.5 - 1.
+        batch_incomplete = batch_pos*(1.-mask)
+        if FLAGS.guided:
+            edge = edge * mask
+            xin = tf.concat([batch_incomplete, edge], axis=3)
+        else:
+            xin = batch_incomplete
+        # inpaint
+        x1, x2, offset_flow = self.build_inpaint_net(
+            xin, mask, reuse=True,
+            training=False, padding=FLAGS.padding)
+        batch_predicted = x2
+        # apply mask and reconstruct
+        batch_complete = batch_predicted*mask + batch_incomplete*(1.-mask)
+        # global image visualization
+        if FLAGS.guided:
+            viz_img = [
+                batch_pos,
+                batch_incomplete + edge,
+                batch_complete]
+        else:
+            viz_img = [batch_pos, batch_incomplete, batch_complete]
+        if offset_flow is not None:
+            viz_img.append(
+                resize(offset_flow, scale=4,
+                       func=tf.image.resize_bilinear))
+        images_summary(
+            tf.concat(viz_img, axis=2),
+            name+'_raw_incomplete_complete', FLAGS.viz_max_out)
+        return batch_complete
+
+    def build_static_infer_graph(self, FLAGS, batch_data, name):
+        """
+        """
+        # generate mask, 1 represents masked point
+        bbox = (tf.constant(FLAGS.height//2), tf.constant(FLAGS.width//2),
+                tf.constant(FLAGS.height), tf.constant(FLAGS.width))
+        return self.build_infer_graph(FLAGS, batch_data, bbox, name)
+
+
+    def build_server_graph(self, FLAGS, batch_data, reuse=False, is_training=False):
+        """
+        """
+        # generate mask, 1 represents masked point
+        if FLAGS.guided:
+            batch_raw, edge, masks_raw = tf.split(batch_data, 3, axis=2)
+            edge = edge[:, :, :, 0:1] / 255.
+            edge = tf.cast(edge > FLAGS.edge_threshold, tf.float32)
+        else:
+            batch_raw, masks_raw = tf.split(batch_data, 2, axis=2)
+        masks = tf.cast(masks_raw[0:1, :, :, 0:1] > 127.5, tf.float32)
+
+        batch_pos = batch_raw / 127.5 - 1.
+        batch_incomplete = batch_pos * (1. - masks)
+        if FLAGS.guided:
+            edge = edge * masks[:, :, :, 0:1]
+            xin = tf.concat([batch_incomplete, edge], axis=3)
+        else:
+            xin = batch_incomplete
+        # inpaint
+        x1, x2, flow = self.build_inpaint_net(
+            xin, masks, reuse=reuse, training=is_training)
+        batch_predict = x2
+        # apply mask and reconstruct
+        batch_complete = batch_predict*masks + batch_incomplete*(1-masks)
+        return batch_complete

inpaint_ops.py

@@ -0,0 +1,553 @@
+import logging
+import math
+
+import cv2
+import numpy as np
+import tensorflow as tf
+from tensorflow.contrib.framework.python.ops import add_arg_scope
+from PIL import Image, ImageDraw
+
+from neuralgym.ops.layers import resize
+from neuralgym.ops.layers import *
+from neuralgym.ops.loss_ops import *
+from neuralgym.ops.gan_ops import *
+from neuralgym.ops.summary_ops import *
+
+
+logger = logging.getLogger()
+np.random.seed(2018)
+
+
+@add_arg_scope
+def gen_conv(x, cnum, ksize, stride=1, rate=1, name='conv',
+             padding='SAME', activation=tf.nn.elu, training=True):
+    """Define conv for generator.
+
+    Args:
+        x: Input.
+        cnum: Channel number.
+        ksize: Kernel size.
+        Stride: Convolution stride.
+        Rate: Rate for or dilated conv.
+        name: Name of layers.
+        padding: Default to SYMMETRIC.
+        activation: Activation function after convolution.
+        training: If current graph is for training or inference, used for bn.
+
+    Returns:
+        tf.Tensor: output
+
+    """
+    assert padding in ['SYMMETRIC', 'SAME', 'REFELECT']
+    if padding == 'SYMMETRIC' or padding == 'REFELECT':
+        p = int(rate*(ksize-1)/2)
+        x = tf.pad(x, [[0,0], [p, p], [p, p], [0,0]], mode=padding)
+        padding = 'VALID'
+    x = tf.layers.conv2d(
+        x, cnum, ksize, stride, dilation_rate=rate,
+        activation=None, padding=padding, name=name)
+    if cnum == 3 or activation is None:
+        # conv for output
+        return x
+    x, y = tf.split(x, 2, 3)
+    x = activation(x)
+    y = tf.nn.sigmoid(y)
+    x = x * y
+    return x
+
+
+@add_arg_scope
+def gen_deconv(x, cnum, name='upsample', padding='SAME', training=True):
+    """Define deconv for generator.
+    The deconv is defined to be a x2 resize_nearest_neighbor operation with
+    additional gen_conv operation.
+
+    Args:
+        x: Input.
+        cnum: Channel number.
+        name: Name of layers.
+        training: If current graph is for training or inference, used for bn.
+
+    Returns:
+        tf.Tensor: output
+
+    """
+    with tf.variable_scope(name):
+        x = resize(x, func=tf.image.resize_nearest_neighbor)
+        x = gen_conv(
+            x, cnum, 3, 1, name=name+'_conv', padding=padding,
+            training=training)
+    return x
+
+
+@add_arg_scope
+def dis_conv(x, cnum, ksize=5, stride=2, name='conv', training=True):
+    """Define conv for discriminator.
+    Activation is set to leaky_relu.
+
+    Args:
+        x: Input.
+        cnum: Channel number.
+        ksize: Kernel size.
+        Stride: Convolution stride.
+        name: Name of layers.
+        training: If current graph is for training or inference, used for bn.
+
+    Returns:
+        tf.Tensor: output
+
+    """
+    x = conv2d_spectral_norm(x, cnum, ksize, stride, 'SAME', name=name)
+    x = tf.nn.leaky_relu(x)
+    return x
+
+
+def random_bbox(FLAGS):
+    """Generate a random tlhw.
+
+    Returns:
+        tuple: (top, left, height, width)
+
+    """
+    img_shape = FLAGS.img_shapes
+    img_height = img_shape[0]
+    img_width = img_shape[1]
+    maxt = img_height - FLAGS.vertical_margin - FLAGS.height
+    maxl = img_width - FLAGS.horizontal_margin - FLAGS.width
+    t = tf.random_uniform(
+        [], minval=FLAGS.vertical_margin, maxval=maxt, dtype=tf.int32)
+    l = tf.random_uniform(
+        [], minval=FLAGS.horizontal_margin, maxval=maxl, dtype=tf.int32)
+    h = tf.constant(FLAGS.height)
+    w = tf.constant(FLAGS.width)
+    return (t, l, h, w)
+
+
+def bbox2mask(FLAGS, bbox, name='mask'):
+    """Generate mask tensor from bbox.
+
+    Args:
+        bbox: tuple, (top, left, height, width)
+
+    Returns:
+        tf.Tensor: output with shape [1, H, W, 1]
+
+    """
+    def npmask(bbox, height, width, delta_h, delta_w):
+        mask = np.zeros((1, height, width, 1), np.float32)
+        h = np.random.randint(delta_h//2+1)
+        w = np.random.randint(delta_w//2+1)
+        mask[:, bbox[0]+h:bbox[0]+bbox[2]-h,
+             bbox[1]+w:bbox[1]+bbox[3]-w, :] = 1.
+        return mask
+    with tf.variable_scope(name), tf.device('/cpu:0'):
+        img_shape = FLAGS.img_shapes
+        height = img_shape[0]
+        width = img_shape[1]
+        mask = tf.py_func(
+            npmask,
+            [bbox, height, width,
+             FLAGS.max_delta_height, FLAGS.max_delta_width],
+            tf.float32, stateful=False)
+        mask.set_shape([1] + [height, width] + [1])
+    return mask
+
+
+def brush_stroke_mask(FLAGS, name='mask'):
+    """Generate mask tensor from bbox.
+
+    Returns:
+        tf.Tensor: output with shape [1, H, W, 1]
+
+    """
+    min_num_vertex = 4
+    max_num_vertex = 12
+    mean_angle = 2*math.pi / 5
+    angle_range = 2*math.pi / 15
+    min_width = 12
+    max_width = 40
+    def generate_mask(H, W):
+        average_radius = math.sqrt(H*H+W*W) / 8
+        mask = Image.new('L', (W, H), 0)
+
+        for _ in range(np.random.randint(1, 4)):
+            num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
+            angle_min = mean_angle - np.random.uniform(0, angle_range)
+            angle_max = mean_angle + np.random.uniform(0, angle_range)
+            angles = []
+            vertex = []
+            for i in range(num_vertex):
+                if i % 2 == 0:
+                    angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
+                else:
+                    angles.append(np.random.uniform(angle_min, angle_max))
+
+            h, w = mask.size
+            vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
+            for i in range(num_vertex):
+                r = np.clip(
+                    np.random.normal(loc=average_radius, scale=average_radius//2),
+                    0, 2*average_radius)
+                new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
+                new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
+                vertex.append((int(new_x), int(new_y)))
+
+            draw = ImageDraw.Draw(mask)
+            width = int(np.random.uniform(min_width, max_width))
+            draw.line(vertex, fill=1, width=width)
+            for v in vertex:
+                draw.ellipse((v[0] - width//2,
+                              v[1] - width//2,
+                              v[0] + width//2,
+                              v[1] + width//2),
+                             fill=1)
+
+        if np.random.normal() > 0:
+            mask.transpose(Image.FLIP_LEFT_RIGHT)
+        if np.random.normal() > 0:
+            mask.transpose(Image.FLIP_TOP_BOTTOM)
+        mask = np.asarray(mask, np.float32)
+        mask = np.reshape(mask, (1, H, W, 1))
+        return mask
+    with tf.variable_scope(name), tf.device('/cpu:0'):
+        img_shape = FLAGS.img_shapes
+        height = img_shape[0]
+        width = img_shape[1]
+        mask = tf.py_func(
+            generate_mask,
+            [height, width],
+            tf.float32, stateful=True)
+        mask.set_shape([1] + [height, width] + [1])
+    return mask
+
+
+def local_patch(x, bbox):
+    """Crop local patch according to bbox.
+
+    Args:
+        x: input
+        bbox: (top, left, height, width)
+
+    Returns:
+        tf.Tensor: local patch
+
+    """
+    x = tf.image.crop_to_bounding_box(x, bbox[0], bbox[1], bbox[2], bbox[3])
+    return x
+
+
+def resize_mask_like(mask, x):
+    """Resize mask like shape of x.
+
+    Args:
+        mask: Original mask.
+        x: To shape of x.
+
+    Returns:
+        tf.Tensor: resized mask
+
+    """
+    mask_resize = resize(
+        mask, to_shape=x.get_shape().as_list()[1:3],
+        func=tf.image.resize_nearest_neighbor)
+    return mask_resize
+
+
+def contextual_attention(f, b, mask=None, ksize=3, stride=1, rate=1,
+                         fuse_k=3, softmax_scale=10., training=True, fuse=True):
+    """ Contextual attention layer implementation.
+
+    Contextual attention is first introduced in publication:
+        Generative Image Inpainting with Contextual Attention, Yu et al.
+
+    Args:
+        x: Input feature to match (foreground).
+        t: Input feature for match (background).
+        mask: Input mask for t, indicating patches not available.
+        ksize: Kernel size for contextual attention.
+        stride: Stride for extracting patches from t.
+        rate: Dilation for matching.
+        softmax_scale: Scaled softmax for attention.
+        training: Indicating if current graph is training or inference.
+
+    Returns:
+        tf.Tensor: output
+
+    """
+    # get shapes
+    raw_fs = tf.shape(f)
+    raw_int_fs = f.get_shape().as_list()
+    raw_int_bs = b.get_shape().as_list()
+    # extract patches from background with stride and rate
+    kernel = 2*rate
+    raw_w = tf.extract_image_patches(
+        b, [1,kernel,kernel,1], [1,rate*stride,rate*stride,1], [1,1,1,1], padding='SAME')
+    raw_w = tf.reshape(raw_w, [raw_int_bs[0], -1, kernel, kernel, raw_int_bs[3]])
+    raw_w = tf.transpose(raw_w, [0, 2, 3, 4, 1])  # transpose to b*k*k*c*hw
+    # downscaling foreground option: downscaling both foreground and
+    # background for matching and use original background for reconstruction.
+    f = resize(f, scale=1./rate, func=tf.image.resize_nearest_neighbor)
+    b = resize(b, to_shape=[int(raw_int_bs[1]/rate), int(raw_int_bs[2]/rate)], func=tf.image.resize_nearest_neighbor)  # https://github.com/tensorflow/tensorflow/issues/11651
+    if mask is not None:
+        mask = resize(mask, scale=1./rate, func=tf.image.resize_nearest_neighbor)
+    fs = tf.shape(f)
+    int_fs = f.get_shape().as_list()
+    f_groups = tf.split(f, int_fs[0], axis=0)
+    # from t(H*W*C) to w(b*k*k*c*h*w)
+    bs = tf.shape(b)
+    int_bs = b.get_shape().as_list()
+    w = tf.extract_image_patches(
+        b, [1,ksize,ksize,1], [1,stride,stride,1], [1,1,1,1], padding='SAME')
+    w = tf.reshape(w, [int_fs[0], -1, ksize, ksize, int_fs[3]])
+    w = tf.transpose(w, [0, 2, 3, 4, 1])  # transpose to b*k*k*c*hw
+    # process mask
+    if mask is None:
+        mask = tf.zeros([1, bs[1], bs[2], 1])
+    m = tf.extract_image_patches(
+        mask, [1,ksize,ksize,1], [1,stride,stride,1], [1,1,1,1], padding='SAME')
+    m = tf.reshape(m, [1, -1, ksize, ksize, 1])
+    m = tf.transpose(m, [0, 2, 3, 4, 1])  # transpose to b*k*k*c*hw
+    m = m[0]
+    mm = tf.cast(tf.equal(tf.reduce_mean(m, axis=[0,1,2], keep_dims=True), 0.), tf.float32)
+    w_groups = tf.split(w, int_bs[0], axis=0)
+    raw_w_groups = tf.split(raw_w, int_bs[0], axis=0)
+    y = []
+    offsets = []
+    k = fuse_k
+    scale = softmax_scale
+    fuse_weight = tf.reshape(tf.eye(k), [k, k, 1, 1])
+    for xi, wi, raw_wi in zip(f_groups, w_groups, raw_w_groups):
+        # conv for compare
+        wi = wi[0]
+        wi_normed = wi / tf.maximum(tf.sqrt(tf.reduce_sum(tf.square(wi), axis=[0,1,2])), 1e-4)
+        yi = tf.nn.conv2d(xi, wi_normed, strides=[1,1,1,1], padding="SAME")
+
+        # conv implementation for fuse scores to encourage large patches
+        if fuse:
+            yi = tf.reshape(yi, [1, fs[1]*fs[2], bs[1]*bs[2], 1])
+            yi = tf.nn.conv2d(yi, fuse_weight, strides=[1,1,1,1], padding='SAME')
+            yi = tf.reshape(yi, [1, fs[1], fs[2], bs[1], bs[2]])
+            yi = tf.transpose(yi, [0, 2, 1, 4, 3])
+            yi = tf.reshape(yi, [1, fs[1]*fs[2], bs[1]*bs[2], 1])
+            yi = tf.nn.conv2d(yi, fuse_weight, strides=[1,1,1,1], padding='SAME')
+            yi = tf.reshape(yi, [1, fs[2], fs[1], bs[2], bs[1]])
+            yi = tf.transpose(yi, [0, 2, 1, 4, 3])
+        yi = tf.reshape(yi, [1, fs[1], fs[2], bs[1]*bs[2]])
+
+        # softmax to match
+        yi *=  mm  # mask
+        yi = tf.nn.softmax(yi*scale, 3)
+        yi *=  mm  # mask
+
+        offset = tf.argmax(yi, axis=3, output_type=tf.int32)
+        offset = tf.stack([offset // fs[2], offset % fs[2]], axis=-1)
+        # deconv for patch pasting
+        # 3.1 paste center
+        wi_center = raw_wi[0]
+        yi = tf.nn.conv2d_transpose(yi, wi_center, tf.concat([[1], raw_fs[1:]], axis=0), strides=[1,rate,rate,1]) / 4.
+        y.append(yi)
+        offsets.append(offset)
+    y = tf.concat(y, axis=0)
+    y.set_shape(raw_int_fs)
+    offsets = tf.concat(offsets, axis=0)
+    offsets.set_shape(int_bs[:3] + [2])
+    # case1: visualize optical flow: minus current position
+    h_add = tf.tile(tf.reshape(tf.range(bs[1]), [1, bs[1], 1, 1]), [bs[0], 1, bs[2], 1])
+    w_add = tf.tile(tf.reshape(tf.range(bs[2]), [1, 1, bs[2], 1]), [bs[0], bs[1], 1, 1])
+    offsets = offsets - tf.concat([h_add, w_add], axis=3)
+    # to flow image
+    flow = flow_to_image_tf(offsets)
+    # # case2: visualize which pixels are attended
+    # flow = highlight_flow_tf(offsets * tf.cast(mask, tf.int32))
+    if rate != 1:
+        flow = resize(flow, scale=rate, func=tf.image.resize_bilinear)
+    return y, flow
+
+
+def test_contextual_attention(args):
+    """Test contextual attention layer with 3-channel image input
+    (instead of n-channel feature).
+
+    """
+    import cv2
+    import os
+    # run on cpu
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+    rate = 2
+    stride = 1
+    grid = rate*stride
+
+    b = cv2.imread(args.imageA)
+    b = cv2.resize(b, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC)
+    h, w, _ = b.shape
+    b = b[:h//grid*grid, :w//grid*grid, :]
+    b = np.expand_dims(b, 0)
+    logger.info('Size of imageA: {}'.format(b.shape))
+
+    f = cv2.imread(args.imageB)
+    h, w, _ = f.shape
+    f = f[:h//grid*grid, :w//grid*grid, :]
+    f = np.expand_dims(f, 0)
+    logger.info('Size of imageB: {}'.format(f.shape))
+
+    with tf.Session() as sess:
+        bt = tf.constant(b, dtype=tf.float32)
+        ft = tf.constant(f, dtype=tf.float32)
+
+        yt, flow = contextual_attention(
+            ft, bt, stride=stride, rate=rate,
+            training=False, fuse=False)
+        y = sess.run(yt)
+        cv2.imwrite(args.imageOut, y[0])
+
+
+def make_color_wheel():
+    RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6)
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros([ncols, 3])
+    col = 0
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
+    col += RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
+    colorwheel[col:col+YG, 1] = 255
+    col += YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
+    col += GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
+    colorwheel[col:col+CB, 2] = 255
+    col += CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
+    col += + BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+
+COLORWHEEL = make_color_wheel()
+
+
+def compute_color(u,v):
+    h, w = u.shape
+    img = np.zeros([h, w, 3])
+    nanIdx = np.isnan(u) | np.isnan(v)
+    u[nanIdx] = 0
+    v[nanIdx] = 0
+    # colorwheel = COLORWHEEL
+    colorwheel = make_color_wheel()
+    ncols = np.size(colorwheel, 0)
+    rad = np.sqrt(u**2+v**2)
+    a = np.arctan2(-v, -u) / np.pi
+    fk = (a+1) / 2 * (ncols - 1) + 1
+    k0 = np.floor(fk).astype(int)
+    k1 = k0 + 1
+    k1[k1 == ncols+1] = 1
+    f = fk - k0
+    for i in range(np.size(colorwheel,1)):
+        tmp = colorwheel[:, i]
+        col0 = tmp[k0-1] / 255
+        col1 = tmp[k1-1] / 255
+        col = (1-f) * col0 + f * col1
+        idx = rad <= 1
+        col[idx] = 1-rad[idx]*(1-col[idx])
+        notidx = np.logical_not(idx)
+        col[notidx] *= 0.75
+        img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))
+    return img
+
+
+
+def flow_to_image(flow):
+    """Transfer flow map to image.
+    Part of code forked from flownet.
+    """
+    out = []
+    maxu = -999.
+    maxv = -999.
+    minu = 999.
+    minv = 999.
+    maxrad = -1
+    for i in range(flow.shape[0]):
+        u = flow[i, :, :, 0]
+        v = flow[i, :, :, 1]
+        idxunknow = (abs(u) > 1e7) | (abs(v) > 1e7)
+        u[idxunknow] = 0
+        v[idxunknow] = 0
+        maxu = max(maxu, np.max(u))
+        minu = min(minu, np.min(u))
+        maxv = max(maxv, np.max(v))
+        minv = min(minv, np.min(v))
+        rad = np.sqrt(u ** 2 + v ** 2)
+        maxrad = max(maxrad, np.max(rad))
+        u = u/(maxrad + np.finfo(float).eps)
+        v = v/(maxrad + np.finfo(float).eps)
+        img = compute_color(u, v)
+        out.append(img)
+    return np.float32(np.uint8(out))
+
+
+def flow_to_image_tf(flow, name='flow_to_image'):
+    """Tensorflow ops for computing flow to image.
+    """
+    with tf.variable_scope(name), tf.device('/cpu:0'):
+        img = tf.py_func(flow_to_image, [flow], tf.float32, stateful=False)
+        img.set_shape(flow.get_shape().as_list()[0:-1]+[3])
+        img = img / 127.5 - 1.
+        return img
+
+
+def highlight_flow(flow):
+    """Convert flow into middlebury color code image.
+    """
+    out = []
+    s = flow.shape
+    for i in range(flow.shape[0]):
+        img = np.ones((s[1], s[2], 3)) * 144.
+        u = flow[i, :, :, 0]
+        v = flow[i, :, :, 1]
+        for h in range(s[1]):
+            for w in range(s[1]):
+                ui = u[h,w]
+                vi = v[h,w]
+                img[ui, vi, :] = 255.
+        out.append(img)
+    return np.float32(np.uint8(out))
+
+
+def highlight_flow_tf(flow, name='flow_to_image'):
+    """Tensorflow ops for highlight flow.
+    """
+    with tf.variable_scope(name), tf.device('/cpu:0'):
+        img = tf.py_func(highlight_flow, [flow], tf.float32, stateful=False)
+        img.set_shape(flow.get_shape().as_list()[0:-1]+[3])
+        img = img / 127.5 - 1.
+        return img
+
+
+def image2edge(image):
+    """Convert image to edges.
+    """
+    out = []
+    for i in range(image.shape[0]):
+        img = cv2.Laplacian(image[i, :, :, :], cv2.CV_64F, ksize=3, scale=2)
+        out.append(img)
+    return np.float32(np.uint8(out))
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--imageA', default='', type=str, help='Image A as background patches to reconstruct image B.')
+    parser.add_argument('--imageB', default='', type=str, help='Image B is reconstructed with image A.')
+    parser.add_argument('--imageOut', default='result.png', type=str, help='Image B is reconstructed with image A.')
+    args = parser.parse_args()
+    test_contextual_attention(args)

main.py

@@ -0,0 +1,58 @@
+import argparse
+
+from PIL import Image
+import cv2
+import numpy as np
+from preprocess_image import preprocess_image
+import tensorflow as tf
+import neuralgym as ng
+
+from inpaint_model import InpaintCAModel
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--image', default='', type=str,
+                    help='The filename of image to be completed.')
+parser.add_argument('--output', default='output.png', type=str,
+                    help='Where to write output.')
+parser.add_argument('--watermark_type', default='istock', type=str,
+                    help='The watermark type')
+parser.add_argument('--checkpoint_dir', default='model/', type=str,
+                    help='The directory of tensorflow checkpoint.')
+
+#checkpoint_dir = 'model/'
+
+
+if __name__ == "__main__":
+    FLAGS = ng.Config('inpaint.yml')
+    # ng.get_gpus(1)
+    args, unknown = parser.parse_known_args()
+
+    model = InpaintCAModel()
+    image = Image.open(args.image)
+    input_image = preprocess_image(image, args.watermark_type)
+    tf.reset_default_graph()
+
+    sess_config = tf.ConfigProto()
+    sess_config.gpu_options.allow_growth = True
+    if (input_image.shape != (0,)):
+        with tf.Session(config=sess_config) as sess:
+            input_image = tf.constant(input_image, dtype=tf.float32)
+            output = model.build_server_graph(FLAGS, input_image)
+            output = (output + 1.) * 127.5
+            output = tf.reverse(output, [-1])
+            output = tf.saturate_cast(output, tf.uint8)
+            # load pretrained model
+            vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
+            assign_ops = []
+            for var in vars_list:
+                vname = var.name
+                from_name = vname
+                var_value = tf.contrib.framework.load_variable(
+                    args.checkpoint_dir, from_name)
+                assign_ops.append(tf.assign(var, var_value))
+            sess.run(assign_ops)
+            print('Model loaded.')
+            result = sess.run(output)
+            cv2.imwrite(args.output, cv2.cvtColor(
+                result[0][:, :, ::-1], cv2.COLOR_BGR2RGB))
+            print('image saved to {}'.format(args.output))

preprocess_image.py

@@ -0,0 +1,53 @@
+import numpy as np
+from PIL import Image
+import cv2
+
+
+def preprocess_image(image, watermark_type):
+    image_type: str = ''
+    preprocessed_mask_image = np.array([])
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    image = np.array(image)
+    image_h = image.shape[0]
+    image_w = image.shape[1]
+    aspectRatioImage = image_w / image_h
+    print("image size: {}".format(image.shape))
+
+    if image_w > image_h:
+        image_type = "landscape"
+    elif image_w == image_h:
+        image_type = "landscape"
+    else:
+        image_type = "potrait"
+
+    mask_image = Image.open(
+        "utils/{}/{}/mask.png".format(watermark_type, image_type))
+    if mask_image.mode != "RGB":
+        mask_image = mask_image.convert("RGB")
+    mask_image = np.array(mask_image)
+    print("mask image size: {}".format(mask_image.shape))
+
+    aspectRatioMaskImage = mask_image.shape[1] / mask_image.shape[0]
+    upperBoundAspectRatio = 1.05 * aspectRatioMaskImage
+    lowerBoundAspectRatio = 0.95 * aspectRatioMaskImage
+
+    if aspectRatioImage >= lowerBoundAspectRatio and aspectRatioImage <= upperBoundAspectRatio:
+        preprocessed_mask_image = cv2.resize(mask_image, (image_w, image_h))
+        print(preprocessed_mask_image.shape)
+    else:
+        print("Image size not supported!!!")
+
+    if (preprocessed_mask_image.shape != (0,)):
+        assert image.shape == preprocessed_mask_image.shape
+        grid = 8
+        image = image[:image_h//grid*grid, :image_w//grid*grid, :]
+        preprocessed_mask_image = preprocessed_mask_image[:image_h //
+                                                          grid*grid, :image_w//grid*grid, :]
+        image = np.expand_dims(image, 0)
+        preprocessed_mask_image = np.expand_dims(preprocessed_mask_image, 0)
+        input_image = np.concatenate([image, preprocessed_mask_image], axis=2)
+        return input_image
+
+    else:
+        return preprocessed_mask_image

requirements.txt

@@ -0,0 +1,2 @@
+tensorflow==1.15.5
+opencv-python==4.9.0.80