-
Notifications
You must be signed in to change notification settings - Fork 2
/
style.py
167 lines (141 loc) · 7.09 KB
/
style.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
import torch
import cv2
import argparse
import os
import os.path as osp
from torchvision.transforms import transforms
import glob
# Cloned modules
import utils
import transformer
# To use this script just simply enter
# python style.py --input video.mp4 --style hokusai-wave.pth
# E.g python style.py -i data\Cannonbal.mp4 -s weights\hokusai-wave.pth
# GPU Memory alert! if you have bigger memory, increase batch size , smaller memory, decrease batch size by using -b, default is batch size 8
# Remove old folders because this script still cannot overwrite the exisiting folder
def stylize_folder(style_path, folder_containing_the_content_folder, save_folder, batch_size=1):
# This function is cloned and modified from the original repo: https://github.com/rrmina/fast-neural-style-pytorch
"""Stylizes images in a folder by batch
If the images are of different dimensions, use transform.resize() or use a batch size of 1
IMPORTANT: Put content_folder inside another folder folder_containing_the_content_folder
"""
# Choose GPU if available else use CPU
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Image loader
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
# Please take note that the original frames must be inside a subfolder inside of the folder that you provide the path to --input
# This is due to the nature of how pytorch's ImageFolder work
image_dataset = utils.ImageFolderWithPaths(folder_containing_the_content_folder, transform=transform)
image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=batch_size)
# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_path))
net = net.to(device)
# Stylize batches of images
with torch.no_grad():
for content_batch, _, path in image_loader:
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
generated_tensor = net(content_batch.to(device)).detach()
# Create folder if not exist
if not os.path.exists(save_folder):
os.makedirs(save_folder)
# Save images
for i in range(len(path)):
generated_image = utils.ttoi(generated_tensor[i])
image_name = os.path.basename(path[i])
utils.saveimg(generated_image, osp.join(save_folder, image_name))
def getFrames(video_path):
# This function is cloned and modified from the original repo: https://github.com/rrmina/fast-neural-style-pytorch
if not os.path.exists(ORIGINAL_FRAMES_FOLDER_NAME):
os.makedirs(osp.join(ORIGINAL_FRAMES_FOLDER_NAME))
if not os.path.exists(osp.join(ORIGINAL_FRAMES_FOLDER_NAME, "content")):
os.makedirs(osp.join(ORIGINAL_FRAMES_FOLDER_NAME, "content"))
cap = cv2.VideoCapture(video_path)
success, image = cap.read()
count = 1
success = True
# "content" is to make the subfolder that will be needed later
while success:
cv2.imwrite(osp.join(ORIGINAL_FRAMES_FOLDER_NAME, "content", "frame" + str(count) + ".jpg"), image)
success, image = cap.read()
count += 1
print("Done extracting all frames")
# Main function
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", type=str, default="None")
parser.add_argument("-s", "--style", type=str, default="None")
parser.add_argument("-o", "--output", type=str, default="output.mp4")
parser.add_argument("-b", "--batchsize", type=int, default=8)
args = parser.parse_args()
# need to chg the input from folder into video
print("Your input video: {}".format(args.input))
print("Your style image: {}".format(args.style))
ORIGINAL_FRAMES_FOLDER_NAME = "original_frames"
STYLED_FRAMES_FOLDER_NAME = "styled_frames"
# Information of your input video
cap = cv2.VideoCapture(args.input)
VIDEO_WIDTH = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
VIDEO_HEIGHT = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
FPS = int(cap.get(cv2.CAP_PROP_FPS))
getFrames(args.input)
# Prepare Part 2 by inference the original frames into the network to generate styled frames
stylize_folder(args.style, ORIGINAL_FRAMES_FOLDER_NAME, STYLED_FRAMES_FOLDER_NAME, batch_size=args.batchsize)
# This is hack-ish because I have not find a way to use glob to order the images in 1,2,3... instead of 1,10,11...
numberOfImgs_Part2 = len(glob.glob(osp.join(STYLED_FRAMES_FOLDER_NAME, "frame*.jpg")))
part2Frames = []
for i in range(1, numberOfImgs_Part2 + 1):
filename = osp.join(STYLED_FRAMES_FOLDER_NAME, "frame{}.jpg").format(i)
part2Frames.append(filename)
# Uncomment this to only produce Part 2 without the Part 1 and Part 3
# fourcc = cv2.VideoWriter_fourcc(*'MP4V')
# out = cv2.VideoWriter(args.output, fourcc, FPS, (int(VIDEO_WIDTH), int(VIDEO_HEIGHT)))
#
# # Write frames from Part2 the video
# for image_name in part2Frames:
# out.write(cv2.imread(image_name))
# out.release()
# print("Done creating the Part 2 of your styled video for your music!")
# Create Part1 that has a image blending effect transition
# Get the first frame from the "original_frames" folder and the the first frame from the "styled_frames" folder
first_ori = cv2.imread("original_frames/content/frame1.jpg")
first_styled = cv2.imread("styled_frames/frame1.jpg")
part1Frames = []
alpha = 0.0
for i in range(0, 100):
alpha = i / 100
beta = (1.0 - alpha)
output = cv2.addWeighted(first_styled, alpha, first_ori, beta, 0.0)
part1Frames.append(output)
# cv2.imwrite("part1_{}.jpg".format(i),output)
# Create Part3 that has an image blending effect transition but in an inverse order
# Get the last frame from the "original_frames" folder and the the last frame from the "styled_frames" folder
lastIndex = len(glob.glob(osp.join(ORIGINAL_FRAMES_FOLDER_NAME, "content/frame*.jpg")))
last_ori = cv2.imread("original_frames/content/frame" + str(lastIndex) + ".jpg")
last_styled = cv2.imread("styled_frames/frame" + str(lastIndex) + ".jpg")
part3Frames = []
alpha = 0.0
for i in range(0, 100):
alpha = i / 100
beta = (1.0 - alpha)
output = cv2.addWeighted(last_ori, alpha, last_styled, beta, 0.0)
part3Frames.append(output)
# cv2.imwrite("part3_{}.jpg".format(i), output)
# It is time to write our frames into a video!
# Part1 and Part3 frames are stored in memory, there are of course better ways to do this, but oh well ...
# Define the codec and create VideoWrite object
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
out = cv2.VideoWriter(args.output, fourcc, FPS, (int(VIDEO_WIDTH), int(VIDEO_HEIGHT)))
for i in (range(len(part1Frames))):
out.write(part1Frames[i])
for image_name in part2Frames:
out.write(cv2.imread(image_name))
for i in (range(len(part3Frames))):
out.write(part3Frames[i])
out.release()
print("Done creating your styled video for your music!")