├── models └── .gitkeep ├── decensor_input_images └── .gitkeep ├── decensor_output_images └── .gitkeep ├── testing_output_images └── .gitkeep ├── training_output_images └── .gitkeep ├── training_data ├── .gitignore └── to_npy.py ├── readme_images ├── guide.png ├── collage.png └── format_results.py ├── load.py ├── LICENSE ├── decensor.py ├── test.py ├── layer.py ├── poisson_blend.py ├── README.md ├── train.py ├── model.py ├── model_mosaic.py └── train_mosaic.py /models/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /decensor_input_images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /decensor_output_images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /testing_output_images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /training_output_images/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /training_data/.gitignore: -------------------------------------------------------------------------------- 1 | images/* 2 | npy/* 3 | !.gitkeep 4 | -------------------------------------------------------------------------------- /readme_images/guide.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/StartleStars/DeepMindBreak/HEAD/readme_images/guide.png -------------------------------------------------------------------------------- /readme_images/collage.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/StartleStars/DeepMindBreak/HEAD/readme_images/collage.png -------------------------------------------------------------------------------- /load.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | 4 | def load(dir_='./training_data/npy'): 5 | x_train = np.load(os.path.join(dir_, 'x_train.npy')) 6 | x_test = np.load(os.path.join(dir_, 'x_test.npy')) 7 | return x_train, x_test 8 | 9 | 10 | if __name__ == '__main__': 11 | x_train, x_test = load() 12 | print(x_train.shape) 13 | print(x_test.shape) 14 | 15 | -------------------------------------------------------------------------------- /readme_images/format_results.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import matplotlib.pyplot as plt 3 | 4 | def format_results(images, dst): 5 | fig = plt.figure() 6 | for i, image in enumerate(images): 7 | text, img = image 8 | fig.add_subplot(1, 3, i + 1) 9 | plt.imshow(img) 10 | plt.tick_params(labelbottom='off') 11 | plt.tick_params(labelleft='off') 12 | plt.gca().get_xaxis().set_ticks_position('none') 13 | plt.gca().get_yaxis().set_ticks_position('none') 14 | plt.xlabel(text) 15 | plt.savefig(dst) 16 | plt.close() 17 | 18 | if __name__ == "__main__": 19 | masked = Image.open("censored.png") 20 | img = Image.open("decensored.png") 21 | raw = Image.open("original.png") 22 | format_results([['Input', masked], ['Output', img], ['Ground Truth', raw]], "result.png") -------------------------------------------------------------------------------- /training_data/to_npy.py: -------------------------------------------------------------------------------- 1 | import glob 2 | import os 3 | #import cv2 4 | from PIL import Image 5 | import numpy as np 6 | 7 | ratio = 0.95 8 | image_size = 128 9 | 10 | x = [] 11 | paths = glob.glob('images/*') 12 | for path in paths: 13 | #img = cv2.imread(path) 14 | #img = Image.open(path) 15 | #img = cv2.resize(img, (image_size, image_size)) 16 | #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 17 | #x.append(img) 18 | temp = Image.open(path) 19 | keep = temp.copy() 20 | keep = np.array(keep) 21 | x.append(keep) 22 | temp.close() 23 | 24 | x = np.array(x, dtype=np.uint8) 25 | #np.random.shuffle(x) 26 | 27 | p = int(ratio * len(x)) 28 | x_train = x[:p] 29 | x_test = x[p:] 30 | 31 | if not os.path.exists('./npy'): 32 | os.mkdir('./npy') 33 | np.save('./npy/x_train.npy', x_train) 34 | np.save('./npy/x_test.npy', x_test) -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /decensor.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | from PIL import Image 4 | import tqdm 5 | import os 6 | import matplotlib.pyplot as plt 7 | import sys 8 | sys.path.append('..') 9 | from model import Model 10 | from poisson_blend import blend 11 | 12 | IMAGE_SIZE = 128 13 | LOCAL_SIZE = 64 14 | HOLE_MIN = 24 15 | HOLE_MAX = 48 16 | BATCH_SIZE = 1 17 | 18 | 19 | image_folder = 'decensor_input_images/' 20 | mask_color = [0, 255, 0] 21 | poisson_blending_enabled = False 22 | 23 | def decensor(): 24 | x = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 25 | mask = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 1]) 26 | local_x = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 27 | global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 28 | local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 29 | is_training = tf.placeholder(tf.bool, []) 30 | 31 | model = Model(x, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) 32 | sess = tf.Session() 33 | init_op = tf.global_variables_initializer() 34 | sess.run(init_op) 35 | 36 | saver = tf.train.Saver() 37 | saver.restore(sess, './models/latest') 38 | 39 | x_decensor = [] 40 | mask_decensor = [] 41 | for subdir, dirs, files in sorted(os.walk(image_folder)): 42 | for file in sorted(files): 43 | file_path = os.path.join(subdir, file) 44 | if os.path.isfile(file_path) and os.path.splitext(file_path)[1] == ".png": 45 | print(file_path) 46 | image = Image.open(file_path).convert('RGB') 47 | image = np.array(image) 48 | image = np.array(image / 127.5 - 1) 49 | x_decensor.append(image) 50 | x_decensor = np.array(x_decensor) 51 | print(x_decensor.shape) 52 | step_num = int(len(x_decensor) / BATCH_SIZE) 53 | 54 | cnt = 0 55 | for i in tqdm.tqdm(range(step_num)): 56 | x_batch = x_decensor[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] 57 | mask_batch = get_mask(x_batch) 58 | completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) 59 | for i in range(BATCH_SIZE): 60 | cnt += 1 61 | img = completion[i] 62 | img = np.array((img + 1) * 127.5, dtype=np.uint8) 63 | original = x_batch[i] 64 | original = np.array((original + 1) * 127.5, dtype=np.uint8) 65 | if (poisson_blending_enabled): 66 | img = blend(original, img, mask_batch[0,:,:,0]) 67 | output = Image.fromarray(img.astype('uint8'), 'RGB') 68 | dst = './decensor_output_images/{}.png'.format("{0:06d}".format(cnt)) 69 | output.save(dst) 70 | 71 | def get_mask(x_batch): 72 | points = [] 73 | mask = [] 74 | for i in range(BATCH_SIZE): 75 | raw = x_batch[i] 76 | raw = np.array((raw + 1) * 127.5, dtype=np.uint8) 77 | m = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 1), dtype=np.uint8) 78 | for x in range(IMAGE_SIZE): 79 | for y in range(IMAGE_SIZE): 80 | if np.array_equal(raw[x][y], [0, 255, 0]): 81 | m[x, y] = 1 82 | mask.append(m) 83 | return np.array(mask) 84 | 85 | if __name__ == '__main__': 86 | decensor() 87 | 88 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | from PIL import Image 4 | import tqdm 5 | import os 6 | import matplotlib.pyplot as plt 7 | import sys 8 | sys.path.append('..') 9 | from model import Model 10 | 11 | IMAGE_SIZE = 128 12 | LOCAL_SIZE = 64 13 | HOLE_MIN = 24 14 | HOLE_MAX = 48 15 | BATCH_SIZE = 16 16 | 17 | test_npy = './lfw.npy' 18 | 19 | def test(): 20 | x = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 21 | mask = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 1]) 22 | local_x = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 23 | global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 24 | local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 25 | is_training = tf.placeholder(tf.bool, []) 26 | 27 | model = Model(x, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) 28 | sess = tf.Session() 29 | init_op = tf.global_variables_initializer() 30 | sess.run(init_op) 31 | 32 | saver = tf.train.Saver() 33 | saver.restore(sess, './models/latest') 34 | 35 | x_test = np.load(test_npy) 36 | np.random.shuffle(x_test) 37 | x_test = np.array([a / 127.5 - 1 for a in x_test]) 38 | 39 | step_num = int(len(x_test) / BATCH_SIZE) 40 | 41 | cnt = 0 42 | for i in tqdm.tqdm(range(step_num)): 43 | x_batch = x_test[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] 44 | _, mask_batch = get_points() 45 | completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) 46 | for i in range(BATCH_SIZE): 47 | cnt += 1 48 | raw = x_batch[i] 49 | raw = np.array((raw + 1) * 127.5, dtype=np.uint8) 50 | masked = raw * (1 - mask_batch[i]) + np.ones_like(raw) * mask_batch[i] * 255 51 | img = completion[i] 52 | img = np.array((img + 1) * 127.5, dtype=np.uint8) 53 | dst = './testing_output_images/{}.jpg'.format("{0:06d}".format(cnt)) 54 | output_image([['Input', masked], ['Output', img], ['Ground Truth', raw]], dst) 55 | 56 | 57 | def get_points(): 58 | points = [] 59 | mask = [] 60 | for i in range(BATCH_SIZE): 61 | x1, y1 = np.random.randint(0, IMAGE_SIZE - LOCAL_SIZE + 1, 2) 62 | x2, y2 = np.array([x1, y1]) + LOCAL_SIZE 63 | points.append([x1, y1, x2, y2]) 64 | 65 | w, h = np.random.randint(HOLE_MIN, HOLE_MAX + 1, 2) 66 | p1 = x1 + np.random.randint(0, LOCAL_SIZE - w) 67 | q1 = y1 + np.random.randint(0, LOCAL_SIZE - h) 68 | p2 = p1 + w 69 | q2 = q1 + h 70 | 71 | m = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 1), dtype=np.uint8) 72 | m[q1:q2 + 1, p1:p2 + 1] = 1 73 | mask.append(m) 74 | 75 | return np.array(points), np.array(mask) 76 | 77 | 78 | def output_image(images, dst): 79 | fig = plt.figure() 80 | for i, image in enumerate(images): 81 | text, img = image 82 | fig.add_subplot(1, 3, i + 1) 83 | plt.imshow(img) 84 | plt.tick_params(labelbottom='off') 85 | plt.tick_params(labelleft='off') 86 | plt.gca().get_xaxis().set_ticks_position('none') 87 | plt.gca().get_yaxis().set_ticks_position('none') 88 | plt.xlabel(text) 89 | plt.savefig(dst) 90 | plt.close() 91 | 92 | 93 | if __name__ == '__main__': 94 | test() 95 | 96 | -------------------------------------------------------------------------------- /layer.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def conv_layer(x, filter_shape, stride): 4 | filters = tf.get_variable( 5 | name='weight', 6 | shape=filter_shape, 7 | dtype=tf.float32, 8 | initializer=tf.contrib.layers.xavier_initializer(), 9 | trainable=True) 10 | return tf.nn.conv2d(x, filters, [1, stride, stride, 1], padding='SAME') 11 | 12 | 13 | def dilated_conv_layer(x, filter_shape, dilation): 14 | filters = tf.get_variable( 15 | name='weight', 16 | shape=filter_shape, 17 | dtype=tf.float32, 18 | initializer=tf.contrib.layers.xavier_initializer(), 19 | trainable=True) 20 | return tf.nn.atrous_conv2d(x, filters, dilation, padding='SAME') 21 | 22 | 23 | def deconv_layer(x, filter_shape, output_shape, stride): 24 | filters = tf.get_variable( 25 | name='weight', 26 | shape=filter_shape, 27 | dtype=tf.float32, 28 | initializer=tf.contrib.layers.xavier_initializer(), 29 | trainable=True) 30 | return tf.nn.conv2d_transpose(x, filters, output_shape, [1, stride, stride, 1]) 31 | 32 | 33 | def batch_normalize(x, is_training, decay=0.99, epsilon=0.001): 34 | def bn_train(): 35 | batch_mean, batch_var = tf.nn.moments(x, axes=[0, 1, 2]) 36 | train_mean = tf.assign(pop_mean, pop_mean * decay + batch_mean * (1 - decay)) 37 | train_var = tf.assign(pop_var, pop_var * decay + batch_var * (1 - decay)) 38 | with tf.control_dependencies([train_mean, train_var]): 39 | return tf.nn.batch_normalization(x, batch_mean, batch_var, beta, scale, epsilon) 40 | 41 | def bn_inference(): 42 | return tf.nn.batch_normalization(x, pop_mean, pop_var, beta, scale, epsilon) 43 | 44 | dim = x.get_shape().as_list()[-1] 45 | beta = tf.get_variable( 46 | name='beta', 47 | shape=[dim], 48 | dtype=tf.float32, 49 | initializer=tf.truncated_normal_initializer(stddev=0.0), 50 | trainable=True) 51 | scale = tf.get_variable( 52 | name='scale', 53 | shape=[dim], 54 | dtype=tf.float32, 55 | initializer=tf.truncated_normal_initializer(stddev=0.1), 56 | trainable=True) 57 | pop_mean = tf.get_variable( 58 | name='pop_mean', 59 | shape=[dim], 60 | dtype=tf.float32, 61 | initializer=tf.constant_initializer(0.0), 62 | trainable=False) 63 | pop_var = tf.get_variable( 64 | name='pop_var', 65 | shape=[dim], 66 | dtype=tf.float32, 67 | initializer=tf.constant_initializer(1.0), 68 | trainable=False) 69 | 70 | return tf.cond(is_training, bn_train, bn_inference) 71 | 72 | 73 | def flatten_layer(x): 74 | input_shape = x.get_shape().as_list() 75 | dim = input_shape[1] * input_shape[2] * input_shape[3] 76 | transposed = tf.transpose(x, (0, 3, 1, 2)) 77 | return tf.reshape(transposed, [-1, dim]) 78 | 79 | 80 | def full_connection_layer(x, out_dim): 81 | in_dim = x.get_shape().as_list()[-1] 82 | W = tf.get_variable( 83 | name='weight', 84 | shape=[in_dim, out_dim], 85 | dtype=tf.float32, 86 | initializer=tf.truncated_normal_initializer(stddev=0.1), 87 | trainable=True) 88 | b = tf.get_variable( 89 | name='bias', 90 | shape=[out_dim], 91 | dtype=tf.float32, 92 | initializer=tf.constant_initializer(0.0), 93 | trainable=True) 94 | return tf.add(tf.matmul(x, W), b) 95 | 96 | -------------------------------------------------------------------------------- /poisson_blend.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | import numpy as np 5 | import scipy.sparse 6 | import PIL.Image 7 | import pyamg 8 | import copy 9 | 10 | # pre-process the mask array so that uint64 types from opencv.imread can be adapted 11 | def prepare_mask(mask): 12 | result = np.ndarray((mask.shape[0], mask.shape[1]), dtype=np.uint8) 13 | for i in range(mask.shape[0]): 14 | for j in range(mask.shape[1]): 15 | if mask[i][j] > 0: 16 | result[i][j] = 1 17 | else: 18 | result[i][j] = 0 19 | mask = result 20 | return mask 21 | 22 | def blend(img_target, img_source, img_mask, offset=(0, 0)): 23 | # compute regions to be blended 24 | region_source = ( 25 | max(-offset[0], 0), 26 | max(-offset[1], 0), 27 | min(img_target.shape[0]-offset[0], img_source.shape[0]), 28 | min(img_target.shape[1]-offset[1], img_source.shape[1])) 29 | region_target = ( 30 | max(offset[0], 0), 31 | max(offset[1], 0), 32 | min(img_target.shape[0], img_source.shape[0]+offset[0]), 33 | min(img_target.shape[1], img_source.shape[1]+offset[1])) 34 | region_size = (region_source[2]-region_source[0], region_source[3]-region_source[1]) 35 | 36 | # clip and normalize mask image 37 | img_mask = img_mask[region_source[0]:region_source[2], region_source[1]:region_source[3]] 38 | #img_mask_copy = copy.deepcopy(img_mask) 39 | # prepare_mask doesn't change anything 40 | # img_mask = prepare_mask(img_mask) 41 | # if np.array_equal(img_mask, img_mask_copy): 42 | # print "eq" 43 | img_mask[img_mask==0] = False 44 | img_mask[img_mask!=False] = True 45 | 46 | # create coefficient matrix 47 | A = scipy.sparse.identity(np.prod(region_size), format='lil') 48 | for y in range(region_size[0]): 49 | for x in range(region_size[1]): 50 | if img_mask[y,x]: 51 | index = x+y*region_size[1] 52 | A[index, index] = 4 53 | if index+1 < np.prod(region_size): 54 | A[index, index+1] = -1 55 | if index-1 >= 0: 56 | A[index, index-1] = -1 57 | if index+region_size[1] < np.prod(region_size): 58 | A[index, index+region_size[1]] = -1 59 | if index-region_size[1] >= 0: 60 | A[index, index-region_size[1]] = -1 61 | A = A.tocsr() 62 | 63 | # create poisson matrix for b 64 | P = pyamg.gallery.poisson(img_mask.shape) 65 | 66 | # for each layer (ex. RGB) 67 | for num_layer in range(img_target.shape[2]): 68 | # get subimages 69 | t = img_target[region_target[0]:region_target[2],region_target[1]:region_target[3],num_layer] 70 | s = img_source[region_source[0]:region_source[2], region_source[1]:region_source[3],num_layer] 71 | t = t.flatten() 72 | s = s.flatten() 73 | 74 | # create b 75 | b = P * s 76 | for y in range(region_size[0]): 77 | for x in range(region_size[1]): 78 | if not img_mask[y,x]: 79 | index = x+y*region_size[1] 80 | b[index] = t[index] 81 | 82 | # solve Ax = b 83 | x = pyamg.solve(A,b,verb=False,tol=1e-10) 84 | 85 | # assign x to target image 86 | x = np.reshape(x, region_size) 87 | x[x>255] = 255 88 | x[x<0] = 0 89 | x = np.array(x, img_target.dtype) 90 | img_target[region_target[0]:region_target[2],region_target[1]:region_target[3],num_layer] = x 91 | 92 | return img_target 93 | 94 | 95 | def test(): 96 | img_mask = np.asarray(PIL.Image.open('./testimages/test1_mask.png')) 97 | img_mask.flags.writeable = True 98 | img_source = np.asarray(PIL.Image.open('./testimages/test1_src.png')) 99 | img_source.flags.writeable = True 100 | img_target = np.asarray(PIL.Image.open('./testimages/test1_target.png')) 101 | img_target.flags.writeable = True 102 | img_ret = blend(img_target, img_source, img_mask, offset=(40,-30)) 103 | img_ret = PIL.Image.fromarray(np.uint8(img_ret)) 104 | img_ret.save('./testimages/test1_ret.png') 105 | 106 | 107 | if __name__ == '__main__': 108 | test() 109 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeepMindBreak 2 | *Decensoring Hentai with Deep Neural Networks* 3 | 4 | This project applies an implementation of [Globally and Locally Consistent Image Completion](http://hi.cs.waseda.ac.jp/%7Eiizuka/projects/completion/data/completion_sig2017.pdf) to the problem of hentai decensorship. Using a deep fully convolutional neural network, DeepMindBreak can replace censored artwork in hentai with plausible reconstructions. The user needs to only specify the censored regions. 5 | 6 | # **THIS PROJECT IS STILL IN DEVELOPMENT. DO NOT BE DISAPPOINTED IF THE RESULTS AREN'T AS GOOD AS YOU EXPECT.** 7 | 8 | ![Censored, decensored](/readme_images/collage.png) 9 | 10 | # Limitations 11 | 12 | This project is LIMITED in capability. It is a proof of concept of ongoing research. 13 | 14 | The decensorship is intended to ONLY work on color hentai images that have minor bar censorship of the penis or vagina. 15 | 16 | It does NOT work with: 17 | - Black and white images 18 | - Monochrome images 19 | - Hentai containing screentones (e.g. printed hentai) 20 | - Real life porn 21 | - Mosaic censorship 22 | - Censorship of nipples 23 | - Censorship of anus 24 | - Animated gifs/videos 25 | 26 | In particular, if a vagina or penis is completely censored out, inpainting will be ineffective. 27 | 28 | Embarrassingly, because the neural network was trained to decensor horizontally and vertically oriented rectangles, it has trouble with angled rectangles. This will be fixed soon. 29 | 30 | # Dependencies 31 | 32 | - Python 2/3 33 | - TensorFlow 1.5 34 | - Pillow 35 | - tqdm 36 | - scipy 37 | - pyamg (only needed if poisson blending is enabled in decensor.py) 38 | - matplotlib (only for running test.py) 39 | 40 | No GPU required! Tested on Ubuntu 16.04 and Windows. 41 | 42 | Poisson blending is disabled by default since it has little effect on output quality. 43 | 44 | Pillow, tqdm, scipy, pyamg, and matplotlib can all be installed using pip. 45 | 46 | # Model 47 | The pretrained model can be downloaded from https://drive.google.com/open?id=1mWHYSj0LDSbJQQxjR4hUMykQkVve2U3Q. 48 | 49 | Unzip the contents into the /models/ folder. 50 | 51 | # Usage 52 | 53 | ## I. Decensoring hentai 54 | 55 | ![Guide](/readme_images/guide.png) 56 | 57 | The decensorship process is fairly involved. A user interface will eventually be released to streamline the process. 58 | 59 | Using image editing software like Photoshop or GIMP, paint the areas you want to decensor the color with RGB values of (0,255,0). For each censored region, crop 128 x 128 size images containing the censored regions from your images and save them as new ".png" images. 60 | 61 | Move the cropped images to the "decensor_input_images" directory. Decensor the images by running 62 | 63 | ``` 64 | $ python decensor.py 65 | ``` 66 | 67 | Decensored images will be saved to the "decensor_output_images" directory. Paste the decensored images back into the original image. 68 | 69 | ## II. Train the pretrained model on custom dataset 70 | 71 | Put your custom dataset for training in the "data/images" directory and convert images to npy format. 72 | 73 | ``` 74 | $ cd training_data 75 | $ python to_npy.py 76 | ``` 77 | 78 | Train pretrained model on your custom dataset. 79 | 80 | ``` 81 | $ python train.py 82 | ``` 83 | 84 | # To do 85 | - ~~Add Python 3 compatibility~~ 86 | - Add random rotations in cropping rectangles 87 | - Retrain for arbitrary shape censors 88 | - Add a user interface 89 | - Incorporate GAN loss into training 90 | - Update the model to the new version 91 | 92 | Contributions are welcome! 93 | 94 | # License 95 | 96 | Example image by dannychoo under [CC BY-NC-SA 2.0 License](https://creativecommons.org/licenses/by-nc-sa/2.0/). The example image is modified from the original, which can be found [here](https://www.flickr.com/photos/dannychoo/16081096643/in/photostream/). 97 | 98 | Model is licensed under CC BY-NC 3.0 License. 99 | 100 | Code is licensed under MIT License and is modified from tadax's project [Globally and Locally Consistent Image Completion with TensorFlow ](https://github.com/tadax/glcic), which is an implementation of the paper [Globally and Locally Consistent Image Completion](http://hi.cs.waseda.ac.jp/%7Eiizuka/projects/completion/data/completion_sig2017.pdf). It also has a modified version of parosky's project [poissonblending](https://github.com/parosky/poissonblending). 101 | 102 | --- 103 | 104 | Copyright (c) 2018 tadax, parosky, deeppomf 105 | 106 | Permission is hereby granted, free of charge, to any person obtaining a copy 107 | of this software and associated documentation files (the "Software"), to deal 108 | in the Software without restriction, including without limitation the rights 109 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 110 | copies of the Software, and to permit persons to whom the Software is 111 | furnished to do so, subject to the following conditions: 112 | 113 | The above copyright notice and this permission notice shall be included in all 114 | copies or substantial portions of the Software. 115 | 116 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 117 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 118 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 119 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 120 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 121 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 122 | SOFTWARE. 123 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | from PIL import Image 4 | import tqdm 5 | from model import Model 6 | import load 7 | import scipy.ndimage 8 | 9 | IMAGE_SIZE = 128 10 | LOCAL_SIZE = 64 11 | HOLE_MIN = 24 12 | HOLE_MAX = 48 13 | LEARNING_RATE = 1e-3 14 | BATCH_SIZE = 16 15 | PRETRAIN_EPOCH = 100 16 | #the chance the rectangle crop will be rotated 17 | ROTATE_CHANCE = 0.5 18 | 19 | def train(): 20 | x = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 21 | mask = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 1]) 22 | local_x = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 23 | global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 24 | local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 25 | is_training = tf.placeholder(tf.bool, []) 26 | 27 | model = Model(x, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) 28 | sess = tf.Session() 29 | global_step = tf.Variable(0, name='global_step', trainable=False) 30 | epoch = tf.Variable(0, name='epoch', trainable=False) 31 | 32 | opt = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE) 33 | g_train_op = opt.minimize(model.g_loss, global_step=global_step, var_list=model.g_variables) 34 | d_train_op = opt.minimize(model.d_loss, global_step=global_step, var_list=model.d_variables) 35 | 36 | init_op = tf.global_variables_initializer() 37 | sess.run(init_op) 38 | 39 | if tf.train.get_checkpoint_state('./models'): 40 | saver = tf.train.Saver() 41 | saver.restore(sess, './models/latest') 42 | 43 | x_train, x_test = load.load() 44 | x_train = np.array([a / 127.5 - 1 for a in x_train]) 45 | x_test = np.array([a / 127.5 - 1 for a in x_test]) 46 | 47 | step_num = int(len(x_train) / BATCH_SIZE) 48 | 49 | while True: 50 | sess.run(tf.assign(epoch, tf.add(epoch, 1))) 51 | print('epoch: {}'.format(sess.run(epoch))) 52 | 53 | np.random.shuffle(x_train) 54 | 55 | # Completion 56 | if sess.run(epoch) <= PRETRAIN_EPOCH: 57 | g_loss_value = 0 58 | for i in tqdm.tqdm(range(step_num)): 59 | x_batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] 60 | points_batch, mask_batch = get_points() 61 | 62 | _, g_loss = sess.run([g_train_op, model.g_loss], feed_dict={x: x_batch, mask: mask_batch, is_training: True}) 63 | g_loss_value += g_loss 64 | 65 | print('Completion loss: {}'.format(g_loss_value)) 66 | 67 | np.random.shuffle(x_test) 68 | x_batch = x_test[:BATCH_SIZE] 69 | completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) 70 | sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) 71 | result = Image.fromarray(sample) 72 | result.save('./training_output_images/{}.jpg'.format("{0:06d}".format(sess.run(epoch)))) 73 | 74 | saver = tf.train.Saver() 75 | saver.save(sess, './models/latest', write_meta_graph=False) 76 | if sess.run(epoch) == PRETRAIN_EPOCH: 77 | saver.save(sess, './models/pretrained', write_meta_graph=False) 78 | 79 | 80 | # Discrimitation 81 | else: 82 | g_loss_value = 0 83 | d_loss_value = 0 84 | for i in tqdm.tqdm(range(step_num)): 85 | x_batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] 86 | points_batch, mask_batch = get_points() 87 | 88 | _, g_loss, completion = sess.run([g_train_op, model.g_loss, model.completion], feed_dict={x: x_batch, mask: mask_batch, is_training: True}) 89 | g_loss_value += g_loss 90 | 91 | local_x_batch = [] 92 | local_completion_batch = [] 93 | for i in range(BATCH_SIZE): 94 | x1, y1, x2, y2 = points_batch[i] 95 | local_x_batch.append(x_batch[i][y1:y2, x1:x2, :]) 96 | local_completion_batch.append(completion[i][y1:y2, x1:x2, :]) 97 | local_x_batch = np.array(local_x_batch) 98 | local_completion_batch = np.array(local_completion_batch) 99 | 100 | _, d_loss = sess.run( 101 | [d_train_op, model.d_loss], 102 | feed_dict={x: x_batch, mask: mask_batch, local_x: local_x_batch, global_completion: completion, local_completion: local_completion_batch, is_training: True}) 103 | d_loss_value += d_loss 104 | 105 | print('Completion loss: {}'.format(g_loss_value)) 106 | print('Discriminator loss: {}'.format(d_loss_value)) 107 | 108 | np.random.shuffle(x_test) 109 | x_batch = x_test[:BATCH_SIZE] 110 | completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False}) 111 | sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) 112 | result = Image.fromarray(sample) 113 | result.save('./training_output_images/{}.jpg'.format("{0:06d}".format(sess.run(epoch)))) 114 | 115 | saver = tf.train.Saver() 116 | saver.save(sess, './models/latest', write_meta_graph=False) 117 | 118 | 119 | def get_points(): 120 | points = [] 121 | mask = [] 122 | for i in range(BATCH_SIZE): 123 | x1, y1 = np.random.randint(0, IMAGE_SIZE - LOCAL_SIZE + 1, 2) 124 | x2, y2 = np.array([x1, y1]) + LOCAL_SIZE 125 | points.append([x1, y1, x2, y2]) 126 | 127 | w, h = np.random.randint(HOLE_MIN, HOLE_MAX + 1, 2) 128 | p1 = x1 + np.random.randint(0, LOCAL_SIZE - w) 129 | q1 = y1 + np.random.randint(0, LOCAL_SIZE - h) 130 | p2 = p1 + w 131 | q2 = q1 + h 132 | 133 | m = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 1), dtype=np.uint8) 134 | m[q1:q2 + 1, p1:p2 + 1] = 1 135 | 136 | if (np.random.random() < ROTATE_CHANCE): 137 | #rotate random amount between 0 and 90 degrees 138 | m = scipy.ndimage.rotate(m, np.random.random()*90, reshape = False) 139 | #set all elements greater than 0 to 1 140 | m[m > 0] = 1 141 | 142 | mask.append(m) 143 | 144 | return np.array(points), np.array(mask) 145 | 146 | 147 | if __name__ == '__main__': 148 | train() -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | from layer import * 2 | 3 | class Model: 4 | def __init__(self, x, mask, local_x, global_completion, local_completion, is_training, batch_size): 5 | self.batch_size = batch_size 6 | self.imitation = self.generator(x * (1 - mask), is_training) 7 | self.completion = self.imitation * mask + x * (1 - mask) 8 | self.real = self.discriminator(x, local_x, reuse=False) 9 | self.fake = self.discriminator(global_completion, local_completion, reuse=True) 10 | self.g_loss = self.calc_g_loss(x, self.completion) 11 | self.d_loss = self.calc_d_loss(self.real, self.fake) 12 | self.g_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator') 13 | self.d_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator') 14 | 15 | 16 | def generator(self, x, is_training): 17 | with tf.variable_scope('generator'): 18 | with tf.variable_scope('conv1'): 19 | x = conv_layer(x, [5, 5, 3, 64], 1) 20 | x = batch_normalize(x, is_training) 21 | x = tf.nn.relu(x) 22 | with tf.variable_scope('conv2'): 23 | x = conv_layer(x, [3, 3, 64, 128], 2) 24 | x = batch_normalize(x, is_training) 25 | x = tf.nn.relu(x) 26 | with tf.variable_scope('conv3'): 27 | x = conv_layer(x, [3, 3, 128, 128], 1) 28 | x = batch_normalize(x, is_training) 29 | x = tf.nn.relu(x) 30 | with tf.variable_scope('conv4'): 31 | x = conv_layer(x, [3, 3, 128, 256], 2) 32 | x = batch_normalize(x, is_training) 33 | x = tf.nn.relu(x) 34 | with tf.variable_scope('conv5'): 35 | x = conv_layer(x, [3, 3, 256, 256], 1) 36 | x = batch_normalize(x, is_training) 37 | x = tf.nn.relu(x) 38 | with tf.variable_scope('conv6'): 39 | x = conv_layer(x, [3, 3, 256, 256], 1) 40 | x = batch_normalize(x, is_training) 41 | x = tf.nn.relu(x) 42 | with tf.variable_scope('dilated1'): 43 | x = dilated_conv_layer(x, [3, 3, 256, 256], 2) 44 | x = batch_normalize(x, is_training) 45 | x = tf.nn.relu(x) 46 | with tf.variable_scope('dilated2'): 47 | x = dilated_conv_layer(x, [3, 3, 256, 256], 4) 48 | x = batch_normalize(x, is_training) 49 | x = tf.nn.relu(x) 50 | with tf.variable_scope('dilated3'): 51 | x = dilated_conv_layer(x, [3, 3, 256, 256], 8) 52 | x = batch_normalize(x, is_training) 53 | x = tf.nn.relu(x) 54 | with tf.variable_scope('dilated4'): 55 | x = dilated_conv_layer(x, [3, 3, 256, 256], 16) 56 | x = batch_normalize(x, is_training) 57 | x = tf.nn.relu(x) 58 | with tf.variable_scope('conv7'): 59 | x = conv_layer(x, [3, 3, 256, 256], 1) 60 | x = batch_normalize(x, is_training) 61 | x = tf.nn.relu(x) 62 | with tf.variable_scope('conv8'): 63 | x = conv_layer(x, [3, 3, 256, 256], 1) 64 | x = batch_normalize(x, is_training) 65 | x = tf.nn.relu(x) 66 | with tf.variable_scope('deconv1'): 67 | x = deconv_layer(x, [4, 4, 128, 256], [self.batch_size, 64, 64, 128], 2) 68 | x = batch_normalize(x, is_training) 69 | x = tf.nn.relu(x) 70 | with tf.variable_scope('conv9'): 71 | x = conv_layer(x, [3, 3, 128, 128], 1) 72 | x = batch_normalize(x, is_training) 73 | x = tf.nn.relu(x) 74 | with tf.variable_scope('deconv2'): 75 | x = deconv_layer(x, [4, 4, 64, 128], [self.batch_size, 128, 128, 64], 2) 76 | x = batch_normalize(x, is_training) 77 | x = tf.nn.relu(x) 78 | with tf.variable_scope('conv10'): 79 | x = conv_layer(x, [3, 3, 64, 32], 1) 80 | x = batch_normalize(x, is_training) 81 | x = tf.nn.relu(x) 82 | with tf.variable_scope('conv11'): 83 | x = conv_layer(x, [3, 3, 32, 3], 1) 84 | x = tf.nn.tanh(x) 85 | 86 | return x 87 | 88 | 89 | def discriminator(self, global_x, local_x, reuse): 90 | def global_discriminator(x): 91 | is_training = tf.constant(True) 92 | with tf.variable_scope('global'): 93 | with tf.variable_scope('conv1'): 94 | x = conv_layer(x, [5, 5, 3, 64], 2) 95 | x = batch_normalize(x, is_training) 96 | x = tf.nn.relu(x) 97 | with tf.variable_scope('conv2'): 98 | x = conv_layer(x, [5, 5, 64, 128], 2) 99 | x = batch_normalize(x, is_training) 100 | x = tf.nn.relu(x) 101 | with tf.variable_scope('conv3'): 102 | x = conv_layer(x, [5, 5, 128, 256], 2) 103 | x = batch_normalize(x, is_training) 104 | x = tf.nn.relu(x) 105 | with tf.variable_scope('conv4'): 106 | x = conv_layer(x, [5, 5, 256, 512], 2) 107 | x = batch_normalize(x, is_training) 108 | x = tf.nn.relu(x) 109 | with tf.variable_scope('conv5'): 110 | x = conv_layer(x, [5, 5, 512, 512], 2) 111 | x = batch_normalize(x, is_training) 112 | x = tf.nn.relu(x) 113 | with tf.variable_scope('fc'): 114 | x = flatten_layer(x) 115 | x = full_connection_layer(x, 1024) 116 | return x 117 | 118 | def local_discriminator(x): 119 | is_training = tf.constant(True) 120 | with tf.variable_scope('local'): 121 | with tf.variable_scope('conv1'): 122 | x = conv_layer(x, [5, 5, 3, 64], 2) 123 | x = batch_normalize(x, is_training) 124 | x = tf.nn.relu(x) 125 | with tf.variable_scope('conv2'): 126 | x = conv_layer(x, [5, 5, 64, 128], 2) 127 | x = batch_normalize(x, is_training) 128 | x = tf.nn.relu(x) 129 | with tf.variable_scope('conv3'): 130 | x = conv_layer(x, [5, 5, 128, 256], 2) 131 | x = batch_normalize(x, is_training) 132 | x = tf.nn.relu(x) 133 | with tf.variable_scope('conv4'): 134 | x = conv_layer(x, [5, 5, 256, 512], 2) 135 | x = batch_normalize(x, is_training) 136 | x = tf.nn.relu(x) 137 | with tf.variable_scope('fc'): 138 | x = flatten_layer(x) 139 | x = full_connection_layer(x, 1024) 140 | return x 141 | 142 | with tf.variable_scope('discriminator', reuse=reuse): 143 | global_output = global_discriminator(global_x) 144 | local_output = local_discriminator(local_x) 145 | with tf.variable_scope('concatenation'): 146 | output = tf.concat((global_output, local_output), 1) 147 | output = full_connection_layer(output, 1) 148 | 149 | return output 150 | 151 | 152 | def calc_g_loss(self, x, completion): 153 | loss = tf.nn.l2_loss(x - completion) 154 | return tf.reduce_mean(loss) 155 | 156 | 157 | def calc_d_loss(self, real, fake): 158 | alpha = 4e-4 159 | d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=real, labels=tf.ones_like(real))) 160 | d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake, labels=tf.zeros_like(fake))) 161 | return tf.add(d_loss_real, d_loss_fake) * alpha 162 | 163 | -------------------------------------------------------------------------------- /model_mosaic.py: -------------------------------------------------------------------------------- 1 | from layer import * 2 | 3 | class Model: 4 | def __init__(self, x, mosaic, mask, local_x, global_completion, local_completion, is_training, batch_size): 5 | self.batch_size = batch_size 6 | self.merged = x * (1 - mask) + mosaic * (mask) 7 | self.imitation = self.generator(self.merged, is_training) 8 | self.completion = self.imitation * mask + x * (1 - mask) 9 | self.real = self.discriminator(x, local_x, reuse=False) 10 | self.fake = self.discriminator(global_completion, local_completion, reuse=True) 11 | self.g_loss = self.calc_g_loss(x, self.completion) 12 | self.d_loss = self.calc_d_loss(self.real, self.fake) 13 | self.g_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='generator') 14 | self.d_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator') 15 | 16 | 17 | def generator(self, x, is_training): 18 | with tf.variable_scope('generator'): 19 | with tf.variable_scope('conv1'): 20 | x = conv_layer(x, [5, 5, 3, 64], 1) 21 | x = batch_normalize(x, is_training) 22 | x = tf.nn.relu(x) 23 | with tf.variable_scope('conv2'): 24 | x = conv_layer(x, [3, 3, 64, 128], 2) 25 | x = batch_normalize(x, is_training) 26 | x = tf.nn.relu(x) 27 | with tf.variable_scope('conv3'): 28 | x = conv_layer(x, [3, 3, 128, 128], 1) 29 | x = batch_normalize(x, is_training) 30 | x = tf.nn.relu(x) 31 | with tf.variable_scope('conv4'): 32 | x = conv_layer(x, [3, 3, 128, 256], 2) 33 | x = batch_normalize(x, is_training) 34 | x = tf.nn.relu(x) 35 | with tf.variable_scope('conv5'): 36 | x = conv_layer(x, [3, 3, 256, 256], 1) 37 | x = batch_normalize(x, is_training) 38 | x = tf.nn.relu(x) 39 | with tf.variable_scope('conv6'): 40 | x = conv_layer(x, [3, 3, 256, 256], 1) 41 | x = batch_normalize(x, is_training) 42 | x = tf.nn.relu(x) 43 | with tf.variable_scope('dilated1'): 44 | x = dilated_conv_layer(x, [3, 3, 256, 256], 2) 45 | x = batch_normalize(x, is_training) 46 | x = tf.nn.relu(x) 47 | with tf.variable_scope('dilated2'): 48 | x = dilated_conv_layer(x, [3, 3, 256, 256], 4) 49 | x = batch_normalize(x, is_training) 50 | x = tf.nn.relu(x) 51 | with tf.variable_scope('dilated3'): 52 | x = dilated_conv_layer(x, [3, 3, 256, 256], 8) 53 | x = batch_normalize(x, is_training) 54 | x = tf.nn.relu(x) 55 | with tf.variable_scope('dilated4'): 56 | x = dilated_conv_layer(x, [3, 3, 256, 256], 16) 57 | x = batch_normalize(x, is_training) 58 | x = tf.nn.relu(x) 59 | with tf.variable_scope('conv7'): 60 | x = conv_layer(x, [3, 3, 256, 256], 1) 61 | x = batch_normalize(x, is_training) 62 | x = tf.nn.relu(x) 63 | with tf.variable_scope('conv8'): 64 | x = conv_layer(x, [3, 3, 256, 256], 1) 65 | x = batch_normalize(x, is_training) 66 | x = tf.nn.relu(x) 67 | with tf.variable_scope('deconv1'): 68 | x = deconv_layer(x, [4, 4, 128, 256], [self.batch_size, 64, 64, 128], 2) 69 | x = batch_normalize(x, is_training) 70 | x = tf.nn.relu(x) 71 | with tf.variable_scope('conv9'): 72 | x = conv_layer(x, [3, 3, 128, 128], 1) 73 | x = batch_normalize(x, is_training) 74 | x = tf.nn.relu(x) 75 | with tf.variable_scope('deconv2'): 76 | x = deconv_layer(x, [4, 4, 64, 128], [self.batch_size, 128, 128, 64], 2) 77 | x = batch_normalize(x, is_training) 78 | x = tf.nn.relu(x) 79 | with tf.variable_scope('conv10'): 80 | x = conv_layer(x, [3, 3, 64, 32], 1) 81 | x = batch_normalize(x, is_training) 82 | x = tf.nn.relu(x) 83 | with tf.variable_scope('conv11'): 84 | x = conv_layer(x, [3, 3, 32, 3], 1) 85 | x = tf.nn.tanh(x) 86 | 87 | return x 88 | 89 | 90 | def discriminator(self, global_x, local_x, reuse): 91 | def global_discriminator(x): 92 | is_training = tf.constant(True) 93 | with tf.variable_scope('global'): 94 | with tf.variable_scope('conv1'): 95 | x = conv_layer(x, [5, 5, 3, 64], 2) 96 | x = batch_normalize(x, is_training) 97 | x = tf.nn.relu(x) 98 | with tf.variable_scope('conv2'): 99 | x = conv_layer(x, [5, 5, 64, 128], 2) 100 | x = batch_normalize(x, is_training) 101 | x = tf.nn.relu(x) 102 | with tf.variable_scope('conv3'): 103 | x = conv_layer(x, [5, 5, 128, 256], 2) 104 | x = batch_normalize(x, is_training) 105 | x = tf.nn.relu(x) 106 | with tf.variable_scope('conv4'): 107 | x = conv_layer(x, [5, 5, 256, 512], 2) 108 | x = batch_normalize(x, is_training) 109 | x = tf.nn.relu(x) 110 | with tf.variable_scope('conv5'): 111 | x = conv_layer(x, [5, 5, 512, 512], 2) 112 | x = batch_normalize(x, is_training) 113 | x = tf.nn.relu(x) 114 | with tf.variable_scope('fc'): 115 | x = flatten_layer(x) 116 | x = full_connection_layer(x, 1024) 117 | return x 118 | 119 | def local_discriminator(x): 120 | is_training = tf.constant(True) 121 | with tf.variable_scope('local'): 122 | with tf.variable_scope('conv1'): 123 | x = conv_layer(x, [5, 5, 3, 64], 2) 124 | x = batch_normalize(x, is_training) 125 | x = tf.nn.relu(x) 126 | with tf.variable_scope('conv2'): 127 | x = conv_layer(x, [5, 5, 64, 128], 2) 128 | x = batch_normalize(x, is_training) 129 | x = tf.nn.relu(x) 130 | with tf.variable_scope('conv3'): 131 | x = conv_layer(x, [5, 5, 128, 256], 2) 132 | x = batch_normalize(x, is_training) 133 | x = tf.nn.relu(x) 134 | with tf.variable_scope('conv4'): 135 | x = conv_layer(x, [5, 5, 256, 512], 2) 136 | x = batch_normalize(x, is_training) 137 | x = tf.nn.relu(x) 138 | with tf.variable_scope('fc'): 139 | x = flatten_layer(x) 140 | x = full_connection_layer(x, 1024) 141 | return x 142 | 143 | with tf.variable_scope('discriminator', reuse=reuse): 144 | global_output = global_discriminator(global_x) 145 | local_output = local_discriminator(local_x) 146 | with tf.variable_scope('concatenation'): 147 | output = tf.concat((global_output, local_output), 1) 148 | output = full_connection_layer(output, 1) 149 | 150 | return output 151 | 152 | 153 | def calc_g_loss(self, x, completion): 154 | loss = tf.nn.l2_loss(x - completion) 155 | return tf.reduce_mean(loss) 156 | 157 | 158 | def calc_d_loss(self, real, fake): 159 | alpha = 4e-4 160 | d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=real, labels=tf.ones_like(real))) 161 | d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake, labels=tf.zeros_like(fake))) 162 | return tf.add(d_loss_real, d_loss_fake) * alpha 163 | 164 | -------------------------------------------------------------------------------- /train_mosaic.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | from PIL import Image, ImageFilter 4 | import tqdm 5 | from model_mosaic import Model 6 | import load 7 | 8 | IMAGE_SIZE = 128 9 | LOCAL_SIZE = 64 10 | HOLE_MIN = 24 11 | HOLE_MAX = 48 12 | MOSAIC_MIN = 8 #Minimum number of mosaic squares across image 13 | MOSAIC_MAX = 32 #Maximum number of mosaic squares across image 14 | MOSAIC_GAUSSIAN_P = 0.5 #represent images that have been compressed post-mosaic 15 | MOSAIC_GAUSSIAN_MIN = 0.2 16 | MOSAIC_GAUSSIAN_MAX = 1.2 17 | LEARNING_RATE = 1e-3 18 | BATCH_SIZE = 16 19 | PRETRAIN_EPOCH = 100 20 | 21 | def train(): 22 | x = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 23 | mosaic = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 24 | mask = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 1]) 25 | local_x = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 26 | global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3]) 27 | local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, LOCAL_SIZE, LOCAL_SIZE, 3]) 28 | is_training = tf.placeholder(tf.bool, []) 29 | 30 | model = Model(x, mosaic, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE) 31 | sess = tf.Session() 32 | global_step = tf.Variable(0, name='global_step', trainable=False) 33 | epoch = tf.Variable(0, name='epoch', trainable=False) 34 | 35 | opt = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE) 36 | g_train_op = opt.minimize(model.g_loss, global_step=global_step, var_list=model.g_variables) 37 | d_train_op = opt.minimize(model.d_loss, global_step=global_step, var_list=model.d_variables) 38 | 39 | init_op = tf.global_variables_initializer() 40 | sess.run(init_op) 41 | 42 | if tf.train.get_checkpoint_state('./models'): 43 | saver = tf.train.Saver() 44 | saver.restore(sess, './models/latest') 45 | 46 | x_train, x_test = load.load() 47 | x_train = np.array([a / 127.5 - 1 for a in x_train]) 48 | x_test = np.array([a / 127.5 - 1 for a in x_test]) 49 | 50 | step_num = int(len(x_train) / BATCH_SIZE) 51 | 52 | while True: 53 | sess.run(tf.assign(epoch, tf.add(epoch, 1))) 54 | print('epoch: {}'.format(sess.run(epoch))) 55 | 56 | np.random.shuffle(x_train) 57 | 58 | # Completion 59 | if sess.run(epoch) <= PRETRAIN_EPOCH: 60 | g_loss_value = 0 61 | for i in tqdm.tqdm(range(step_num)): 62 | x_batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] 63 | points_batch, mask_batch = get_points() 64 | mosaic_batch = get_mosaic(x_batch) 65 | 66 | _, g_loss = sess.run([g_train_op, model.g_loss], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: True}) 67 | g_loss_value += g_loss 68 | 69 | print('Completion loss: {}'.format(g_loss_value)) 70 | 71 | f = open("loss.csv","a+") 72 | f.write(str(sess.run(epoch)) + "," + str(g_loss_value) + "," + "0" + "\n") 73 | f.close() 74 | 75 | np.random.shuffle(x_test) 76 | x_batch = x_test[:BATCH_SIZE] 77 | mosaic_batch = get_mosaic(x_batch) 78 | merged, completion = sess.run([model.merged, model.completion], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: False}) 79 | sample = np.array((merged[0] + 1) * 127.5, dtype=np.uint8) 80 | result = Image.fromarray(sample) 81 | result.save('./training_output_images/{}_0.png'.format("{0:06d}".format(sess.run(epoch)))) 82 | sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) 83 | result = Image.fromarray(sample) 84 | result.save('./training_output_images/{}_1.png'.format("{0:06d}".format(sess.run(epoch)))) 85 | 86 | saver = tf.train.Saver() 87 | saver.save(sess, './models/latest', write_meta_graph=False) 88 | if sess.run(epoch) == PRETRAIN_EPOCH: 89 | saver.save(sess, './models/pretrained', write_meta_graph=False) 90 | 91 | 92 | # Discrimitation 93 | else: 94 | g_loss_value = 0 95 | d_loss_value = 0 96 | for i in tqdm.tqdm(range(step_num)): 97 | x_batch = x_train[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] 98 | points_batch, mask_batch = get_points() 99 | mosaic_batch = get_mosaic(x_batch) 100 | 101 | _, g_loss, completion = sess.run([g_train_op, model.g_loss, model.completion], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: True}) 102 | g_loss_value += g_loss 103 | 104 | local_x_batch = [] 105 | local_completion_batch = [] 106 | for i in range(BATCH_SIZE): 107 | x1, y1, x2, y2 = points_batch[i] 108 | local_x_batch.append(x_batch[i][y1:y2, x1:x2, :]) 109 | local_completion_batch.append(completion[i][y1:y2, x1:x2, :]) 110 | local_x_batch = np.array(local_x_batch) 111 | local_completion_batch = np.array(local_completion_batch) 112 | 113 | _, d_loss = sess.run( 114 | [d_train_op, model.d_loss], 115 | feed_dict={x: x_batch, mask: mask_batch, local_x: local_x_batch, global_completion: completion, local_completion: local_completion_batch, is_training: True}) 116 | d_loss_value += d_loss 117 | 118 | print('Completion loss: {}'.format(g_loss_value)) 119 | print('Discriminator loss: {}'.format(d_loss_value)) 120 | 121 | np.random.shuffle(x_test) 122 | x_batch = x_test[:BATCH_SIZE] 123 | mosaic_batch = get_mosaic(x_batch) 124 | merged, completion = sess.run([model.merged, model.completion], feed_dict={x: x_batch, mask: mask_batch, mosaic: mosaic_batch, is_training: False}) 125 | sample = np.array((merged[0] + 1) * 127.5, dtype=np.uint8) 126 | result = Image.fromarray(sample) 127 | result.save('./training_output_images/{}_0.png'.format("{0:06d}".format(sess.run(epoch)))) 128 | sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8) 129 | result = Image.fromarray(sample) 130 | result.save('./training_output_images/{}_1.png'.format("{0:06d}".format(sess.run(epoch)))) 131 | 132 | saver = tf.train.Saver() 133 | saver.save(sess, './models/latest', write_meta_graph=False) 134 | 135 | 136 | def get_points(): 137 | points = [] 138 | mask = [] 139 | for i in range(BATCH_SIZE): 140 | x1, y1 = np.random.randint(0, IMAGE_SIZE - LOCAL_SIZE + 1, 2) 141 | x2, y2 = np.array([x1, y1]) + LOCAL_SIZE 142 | points.append([x1, y1, x2, y2]) 143 | 144 | w, h = np.random.randint(HOLE_MIN, HOLE_MAX + 1, 2) 145 | p1 = x1 + np.random.randint(0, LOCAL_SIZE - w) 146 | q1 = y1 + np.random.randint(0, LOCAL_SIZE - h) 147 | p2 = p1 + w 148 | q2 = q1 + h 149 | 150 | m = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 1), dtype=np.uint8) 151 | m[q1:q2 + 1, p1:p2 + 1] = 1 152 | mask.append(m) 153 | 154 | return np.array(points), np.array(mask) 155 | 156 | 157 | def get_mosaic(x_batch): 158 | mosaic = [] 159 | for i in range(BATCH_SIZE): 160 | im = np.array((x_batch[i] + 1) * 127.5, dtype=np.uint8) 161 | im = Image.fromarray(im) 162 | size = np.random.randint(MOSAIC_MIN, MOSAIC_MAX) 163 | im = im.resize((size,size),Image.LANCZOS) 164 | im = im.resize((IMAGE_SIZE,IMAGE_SIZE),Image.NEAREST) 165 | if np.random.rand() < MOSAIC_GAUSSIAN_P: 166 | im = im.filter(ImageFilter.GaussianBlur(np.random.uniform(MOSAIC_GAUSSIAN_MIN, MOSAIC_GAUSSIAN_MAX))) 167 | 168 | mosaic.append(np.array(im)) 169 | 170 | mosaic = np.array([a / 127.5 - 1 for a in mosaic]) 171 | return mosaic 172 | 173 | 174 | if __name__ == '__main__': 175 | train() 176 | 177 | --------------------------------------------------------------------------------