temporarily revert batch_size usage in decensor.py and fix image sample saving in train.py

This commit is contained in:
deeppomf 2018-02-27 23:10:44 -05:00
parent 74e9d08697
commit faca2f76fd
3 changed files with 20 additions and 16 deletions

View File

@ -18,17 +18,18 @@ parser.add_argument('--local_input_size', dest='local_input_size', default=64, h
parser.add_argument('--input_channel_size', dest='input_channel_size', default=3, help='input image channel')
parser.add_argument('--min_mask_size', dest='min_mask_size', default=24, help='minimum mask size')
parser.add_argument('--max_mask_size', dest='max_mask_size', default=48, help='maximum mask size')
parser.add_argument('--rotate_chance', dest='rotate_chance', default=0.5, help='chance the mask will be randomly rotated')
parser.add_argument('--rotate_chance', dest='rotate_chance', default=0.7, help='chance the mask will be randomly rotated')
parser.add_argument('--train_mosaic', dest ='train_mosaic', default=False, help='train neural network to decensor mosaics')
# parser.add_argument('--input_dim', dest='input_dim', default=100, help='input z size')
# #Training Settings
parser.add_argument('--continue_training', dest='continue_training', default=False, type=str2bool, help='flag to continue training')
parser.add_argument('--training_samples_path', dest='training_samples_path', default='./training_samples/', help='samples images generated during training path')
parser.add_argument('--batch_size', dest='batch_size', default=16, help='batch size')
# parser.add_argument('--data', dest='data', default='../ambientGAN_TF/data', help='cats image train path')
parser.add_argument('--batch_size', dest='batch_size', default=16, help='batch size')
# parser.add_argument('--train_step', dest='train_step', default=400, help='total number of train_step')
# parser.add_argument('--Tc', dest='Tc', default=100, help='Tc to train Completion Network')
# parser.add_argument('--Td', dest='Td', default=1, help='Td to train Discriminator Network')
@ -49,7 +50,9 @@ parser.add_argument('--learning_rate', dest='learning_rate', default=0.001, help
# parser.add_argument('--checkpoints_path', dest='checkpoints_path', default='./checkpoints/', help='saved model checkpoint path')
# parser.add_argument('--graph_path', dest='graph_path', default='./graphs/', help='tensorboard graph')
# parser.add_argument('--images_path', dest='images_path', default='./images/', help='result images path')
parser.add_argument('--training_samples_path', dest='training_samples_path', default='./training_samples/', help='samples images generated during training path')
parser.add_argument('--testing_output_path', dest='testing_output_path', default='./testing_output/', help='output images generated from running test.py path')
parser.add_argument('--decensor_input_path', dest='decensor_input_path', default='./decensor_input/', help='input images to be decensored by decensor.py path')
parser.add_argument('--decensor_output_path', dest='decensor_output_path', default='./decensor_output/', help='output images generated from running decensor.py path')
args = parser.parse_args()

View File

@ -13,16 +13,15 @@ from config import *
BATCH_SIZE = 1
image_folder = 'decensor_input_images/'
mask_color = [0, 255, 0]
poisson_blending_enabled = False
def decensor(args):
x = tf.placeholder(tf.float32, [args.batch_size, args.image_size, args.image_size, args.input_channel_size])
mask = tf.placeholder(tf.float32, [args.batch_size, args.image_size, args.image_size, 1])
local_x = tf.placeholder(tf.float32, [args.batch_size, args.local_image_size, args.local_image_size, args.input_channel_size])
global_completion = tf.placeholder(tf.float32, [args.batch_size, args.image_size, args.image_size, args.input_channel_size])
local_completion = tf.placeholder(tf.float32, [args.batch_size, args.local_image_size, args.local_image_size, args.input_channel_size])
x = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, args.input_channel_size])
mask = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, 1])
local_x = tf.placeholder(tf.float32, [BATCH_SIZE, args.local_input_size, args.local_input_size, args.input_channel_size])
global_completion = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, args.input_channel_size])
local_completion = tf.placeholder(tf.float32, [BATCH_SIZE, args.local_input_size, args.local_input_size, args.input_channel_size])
is_training = tf.placeholder(tf.bool, [])
model = Model(x, mask, local_x, global_completion, local_completion, is_training, batch_size=BATCH_SIZE)
@ -35,7 +34,7 @@ def decensor(args):
x_decensor = []
mask_decensor = []
for subdir, dirs, files in sorted(os.walk(image_folder)):
for subdir, dirs, files in sorted(os.walk(args.decensor_input_path)):
for file in sorted(files):
file_path = os.path.join(subdir, file)
if os.path.isfile(file_path) and os.path.splitext(file_path)[1] == ".png":
@ -54,7 +53,6 @@ def decensor(args):
mask_batch = get_mask(x_batch)
completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False})
for i in range(BATCH_SIZE):
cnt += 1
img = completion[i]
img = np.array((img + 1) * 127.5, dtype=np.uint8)
original = x_batch[i]
@ -62,8 +60,9 @@ def decensor(args):
if (poisson_blending_enabled):
img = blend(original, img, mask_batch[0,:,:,0])
output = Image.fromarray(img.astype('uint8'), 'RGB')
dst = './decensor_output_images/{}.png'.format("{0:06d}".format(cnt))
dst = args.decensor_output_path + '{}.png'.format("{0:06d}".format(cnt))
output.save(dst)
cnt += 1
def get_mask(x_batch):
points = []
@ -71,14 +70,16 @@ def get_mask(x_batch):
for i in range(BATCH_SIZE):
raw = x_batch[i]
raw = np.array((raw + 1) * 127.5, dtype=np.uint8)
m = np.zeros((args.image_size, args.image_size, 1), dtype=np.uint8)
for x in range(args.image_size):
for y in range(args.image_size):
m = np.zeros((args.input_size, args.input_size, 1), dtype=np.uint8)
for x in range(args.input_size):
for y in range(args.input_size):
if np.array_equal(raw[x][y], mask_color):
m[x, y] = 1
mask.append(m)
return np.array(mask)
if __name__ == '__main__':
if not os.path.exists(args.decensor_output_path):
os.makedirs(args.decensor_output_path)
decensor(args)

View File

@ -68,7 +68,7 @@ def train(args):
completion = sess.run(model.completion, feed_dict={x: x_batch, mask: mask_batch, is_training: False})
sample = np.array((completion[0] + 1) * 127.5, dtype=np.uint8)
result = Image.fromarray(sample)
result.save('./training_output_images/{}.jpg'.format("{0:06d}".format(sess.run(epoch))))
result.save(args.training_samples_path + '/{}.jpg'.format("{0:06d}".format(sess.run(epoch))))
saver = tf.train.Saver()
saver.save(sess, './models/latest', write_meta_graph=False)