def save_debug_image(tensor_orig, tensor_transformed, filename):
assert tensor_orig.size() == tensor_transformed.size()
result = Image.fromarray(
recover_image(tensor_transformed.cpu().numpy())[0])
orig = Image.fromarray(recover_image(tensor_orig.cpu().numpy())[0])
new_im = Image.new('RGB', (result.size[0] * 2 + 5, result.size[1]))
new_im.paste(orig, (0, 0))
new_im.paste(result, (result.size[0] + 5, 0))
new_im.save(filename)
def transform(model_file, image_file, target_path, size):
transformer = load_model(model_file)
img_var = load_and_preprocess(image_file, size)
img_output = transformer(img_var)
output_img = Image.fromarray(
recover_image(img_output.data.numpy())[0])
output_img.save(target_path)
BATCH_SIZE = 8
mask = cv2.imread(mask_path).astype(np.bool)
save_model_path = "./models/" + style_name + "_10000_unstable_vgg19.pth"
transformer.load_state_dict(torch.load(save_model_path))
transformer.eval()
batch = []
videogen = skvideo.io.FFmpegReader(data_path)
writer = skvideo.io.FFmpegWriter("/tmp2/vincentwu929/DIP_final/" +\
video_name + "_" + style_name + ".avi")
try:
with torch.no_grad():
for frame in tqdm(videogen.nextFrame()):
batch.append(
preprocess(Image.fromarray(frame[0:4000,
1000:5000])).unsqueeze(0))
if len(batch) == BATCH_SIZE:
for frame_out in recover_image(
transformer(torch.cat(batch, 0).cuda()).cpu().numpy()):
frame_out = cv2.resize(frame_out, (4000, 4000),
cv2.INTER_CUBIC)
out_img = frame.copy()
out_img[0:4000, 1000:5000] = frame_out * mask + frame[
0:4000, 1000:5000] * (~mask)
writer.writeFrame(out_img)
batch = []
except RuntimeError as e:
print(e)
pass
writer.close()
transforms.Resize(1024),
transforms.ToTensor(),
tensor_normalizer()
])
frame = frames[0][0:4000, 1000:5000]
img = Image.fromarray(frame)
img_tensor = transform(img).unsqueeze(0)
if torch.cuda.is_available():
img_tensor = img_tensor.cuda()
with torch.no_grad():
img_output = transformer(Variable(img_tensor))
profile_img = recover_image(img_output.data.cpu().numpy())[0]
profile_img = cv2.resize(profile_img, (4000, 4000), cv2.INTER_CUBIC)
if opts.with_mask:
out_img = frames[0]
out_img[0:4000, 1000:5000] = profile_img * mask + frames[0][
0:4000, 1000:5000] * (~mask)
output = Image.fromarray(out_img)
else:
output = Image.fromarray(profile_img)
if not os.path.exists('./fast-neural-style/results'):
os.makedirs('./fast-neural-style/results')
output.save(
'./fast-neural-style/results/frame_00000_{}_{}-stylized.jpg'.
format(opts.output_name, opts.style_name))
STYLE_IMAGE = "../style_images/mosaic.jpg"
style_img = Image.open(STYLE_IMAGE).convert('RGB')
with torch.no_grad():
style_img_tensor = transforms.Compose([
# transforms.Resize(IMAGE_SIZE* 2),
transforms.ToTensor(),
tensor_normalizer()
])(style_img).unsqueeze(0)
# assert np.sum(style_img - recover_image(style_img_tensor.numpy())[0].astype(np.uint8)) < 3 * style_img_tensor.size()[2] * style_img_tensor.size()[3]
style_img_tensor = style_img_tensor.to(device)
# Sanity check:
# In[7]:
plt.imshow(recover_image(style_img_tensor.cpu().numpy())[0])
# Precalculate gram matrices of the style image:
# In[8]:
# http://pytorch.org/docs/master/notes/autograd.html#volatile
with torch.no_grad():
style_loss_features = loss_network(style_img_tensor)
gram_style = [gram_matrix(y) for y in style_loss_features]
# In[9]:
style_loss_features._fields
# In[10]:
skvideo.io.ffprobe("videos/cat.mp4")
# In[2]:
transformer.load_state_dict(torch.load("../models/udine_10000.pth"))
# In[ ]:
frames = []
frames_orig = []
videogen = skvideo.io.vreader("videos/cat.mp4")
for frame in videogen:
frames_orig.append(Image.fromarray(frame))
frames.append(
recover_image(
transformer(Variable(preprocess(frame).unsqueeze(0),
volatile=True)).data.numpy())[0])
# In[ ]:
Image.fromarray(frames[3])
# In[ ]:
writer = skvideo.io.FFmpegWriter(
"cat.mp4") # tuple([len(frames)] + list(frames[0].shape)))
for frame in frames:
writer.writeFrame(frame)
writer.close()
# ## Higher Resolution Videos