def test_shape_bigger_image_same_aspect(self):
shape = (300, 400, 3)
h, w, _ = shape
in_img = Image.new('RGB', (600, 800))
out_img = resize_and_crop_image(image=in_img, size=(w, h))
self.assertSequenceEqual(out_img.shape,
shape,
msg='Resized when bigger - same aspect ratio')
def test_shape_smaller_image_diff_aspect(self):
shape = (300, 400, 3)
h, w, _ = shape
in_img = Image.new('RGB', (150, 175))
out_img = resize_and_crop_image(image=in_img, size=(w, h))
self.assertSequenceEqual(
out_img.shape,
shape,
msg='Resized when smaller - different aspect ratio')
def detect_img(img, infer_func):
cv.cvtColor(img, cv.COLOR_BGR2RGB, dst=img)
h, w, c = img.shape
out_size = (IMG_SIZE, IMG_SIZE)
img = tf.image.convert_image_dtype(img, tf.float32)
img, info = resize_and_crop_image(img, out_size, out_size)
img = tf.expand_dims(img, 0)
out = infer_func(img)['output_0']
nms_info, val_idx = nms_gpu_v2(out, args.iou_threshold, args.conf_thresh,
args.max_detct_num)
nms_info = nms_info[0][:val_idx[0]].numpy()
box, score, cls = np.split(nms_info, [4, 5], axis=-1)
# convert letterbox coord to origin img coord
w, h, c = img.shape
box = correct_box(box, w, h, IMG_SIZE)
cls = cls.astype(np.int)
return box, score, cls
def test_image_unchanged_when_default_size(self):
in_img = Image.new('RGB', (300, 400))
out_img = resize_and_crop_image(image=in_img, size=(300, 400))
self.assertTrue(np.array_equal(np.array(in_img), out_img),
msg='Unchanged when default size')
def main(args=None, argv=sys.argv):
tf.logging.info(args)
if args is None:
if FLAGS is not None:
args = FLAGS
else:
tf.logging.error('Error: Model requires args.')
sys.exit(1)
if not os.path.exists(args.output_img_dir):
os.makedirs(args.output_img_dir)
timestamp = int(time.time())
output_img_base = '{}-{}'.format(args.img_base_name, timestamp)
style_profile = list(zip(STYLE_LAYERS, args.style_weights))
tf.reset_default_graph()
sess = tf.InteractiveSession()
content_image = Image.open(args.content_image)
content_image = resize_and_crop_image(content_image)
content_image = reshape_and_normalize_image(content_image)
style_image = Image.open(args.style_image)
style_image = resize_and_crop_image(style_image)
style_image = reshape_and_normalize_image(style_image)
generated_image = generate_noise_image(content_image, args.noise_ratio)
model = load_vgg_model(args.pretrained_model)
sess.run(model['input'].assign(content_image))
out = model[args.output_layer]
a_C = sess.run(out)
a_G = out
sess.run(model['input'].assign(style_image))
J_content = compute_content_cost(a_C, a_G)
J_style = compute_style_cost(sess, model, style_profile)
J = total_cost(J_content, J_style, args.alpha, args.beta)
optimizer = tf.train.AdamOptimizer(args.learning_rate)
train_step = optimizer.minimize(J)
with tempfile.TemporaryDirectory() as tmpdir:
try:
transfer_style(
sess=sess,
model=model,
train_step=train_step,
J_content=J_content,
J_style=J_style,
J_total=J,
input_image=generated_image,
num_iterations=args.train_steps,
log_interval=args.log_interval,
output_img_dir=tmpdir,
output_img_base=output_img_base)
except OSError as err:
tf.logging.error(str(err))
finally:
img_dir = os.path.join(args.output_img_dir, output_img_base)
if args.drop_intermediate_images:
output_images = glob.glob(os.path.join(tmpdir, '*.png'))
latest_image = max(output_images, key=os.path.getctime)
os.makedirs(img_dir)
shutil.copy2(
os.path.join(tmpdir, latest_image),
os.path.join(img_dir))
else:
shutil.copytree(tmpdir, img_dir)
yaml_stream = open(os.path.join(img_dir, 'config.yaml'), 'w')
yaml.dump(args, yaml_stream)
tf.logging.info(
'Neural-style transfer complete. Image(s) transferred to {}'.format(
img_dir))