def post_params():
myid = str(int(time.time() * 10000)) + ".jpg"
tmp_img_post_path = RESULT_IMG_PATH + myid
style_img_url = STYLE_IMG_PATH + request.form['style_img_url']
content_img_url = CONTENT_IMG_PATH + request.form['content_img_url']
alpha = float(request.form['alpha'])
content_img = get_img(content_img_url)
style_size = int(float(request.form['style_scale']) * 512)
keep_colors = False
# style_img = get_img_crop(style_img_url, resize=style_size)
style_img = get_img_crop(style_img_url)
if style_size > 0:
style_img = resize_to(style_img, style_size)
if keep_colors:
style_img = preserve_colors_np(style_img, content_img)
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_img, style_img, alpha, False, 0.6,
False)
save_img(tmp_img_post_path, stylized_rgb)
return Response(myid, status=200, mimetype='application/json')
def stylize_output(wct_model, content_img, styles):
if args.crop_size > 0:
styles = [center_crop(style) for style in styles]
if args.keep_colors:
styles = [preserve_colors_np(style, content_img) for style in styles]
# Run the frame through the style network
stylized = content_img
for _ in range(args.passes):
stylized = wct_model.predict(stylized, styles)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(
styles[0], (stylized.shape[0], stylized.shape[0]))
stylized = np.hstack([style_img_resized, stylized])
return stylized
def main():
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Create needed dirs
in_dir = os.path.join(args.tmp_dir, 'input')
out_dir = os.path.join(args.tmp_dir, 'sytlized')
if not os.path.exists(in_dir):
os.makedirs(in_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
in_args = ['ffmpeg', '-i', args.in_path, '%s/frame_%%d.png' % in_dir]
subprocess.call(" ".join(in_args), shell=True)
base_names = os.listdir(in_dir)
in_files = [os.path.join(in_dir, x) for x in base_names]
out_files = [os.path.join(out_dir, x) for x in base_names]
style_img = get_img(args.style_path)
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
s = time.time()
for in_f, out_f in zip(in_files, out_files):
print('{} -> {}'.format(in_f, out_f))
content_img = get_img(in_f)
if args.keep_colors:
style_rgb = preserve_colors_np(style_img, content_img)
else:
style_rgb = style_img
stylized = wct_model.predict(content_img, style_rgb, args.alpha,
args.swap5, args.ss_alpha)
if args.passes > 1:
for _ in range(args.passes - 1):
stylized = wct_model.predict(stylized, style_rgb, args.alpha)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(
style_rgb, (stylized.shape[0], stylized.shape[0]))
stylized = np.hstack([style_img_resized, stylized])
save_img(out_f, stylized)
fr = 30
out_args = [
'ffmpeg', '-i',
'%s/frame_%%d.png' % out_dir, '-f', 'mp4', '-q:v', '0', '-vcodec',
'mpeg4', '-r',
str(fr), args.out_path
]
subprocess.call(" ".join(out_args), shell=True)
print('Video at: %s' % args.out_path)
if args.keep_tmp is False:
shutil.rmtree(args.tmp_dir)
print('Processed in:', (time.time() - s))
def main():
start = time.time()
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Get content & style full paths
if os.path.isdir(args.content_path):
content_files = get_files(args.content_path)
else: # Single image file
content_files = [args.content_path]
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
if args.random > 0:
style_files = np.random.choice(style_files, args.random)
else: # Single image file
style_files = [args.style_path]
os.makedirs(args.out_path, exist_ok=True)
count = 0
# Apply each style to each content image
for content_fullpath in content_files:
content_prefix, content_ext = os.path.splitext(content_fullpath)
content_prefix = os.path.basename(
content_prefix) # Extract filename prefix without ext
content_img = get_img(content_fullpath)
if args.content_size > 0:
content_img = resize_to(content_img, args.content_size)
for style_fullpath in style_files:
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(
style_prefix) # Extract filename prefix without ext
# style_img = get_img_crop(style_fullpath, resize=args.style_size, crop=args.crop_size)
# style_img = resize_to(get_img(style_fullpath), content_img.shape[0])
style_img = get_img(style_fullpath)
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
if args.keep_colors:
style_img = preserve_colors_np(style_img, content_img)
# if args.noise: # Generate textures from noise instead of images
# frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
# frame_resize = gaussian_filter(frame_resize, sigma=0.5)
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_img, style_img,
args.alpha, args.swap5,
args.ss_alpha, args.adain)
if args.passes > 1:
for _ in range(args.passes - 1):
stylized_rgb = wct_model.predict(stylized_rgb, style_img,
args.alpha, args.swap5,
args.ss_alpha, args.adain)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(
style_img, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
# margin = np.ones((style_img_resized.shape[0], 10, 3)) * 255
stylized_rgb = np.hstack([style_img_resized, stylized_rgb])
# Format for out filename: {out_path}/{content_prefix}_{style_prefix}.{content_ext}
out_f = os.path.join(
args.out_path, '{}_{}{}'.format(content_prefix, style_prefix,
content_ext))
# out_f = f'{content_prefix}_{style_prefix}.{content_ext}'
save_img(out_f, stylized_rgb)
count += 1
print("{}: Wrote stylized output image to {}".format(count, out_f))
print("Finished stylizing {} outputs in {}s".format(
count,
time.time() - start))
def main():
# Load the AdaIN model
ada_in = AdaINference(args.checkpoint, args.vgg_path, device=args.device)
# Load a panel to control style settings
style_window = StyleWindow(args.style_path, args.style_size, args.scale, args.alpha, args.interpolate)
# Start the webcam stream
cap = WebcamVideoStream(args.video_source, args.width, args.height).start()
_, frame = cap.read()
# Grab a sample frame to calculate frame size
frame_resize = cv2.resize(frame, None, fx=args.scale, fy=args.scale)
img_shape = frame_resize.shape
# Setup video out writer
if args.video_out is not None:
fourcc = cv2.VideoWriter_fourcc(*'XVID')
if args.concat:
out_shape = (img_shape[1]+img_shape[0],img_shape[0]) # Make room for the style img
else:
out_shape = (img_shape[1],img_shape[0])
print('Video Out Shape:', out_shape)
video_writer = cv2.VideoWriter(args.video_out, fourcc, args.fps, out_shape)
fps = FPS().start() # Track FPS processing speed
keep_colors = args.keep_colors
count = 0
while(True):
ret, frame = cap.read()
if ret is True:
frame_resize = cv2.resize(frame, None, fx=style_window.scale, fy=style_window.scale)
if args.noise: # Generate textures from noise instead of images
frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
frame_resize = gaussian_filter(frame_resize, sigma=0.5)
count += 1
print("Frame:",count,"Orig shape:",frame.shape,"New shape",frame_resize.shape)
content_rgb = cv2.cvtColor(frame_resize, cv2.COLOR_BGR2RGB) # OpenCV uses BGR, we need RGB
if args.random > 0 and count % args.random == 0:
style_window.set_style(random=True, style_idx=0)
if keep_colors:
style_rgb = preserve_colors_np(style_window.style_rgbs[0], content_rgb)
else:
style_rgb = style_window.style_rgbs[0]
if args.interpolate is False:
# Run the frame through the style network
stylized_rgb = ada_in.predict(content_rgb, style_rgb, style_window.alpha)
else:
interp_weights = [style_window.interp_weight, 1 - style_window.interp_weight]
stylized_rgb = ada_in.predict_interpolate(content_rgb,
style_window.style_rgbs,
interp_weights,
style_window.alpha)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat and args.interpolate is False:
# Resize style img to same height as frame
style_rgb_resized = cv2.resize(style_rgb, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
stylized_rgb = np.hstack([style_rgb_resized, stylized_rgb])
stylized_bgr = cv2.cvtColor(stylized_rgb, cv2.COLOR_RGB2BGR)
if args.video_out is not None:
stylized_bgr = cv2.resize(stylized_bgr, out_shape) # Make sure frame matches video size
video_writer.write(stylized_bgr)
cv2.imshow('AdaIN Style', stylized_bgr)
fps.update()
key = cv2.waitKey(10)
if key & 0xFF == ord('r'): # Load new random style
style_window.set_style(random=True, style_idx=0)
if args.interpolate: # Load a a second style if interpolating
style_window.set_style(random=True, style_idx=1, window='style2')
elif key & 0xFF == ord('c'):
keep_colors = not keep_colors
print("Switching to keep_colors",keep_colors)
elif key & 0xFF == ord('q'): # Quit
break
else:
break
fps.stop()
print('[INFO] elapsed time (total): {:.2f}'.format(fps.elapsed()))
print('[INFO] approx. FPS: {:.2f}'.format(fps.fps()))
cap.stop()
if args.video_out is not None:
video_writer.release()
cv2.destroyAllWindows()
def main():
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Create needed dirs
in_dir = os.path.join(args.tmp_dir, 'input')
out_dir = os.path.join(args.tmp_dir, 'sytlized')
if not os.path.exists(in_dir):
os.makedirs(in_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
if os.path.isdir(args.in_path):
in_path = get_files(args.in_path)
else: # Single image file
in_path = [args.in_path]
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
else: # Single image file
style_files = [args.style_path]
print(style_files)
import time
# time.sleep(999)
in_args = [
'ffmpeg',
'-i', args.in_path,
'%s/frame_%%d.png' % in_dir
]
subprocess.call(" ".join(in_args), shell=True)
base_names = os.listdir(in_dir)
in_files = [os.path.join(in_dir, x) for x in base_names]
out_files = [os.path.join(out_dir, x) for x in base_names]
s = time.time()
for content_fullpath in in_path:
content_prefix, content_ext = os.path.splitext(content_fullpath)
content_prefix = os.path.basename(content_prefix)
try:
for style_fullpath in style_files:
style_img = get_img(style_fullpath)
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(style_prefix)
# print("ARRAY: ", style_img)
out_v = os.path.join(args.out_path, '{}_{}{}'.format(content_prefix, style_prefix, content_ext))
print("OUT:",out_v)
if os.path.isfile(out_v):
print("SKIP" , out_v)
continue
for in_f, out_f in zip(in_files, out_files):
print('{} -> {}'.format(in_f, out_f))
content_img = get_img(in_f)
if args.keep_colors:
style_rgb = preserve_colors_np(style_img, content_img)
else:
style_rgb = style_img
stylized = wct_model.predict(content_img, style_rgb, args.alpha, args.swap5, args.ss_alpha)
if args.passes > 1:
for _ in range(args.passes-1):
stylized = wct_model.predict(stylized, style_rgb, args.alpha)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(style_rgb, (stylized.shape[0], stylized.shape[0]))
stylized = np.hstack([style_img_resized, stylized])
save_img(out_f, stylized)
fr = 30
out_args = [
'ffmpeg',
'-i', '%s/frame_%%d.png' % out_dir,
'-f', 'mp4',
'-q:v', '0',
'-vcodec', 'mpeg4',
'-r', str(fr),
'"' + out_v + '"'
]
print(out_args)
subprocess.call(" ".join(out_args), shell=True)
print('Video at: %s' % out_v)
if args.keep_tmp is True or len(style_files) > 1:
continue
else:
shutil.rmtree(args.tmp_dir)
print('Processed in:',(time.time() - s))
print('Processed in:',(time.time() - s))
except Exception as e:
print("EXCEPTION: ",e)
def main():
start = time.time()
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Get content & style full paths
if os.path.isdir(args.content_path):
content_files = get_files(args.content_path)
else: # Single image file
content_files = [args.content_path]
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
if args.random > 0:
style_files = np.random.choice(style_files, args.random)
else: # Single image file
style_files = [args.style_path]
os.makedirs(args.out_path, exist_ok=True)
count = 0
### Apply each style to each content image
for content_fullpath in content_files:
content_prefix, content_ext = os.path.splitext(content_fullpath)
content_prefix = os.path.basename(content_prefix) # Extract filename prefix without ext
content_img = get_img(content_fullpath)
if args.content_size > 0:
content_img = resize_to(content_img, args.content_size)
for style_fullpath in style_files:
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(style_prefix) # Extract filename prefix without ext
# style_img = get_img_crop(style_fullpath, resize=args.style_size, crop=args.crop_size)
# style_img = resize_to(get_img(style_fullpath), content_img.shape[0])
style_img = get_img(style_fullpath)
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
if args.keep_colors:
style_img = preserve_colors_np(style_img, content_img)
# if args.noise: # Generate textures from noise instead of images
# frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
# frame_resize = gaussian_filter(frame_resize, sigma=0.5)
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_img, style_img, args.alpha, args.swap5, args.ss_alpha, args.adain)
if args.passes > 1:
for _ in range(args.passes-1):
stylized_rgb = wct_model.predict(stylized_rgb, style_img, args.alpha, args.swap5, args.ss_alpha, args.adain)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(style_img, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
# margin = np.ones((style_img_resized.shape[0], 10, 3)) * 255
stylized_rgb = np.hstack([style_img_resized, stylized_rgb])
# Format for out filename: {out_path}/{content_prefix}_{style_prefix}.{content_ext}
out_f = os.path.join(args.out_path, '{}_{}{}'.format(content_prefix, style_prefix, content_ext))
# out_f = f'{content_prefix}_{style_prefix}.{content_ext}'
save_img(out_f, stylized_rgb)
count += 1
print("{}: Wrote stylized output image to {}".format(count, out_f))
print("Finished stylizing {} outputs in {}s".format(count, time.time() - start))
def main():
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device)
# Load a panel to control style settings
style_window = StyleWindow(args.style_path, args.style_size,
args.crop_size, args.scale, args.alpha)
# Start the webcam stream
cap = WebcamVideoStream(args.video_source, args.width, args.height).start()
_, frame = cap.read()
# Grab a sample frame to calculate frame size
frame_resize = cv2.resize(frame, None, fx=args.scale, fy=args.scale)
img_shape = frame_resize.shape
# Setup video out writer
if args.video_out is not None:
fourcc = cv2.VideoWriter_fourcc(*'XVID')
if args.concat:
out_shape = (img_shape[1] + img_shape[0], img_shape[0]
) # Make room for the style img
else:
out_shape = (img_shape[1], img_shape[0])
print('Video Out Shape:', out_shape)
video_writer = cv2.VideoWriter(args.video_out, fourcc, args.fps,
out_shape)
fps = FPS().start() # Track FPS processing speed
keep_colors = args.keep_colors
count = 0
while (True):
ret, frame = cap.read()
if ret is True:
frame_resize = cv2.resize(frame,
None,
fx=style_window.scale,
fy=style_window.scale)
if args.noise: # Generate textures from noise instead of images
frame_resize = np.random.randint(0, 256, frame_resize.shape,
np.uint8)
count += 1
print("Frame:", count, "Orig shape:", frame.shape, "New shape",
frame_resize.shape)
content_rgb = cv2.cvtColor(
frame_resize,
cv2.COLOR_BGR2RGB) # OpenCV uses BGR, we need RGB
if args.random > 0 and count % args.random == 0:
style_window.set_style(random=True)
if keep_colors:
style_rgb = preserve_colors_np(style_window.style_rgb,
content_rgb)
else:
style_rgb = style_window.style_rgb
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_rgb, style_rgb,
style_window.alpha)
if args.passes > 1:
for i in range(args.passes - 1):
stylized_rgb = wct_model.predict(stylized_rgb, style_rgb,
style_window.alpha)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_rgb_resized = cv2.resize(
style_rgb, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
stylized_rgb = np.hstack([style_rgb_resized, stylized_rgb])
stylized_bgr = cv2.cvtColor(stylized_rgb, cv2.COLOR_RGB2BGR)
if args.video_out is not None:
stylized_bgr = cv2.resize(
stylized_bgr,
out_shape) # Make sure frame matches video size
video_writer.write(stylized_bgr)
cv2.imshow('WCT Universal Style Transfer', stylized_bgr)
fps.update()
key = cv2.waitKey(10)
if key & 0xFF == ord('r'): # Load new random style
style_window.set_style(random=True)
elif key & 0xFF == ord('c'): # Toggle color preservation
keep_colors = not keep_colors
print('Switching to keep_colors', keep_colors)
elif key & 0xFF == ord('s'): # Save stylized frame
out_f = "{}.png".format(time.time())
save_img(out_f, stylized_rgb)
print('Saved image to', out_f)
elif key & 0xFF == ord('q'): # Quit gracefully
break
else:
break
fps.stop()
print('[INFO] elapsed time (total): {:.2f}'.format(fps.elapsed()))
print('[INFO] approx. FPS: {:.2f}'.format(fps.fps()))
cap.stop()
if args.video_out is not None:
video_writer.release()
cv2.destroyAllWindows()
def main():
# start = time.time()
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Get content & style full paths
if os.path.isdir(args.content_path):
content_files = get_files(args.content_path)
else: # Single image file
content_files = [args.content_path]
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
if args.random > 0:
style_files = np.random.choice(style_files, args.random)
else: # Single image file
style_files = [args.style_path]
os.makedirs(args.out_path, exist_ok=True)
thetotal = len(content_files) * len(style_files)
count = 1
### Apply each style to each content image
for content_fullpath in content_files:
content_prefix, content_ext = os.path.splitext(content_fullpath)
content_prefix = os.path.basename(
content_prefix) # Extract filename prefix without ext
content_img = get_img(content_fullpath)
if args.content_size > 0:
content_img = resize_to(content_img, args.content_size)
for style_fullpath in style_files:
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(
style_prefix) # Extract filename prefix without ext
if args.mask:
mask_prefix_, _ = os.path.splitext(args.mask)
mask_prefix = os.path.basename(mask_prefix_)
if args.keep_colors:
style_prefix = "KPT_" + style_prefix
if args.concat:
style_prefix = "CON_" + style_prefix
if args.adain:
style_prefix = "ADA_" + style_prefix
if args.swap5:
style_prefix = "SWP_" + style_prefix
if args.mask:
style_prefix = "MSK_" + mask_prefix + '_' + style_prefix
if args.remaster:
style_prefix = style_prefix + '_REMASTERED'
out_f = os.path.join(
args.out_path, '{}_{}{}'.format(content_prefix, style_prefix,
content_ext))
if os.path.isfile(out_f):
print("SKIP", out_f)
count += 1
continue
# style_img = get_img_crop(style_fullpath, resize=args.style_size, crop=args.crop_size)
# style_img = resize_to(get_img(style_fullpath), content_img.shape[0])
style_img = get_img(style_fullpath)
if style_img == ("IMAGE IS BROKEN"):
continue
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
if args.keep_colors:
style_img = preserve_colors_np(style_img, content_img)
# if args.noise: # Generate textures from noise instead of images
# frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
# frame_resize = gaussian_filter(frame_resize, sigma=0.5)
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_img, style_img,
args.alpha, args.swap5,
args.ss_alpha, args.adain)
if args.passes > 1:
for _ in range(args.passes - 1):
stylized_rgb = wct_model.predict(stylized_rgb, style_img,
args.alpha, args.swap5,
args.ss_alpha, args.adain)
if args.mask:
import cv2
cv2.imwrite('./tmp.png', stylized_rgb)
stylized_rgb = cv2.imread('./tmp.png')
mask = cv2.imread(args.mask, cv2.IMREAD_GRAYSCALE)
# from scipy.misc import bytescale
# mask = bytescale(mask)
# mask = scipy.ndimage.imread(args.mask,flatten=True,mode='L')
height, width = stylized_rgb.shape[:2]
# print(height, width)
# Resize the mask to fit the image.
mask = scipy.misc.imresize(mask, (height, width),
interp='bilinear')
stylized_rgb = cv2.bitwise_and(stylized_rgb,
stylized_rgb,
mask=mask)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
# style_prefix = style_prefix + "CON_"
# content_img_resized = scipy.misc.imresize(content_img, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
style_img_resized = scipy.misc.imresize(
style_img, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
# margin = np.ones((style_img_resized.shape[0], 10, 3)) * 255
stylized_rgb = np.hstack([style_img_resized, stylized_rgb])
# Format for out filename: {out_path}/{content_prefix}_{style_prefix}.{content_ext}
# out_f = f'{content_prefix}_{style_prefix}.{content_ext}'
out_f = os.path.join(
args.out_path, '{}_{}{}'.format(content_prefix, style_prefix,
content_ext))
if args.remaster:
# outf = os.path.join(args.out_path, '{}_{}_REMASTERED{}'.format(content_prefix, style_prefix, content_ext))
stylized_rgb = remaster_pic(stylized_rgb)
save_img(out_f, stylized_rgb)
totalfiles = len([
name for name in os.listdir(args.out_path)
if os.path.isfile(os.path.join(args.out_path, name))
])
# percent = math.floor(float(totalfiles/thetotal))
print("{}/{} TOTAL FILES".format(count, thetotal))
count += 1
print("{}: Wrote stylized output image to {}".format(count, out_f))
def main():
start = time.time()
session_conf = tf.ConfigProto(
allow_soft_placement=args.allow_soft_placement,
log_device_placement=args.log_device_placement)
session_conf.gpu_options.per_process_gpu_memory_fraction = args.gpu_fraction
with tf.Graph().as_default():
# with tf.Session(config=session_conf) as sess:
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
with wct_model.sess as sess:
# 训练的时候不需要style_img，所以在inference的时候重新保存一次checkpoint
saver = tf.train.Saver()
log_path = args.log_path if args.log_path is not None else os.path.join(
args.checkpoint, 'log')
summary_writer = tf.summary.FileWriter(log_path, sess.graph)
# 部分node没有保存在checkpoint中，需要重新初始化
# sess.run(tf.global_variables_initializer())
# Get content & style full paths
if os.path.isdir(args.content_path):
content_files = get_files(args.content_path)
else: # Single image file
content_files = [args.content_path]
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
if args.random > 0:
style_files = np.random.choice(style_files, args.random)
else: # Single image file
style_files = [args.style_path]
os.makedirs(args.out_path, exist_ok=True)
count = 0
# Apply each style to each content image
for content_fullpath in content_files:
content_prefix, content_ext = os.path.splitext(
content_fullpath)
content_prefix = os.path.basename(
content_prefix) # Extract filename prefix without ext
content_img = get_img(content_fullpath)
if args.content_size > 0:
content_img = resize_to(content_img, args.content_size)
for style_fullpath in style_files:
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(
style_prefix) # Extract filename prefix without ext
# style_img = get_img_crop(style_fullpath, resize=args.style_size, crop=args.crop_size)
# style_img = resize_to(get_img(style_fullpath), content_img.shape[0])
style_img = get_img(style_fullpath)
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
if args.keep_colors:
style_img = preserve_colors_np(style_img, content_img)
# if args.noise: # Generate textures from noise instead of images
# frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
# frame_resize = gaussian_filter(frame_resize, sigma=0.5)
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_img, style_img,
args.alpha, args.swap5,
args.ss_alpha, args.adain)
if args.passes > 1:
for _ in range(args.passes - 1):
stylized_rgb = wct_model.predict(
stylized_rgb, style_img, args.alpha,
args.swap5, args.ss_alpha, args.adain)
save_path = saver.save(
sess, os.path.join(args.checkpoint, 'model.ckpt'))
print("Model saved in file: %s" % save_path)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(
style_img,
(stylized_rgb.shape[0], stylized_rgb.shape[0]))
# margin = np.ones((style_img_resized.shape[0], 10, 3)) * 255
stylized_rgb = np.hstack(
[style_img_resized, stylized_rgb])
# Format for out filename: {out_path}/{content_prefix}_{style_prefix}.{content_ext}
out_f = os.path.join(
args.out_path,
'{}_{}{}'.format(content_prefix, style_prefix,
content_ext))
# out_f = f'{content_prefix}_{style_prefix}.{content_ext}'
# print(stylized_rgb, stylized_rgb.shape, type(stylized_rgb))
# print(out_f)
save_img(out_f, stylized_rgb)
count += 1
print("{}: Wrote stylized output image to {} at {}".format(
count, out_f, time.time()))
print("breaking...")
break
break
print("Finished stylizing {} outputs in {}s".format(
count,
time.time() - start))
def main():
start = time.time()
# Load the WCT model
wct_model = WCT(input_checkpoint=args.input_checkpoint,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
print("model construct end !")
# Get content & style full paths
if os.path.isdir(args.content_path):
content_files = get_files(args.content_path)
else: # Single image file
content_files = [args.content_path]
content_seg_files = list(map(lambda x: x.replace("2017", "2017_seg").replace("jpg", "png"), content_files))
assert reduce(lambda a, b : a + b, map(lambda x: int(os.path.exists(x)),content_seg_files + content_files)) > 0
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
if args.random > 0:
style_files = np.random.choice(style_files, args.random)
else: # Single image file
style_files = [args.style_path]
style_seg_files = list(map(lambda x: x.replace("2017", "2017_seg").replace("jpg", "png"), style_files))
assert reduce(lambda a, b : a + b, map(lambda x: int(os.path.exists(x)),style_seg_files + style_files)) > 0
os.makedirs(args.out_path, exist_ok=True)
count = 0
### Apply each style to each content image
for i in range(len(content_files)):
content_fullpath = content_files[i]
content_prefix, content_ext = os.path.splitext(content_fullpath)
content_prefix = os.path.basename(content_prefix) # Extract filename prefix without ext
content_img = get_img(content_fullpath)
if args.content_size > 0:
content_img = np.asarray(Image.fromarray(content_img.astype(np.uint8)).resize((args.content_size, args.content_size))).astype(np.float32)
content_seg = get_img_ori(content_seg_files[i])
if args.content_size > 0:
content_seg = np.asarray(Image.fromarray(content_seg.astype(np.uint8)).resize((args.content_size, args.content_size))).astype(np.float32)
content_seg = image_to_pixel_map(content_seg)[...,np.newaxis]
content_img = np.concatenate([content_img, content_seg], axis=-1)
for i in range(len(style_files)):
style_fullpath = style_files[i]
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(style_prefix) # Extract filename prefix without ext
style_img = get_img(style_fullpath)
style_seg = get_img_ori(style_seg_files[i])
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
style_seg = resize_to(style_seg, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
style_seg = center_crop(style_seg, args.crop_size)
style_seg = image_to_pixel_map(style_seg)[...,np.newaxis]
style_img = np.concatenate([style_img, style_seg], axis=-1)
assert not args.keep_colors
if args.keep_colors:
style_img = preserve_colors_np(style_img, content_img)
# if args.noise: # Generate textures from noise instead of images
# frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
# frame_resize = gaussian_filter(frame_resize, sigma=0.5)
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_img, style_img, args.alpha, args.swap5, args.ss_alpha, args.adain)
if args.passes > 1:
for _ in range(args.passes-1):
stylized_rgb = np.concatenate([stylized_rgb ,content_seg], axis=-1)
stylized_rgb = wct_model.predict(stylized_rgb, style_img, args.alpha, args.swap5, args.ss_alpha, args.adain)
# Stitch the style + stylized output together, but only if there's one style image
assert not args.concat
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(style_img, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
# margin = np.ones((style_img_resized.shape[0], 10, 3)) * 255
stylized_rgb = np.hstack([style_img_resized, stylized_rgb])
####+++++++++++++++++++++++++++++++++++
# Format for out filename: {out_path}/{content_prefix}_{style_prefix}.{content_ext}
out_f = os.path.join(args.out_path, '{}_{}{}'.format(content_prefix, style_prefix, content_ext))
save_img(out_f, stylized_rgb)
count += 1
print("{}: Wrote stylized output image to {}".format(count, out_f))
print("Finished stylizing {} outputs in {}s".format(count, time.time() - start))
def main():
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Create needed dirs
in_dir = os.path.join(args.tmp_dir, 'input')
out_dir = os.path.join(args.tmp_dir, 'sytlized')
if not os.path.exists(in_dir):
os.makedirs(in_dir)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
if os.path.isdir(args.in_path):
in_path = get_files(args.in_path)
else: # Single image file
in_path = [args.in_path]
if os.path.isdir(args.style_path):
style_files = get_files(args.style_path)
else: # Single image file
style_files = [args.style_path]
print(style_files)
# time.sleep(999)
in_args = [
'ffmpeg',
'-i', args.in_path,
'%s/frame_%%d.png' % in_dir
]
subprocess.call(" ".join(in_args), shell=True)
base_names = os.listdir(in_dir)
in_files = [os.path.join(in_dir, x) for x in base_names]
out_files = [os.path.join(out_dir, x) for x in base_names]
s = time.time()
for content_fullpath in in_path:
content_prefix, content_ext = os.path.splitext(content_fullpath)
content_prefix = os.path.basename(content_prefix)
try:
for style_fullpath in style_files:
style_img = get_img(style_fullpath)
if args.style_size > 0:
style_img = resize_to(style_img, args.style_size)
if args.crop_size > 0:
style_img = center_crop(style_img, args.crop_size)
style_prefix, _ = os.path.splitext(style_fullpath)
style_prefix = os.path.basename(style_prefix)
# print("ARRAY: ", style_img)
out_v = os.path.join(args.out_path, '{}_{}{}'.format(content_prefix, style_prefix, content_ext))
print("OUT:", out_v)
if os.path.isfile(out_v):
print("SKIP", out_v)
continue
for in_f, out_f in zip(in_files, out_files):
print('{} -> {}'.format(in_f, out_f))
content_img = get_img(in_f)
if args.keep_colors:
style_rgb = preserve_colors_np(style_img, content_img)
else:
style_rgb = style_img
stylized = wct_model.predict(content_img, style_rgb, args.alpha, args.swap5, args.ss_alpha)
if args.passes > 1:
for _ in range(args.passes-1):
stylized = wct_model.predict(stylized, style_rgb, args.alpha)
# Stitch the style + stylized output together, but only if there's one style image
if args.concat:
# Resize style img to same height as frame
style_img_resized = scipy.misc.imresize(style_rgb, (stylized.shape[0], stylized.shape[0]))
stylized = np.hstack([style_img_resized, stylized])
save_img(out_f, stylized)
fr = 30
out_args = [
'ffmpeg',
'-i', '%s/frame_%%d.png' % out_dir,
'-f', 'mp4',
'-q:v', '0',
'-vcodec', 'mpeg4',
'-r', str(fr),
'"' + out_v + '"'
]
print(out_args)
subprocess.call(" ".join(out_args), shell=True)
print('Video at: %s' % out_v)
if args.keep_tmp is True or len(style_files) > 1:
continue
else:
shutil.rmtree(args.tmp_dir)
print('Processed in:', (time.time() - s))
print('Processed in:', (time.time() - s))
except Exception as e:
print("EXCEPTION: ", e)
def main():
# Load the WCT model
wct_model = WCT(checkpoints=args.checkpoints,
relu_targets=args.relu_targets,
vgg_path=args.vgg_path,
device=args.device,
ss_patch_size=args.ss_patch_size,
ss_stride=args.ss_stride)
# Load a panel to control style settings
style_window = StyleWindow(args.style_path,
args.style_size,
args.crop_size,
args.scale,
args.alpha,
args.swap5,
args.ss_alpha,
args.passes)
# Start the webcam stream
cap = WebcamVideoStream(args.video_source, args.width, args.height).start()
_, frame = cap.read()
# Grab a sample frame to calculate frame size
frame_resize = cv2.resize(frame, None, fx=args.scale, fy=args.scale)
img_shape = frame_resize.shape
# Setup video out writer
if args.video_out is not None:
fourcc = cv2.VideoWriter_fourcc(*'XVID')
if args.concat:
out_shape = (img_shape[1]+img_shape[0],img_shape[0]) # Make room for the style img
else:
out_shape = (img_shape[1],img_shape[0])
print('Video Out Shape:', out_shape)
video_writer = cv2.VideoWriter(args.video_out, fourcc, args.fps, out_shape)
fps = FPS().start() # Track FPS processing speed
# Toggles changed with kb shortcuts
keep_colors = args.keep_colors
swap_style = args.swap5
use_adain = args.adain
count = 0
while(True):
ret, frame = cap.read()
if ret is True:
frame_resize = cv2.resize(frame, None, fx=style_window.scale, fy=style_window.scale)
if args.noise: # Generate textures from noise instead of images
frame_resize = np.random.randint(0, 256, frame_resize.shape, np.uint8)
count += 1
print("Frame:",count,"Orig shape:",frame.shape,"New shape",frame_resize.shape)
content_rgb = cv2.cvtColor(frame_resize, cv2.COLOR_BGR2RGB) # OpenCV uses BGR, we need RGB
if args.random > 0 and count % args.random == 0:
style_window.set_style(random=True)
if keep_colors:
style_rgb = preserve_colors_np(style_window.style_rgb, content_rgb)
else:
style_rgb = style_window.style_rgb
# Run the frame through the style network
stylized_rgb = wct_model.predict(content_rgb, style_rgb, style_window.alpha, swap_style, style_window.ss_alpha, use_adain)
# Repeat stylization pipeline
if style_window.passes > 1:
for i in range(style_window.passes-1):
stylized_rgb = wct_model.predict(stylized_rgb, style_rgb, style_window.alpha, swap_style, style_window.ss_alpha, use_adain)
# Stitch the style + stylized output together
if args.concat:
# Resize style img to same height as frame
style_rgb_resized = cv2.resize(style_rgb, (stylized_rgb.shape[0], stylized_rgb.shape[0]))
stylized_rgb = np.hstack([style_rgb_resized, stylized_rgb])
stylized_bgr = cv2.cvtColor(stylized_rgb, cv2.COLOR_RGB2BGR)
if args.video_out is not None:
stylized_bgr = cv2.resize(stylized_bgr, out_shape) # Make sure frame matches video size
video_writer.write(stylized_bgr)
cv2.imshow('WCT Universal Style Transfer', stylized_bgr)
fps.update()
key = cv2.waitKey(10)
if key & 0xFF == ord('r'): # Load new random style
style_window.set_style(random=True)
elif key & 0xFF == ord('c'): # Toggle color preservation
keep_colors = not keep_colors
print('Switching to keep_colors:',keep_colors)
elif key & 0xFF == ord('s'): # Toggle style swap
swap_style = not swap_style
print('New value for flag swap_style:',swap_style)
elif key & 0xFF == ord('a'): # Toggle AdaIN
use_adain = not use_adain
print('New value for flag use_adain:',use_adain)
elif key & 0xFF == ord('w'): # Write stylized frame
out_f = "{}.png".format(time.time())
save_img(out_f, stylized_rgb)
print('Saved image to:',out_f)
elif key & 0xFF == ord('q'): # Quit gracefully
break
else:
break
fps.stop()
print('[INFO] elapsed time (total): {:.2f}'.format(fps.elapsed()))
print('[INFO] approx. FPS: {:.2f}'.format(fps.fps()))
cap.stop()
if args.video_out is not None:
video_writer.release()
cv2.destroyAllWindows()