def webcam(style_transform_path, width=1280, height=720):
"""
Captures and saves an image, perform style transfer, and again saves the styled image.
Reads the styled image and show in window.
Saving and loading SHOULD BE eliminated, however this produces too much whitening in
the "generated styled image". This may be caused by the async nature of VideoCapture,
and I don't know how to fix it.
"""
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
print("Loading Transformer Network")
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_transform_path))
net = net.to(device)
print("Done Loading Transformer Network")
# Set webcam settings
cam = cv2.VideoCapture(0)
cam.set(3, width)
cam.set(4, height)
# Main loop
with torch.no_grad():
count = 1
while True:
# Get webcam input
ret_val, img = cam.read()
# Mirror
img = cv2.flip(img, 1)
utils.saveimg(img, str(count) + ".png")
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
content_image = utils.load_image(str(count) + ".png")
content_tensor = utils.itot(content_image).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image,
generated_image)
utils.saveimg(generated_image, str(count + 1) + ".png")
img2 = cv2.imread(str(count + 1) + ".png")
count += 2
# Show webcam
cv2.imshow('Demo webcam', img2)
if cv2.waitKey(1) == 27:
break # esc to quit
# Free-up memories
cam.release()
cv2.destroyAllWindows()
def stylize_folder(style_path, folder_containing_the_content_folder, save_folder, batch_size=1):
"""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
folder_containing_the_content_folder
content_folder
pic1.ext
pic2.ext
pic3.ext
...
and saves as the styled images in save_folder as follow:
save_folder
pic1.ext
pic2.ext
pic3.ext
...
"""
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Image loader
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
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()
# Save images
for i in range(len(path)):
generated_image = utils.ttoi(generated_tensor[i])
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image, generated_image)
image_name = os.path.basename(path[i])
utils.saveimg(generated_image, save_folder + image_name)
#stylize()
def webcam(style_transform_path, width=1280, height=720):
"""
Captures and saves an image, perform style transfer, and again saves the styled image.
Reads the styled image and show in window.
"""
# Device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load Transformer Network
print("Loading Transformer Network")
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_transform_path))
net = net.to(device)
print("Done Loading Transformer Network")
# Set webcam settings
cam = cv2.VideoCapture(0)
cam.set(3, width)
cam.set(4, height)
# Main loop
with torch.no_grad():
while True:
# Get webcam input
ret_val, img = cam.read()
# Mirror
img = cv2.flip(img, 1)
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
content_tensor = utils.itot(img).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if PRESERVE_COLOR:
generated_image = utils.transfer_color(img, generated_image)
generated_image = generated_image / 255
# Show webcam
cv2.imshow("Demo webcam", generated_image)
if cv2.waitKey(1) == 27:
break # esc to quit
# Free-up memories
cam.release()
cv2.destroyAllWindows()
def webcam(style_transform_path, width=1280, height=720):
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
print("Loading Transformer Network")
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_transform_path, map_location=device))
net = net.to(device)
print("Done Loading Transformer Network")
# Set webcam settings
cam = cv2.VideoCapture(0)
cam.set(3, width)
cam.set(4, height)
# Main loop
with torch.no_grad():
count = 1
while True:
# Get webcam input
ret_val, img = cam.read()
# AI Mirror
img = cv2.flip(img, 1)
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Generate image
content_tensor = utils.itot(img).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image,
generated_image)
img2 = cv2.imdecode(
cv2.imencode(".png", generated_image)[1], cv2.IMREAD_UNCHANGED)
count += 2
# Show webcam
cv2.imshow('Demo webcam', img2)
if cv2.waitKey(1) == 27:
break # esc to quit
# Free-up memories
cam.release()
cv2.destroyAllWindows()
def stylize_folder_single(style_path, content_folder, save_folder):
"""
Reads frames/pictures as follows:
content_folder
pic1.ext
pic2.ext
pic3.ext
...
and saves as the styled images in save_folder as follow:
save_folder
pic1.ext
pic2.ext
pic3.ext
...
"""
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(style_path))
net = net.to(device)
# Stylize every frame
images = [img for img in os.listdir(content_folder) if img.endswith(".jpg")]
with torch.no_grad():
for image_name in images:
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Load content image
content_image = utils.load_image(content_folder + image_name)
content_tensor = utils.itot(content_image).to(device)
# Generate image
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image, generated_image)
# Save image
utils.saveimg(generated_image, save_folder + image_name)
def stylize():
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Load Transformer Network
net = transformer.TransformerNetwork()
net.load_state_dict(torch.load(STYLE_TRANSFORM_PATH))
net = net.to(device)
with torch.no_grad():
while(1):
torch.cuda.empty_cache()
print("Stylize Image~ Press Ctrl+C and Enter to close the program")
content_image_path = input("Enter the image path: ")
content_image = utils.load_image(content_image_path)
starttime = time.time()
content_tensor = utils.itot(content_image).to(device)
generated_tensor = net(content_tensor)
generated_image = utils.ttoi(generated_tensor.detach())
if (PRESERVE_COLOR):
generated_image = utils.transfer_color(content_image, generated_image)
print("Transfer Time: {}".format(time.time() - starttime))
utils.show(generated_image)
utils.saveimg(generated_image, "helloworld.jpg")
def train():
# Seeds
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)
random.seed(SEED)
# Device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Dataset and Dataloader
transform = transforms.Compose([
transforms.Resize(TRAIN_IMAGE_SIZE),
transforms.CenterCrop(TRAIN_IMAGE_SIZE),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255)),
])
train_dataset = datasets.ImageFolder(DATASET_PATH, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True)
# Load networks
TransformerNetwork = transformer.TransformerNetwork().to(device)
VGG = vgg.VGG16().to(device)
# Get Style Features
imagenet_neg_mean = (torch.tensor([-103.939, -116.779, -123.68],
dtype=torch.float32).reshape(
1, 3, 1, 1).to(device))
style_image = utils.load_image(STYLE_IMAGE_PATH)
style_tensor = utils.itot(style_image).to(device)
style_tensor = style_tensor.add(imagenet_neg_mean)
B, C, H, W = style_tensor.shape
style_features = VGG(style_tensor.expand([BATCH_SIZE, C, H, W]))
style_gram = {}
for key, value in style_features.items():
style_gram[key] = utils.gram(value)
# Optimizer settings
optimizer = optim.Adam(TransformerNetwork.parameters(), lr=ADAM_LR)
# Loss trackers
content_loss_history = []
style_loss_history = []
total_loss_history = []
batch_content_loss_sum = 0
batch_style_loss_sum = 0
batch_total_loss_sum = 0
# Optimization/Training Loop
batch_count = 1
start_time = time.time()
for epoch in range(NUM_EPOCHS):
print("========Epoch {}/{}========".format(epoch + 1, NUM_EPOCHS))
for content_batch, _ in train_loader:
# Get current batch size in case of odd batch sizes
curr_batch_size = content_batch.shape[0]
# Free-up unneeded cuda memory
torch.cuda.empty_cache()
# Zero-out Gradients
optimizer.zero_grad()
# Generate images and get features
content_batch = content_batch[:, [2, 1, 0]].to(device)
generated_batch = TransformerNetwork(content_batch)
content_features = VGG(content_batch.add(imagenet_neg_mean))
generated_features = VGG(generated_batch.add(imagenet_neg_mean))
# Content Loss
MSELoss = nn.MSELoss().to(device)
content_loss = CONTENT_WEIGHT * MSELoss(
generated_features["relu2_2"], content_features["relu2_2"])
batch_content_loss_sum += content_loss
# Style Loss
style_loss = 0
for key, value in generated_features.items():
s_loss = MSELoss(utils.gram(value),
style_gram[key][:curr_batch_size])
style_loss += s_loss
style_loss *= STYLE_WEIGHT
batch_style_loss_sum += style_loss.item()
# Total Loss
total_loss = content_loss + style_loss
batch_total_loss_sum += total_loss.item()
# Backprop and Weight Update
total_loss.backward()
optimizer.step()
# Save Model and Print Losses
if ((batch_count - 1) % SAVE_MODEL_EVERY
== 0) or (batch_count == NUM_EPOCHS * len(train_loader)):
# Print Losses
print("========Iteration {}/{}========".format(
batch_count, NUM_EPOCHS * len(train_loader)))
print("\tContent Loss:\t{:.2f}".format(batch_content_loss_sum /
batch_count))
print("\tStyle Loss:\t{:.2f}".format(batch_style_loss_sum /
batch_count))
print("\tTotal Loss:\t{:.2f}".format(batch_total_loss_sum /
batch_count))
print("Time elapsed:\t{} seconds".format(time.time() -
start_time))
# Save Model
checkpoint_path = (SAVE_MODEL_PATH + "checkpoint_" +
str(batch_count - 1) + ".pth")
torch.save(TransformerNetwork.state_dict(), checkpoint_path)
print("Saved TransformerNetwork checkpoint file at {}".format(
checkpoint_path))
# Save sample generated image
sample_tensor = generated_batch[0].clone().detach().unsqueeze(
dim=0)
sample_image = utils.ttoi(sample_tensor.clone().detach())
sample_image_path = (SAVE_IMAGE_PATH + "sample0_" +
str(batch_count - 1) + ".png")
utils.saveimg(sample_image, sample_image_path)
print("Saved sample tranformed image at {}".format(
sample_image_path))
# Save loss histories
content_loss_history.append(batch_total_loss_sum / batch_count)
style_loss_history.append(batch_style_loss_sum / batch_count)
total_loss_history.append(batch_total_loss_sum / batch_count)
# Iterate Batch Counter
batch_count += 1
stop_time = time.time()
# Print loss histories
print("Done Training the Transformer Network!")
print("Training Time: {} seconds".format(stop_time - start_time))
print("========Content Loss========")
print(content_loss_history)
print("========Style Loss========")
print(style_loss_history)
print("========Total Loss========")
print(total_loss_history)
# Save TransformerNetwork weights
TransformerNetwork.eval()
TransformerNetwork.cpu()
final_path = SAVE_MODEL_PATH + "transformer_weight.pth"
print("Saving TransformerNetwork weights at {}".format(final_path))
torch.save(TransformerNetwork.state_dict(), final_path)
print("Done saving final model")
# Plot Loss Histories
if PLOT_LOSS:
utils.plot_loss_hist(content_loss_history, style_loss_history,
total_loss_history)
def train():
# Seeds
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)
random.seed(SEED)
# Device
device = ("cuda" if torch.cuda.is_available() else "cpu")
# Dataset and Dataloader
transform = transforms.Compose([
transforms.Resize(TRAIN_IMAGE_SIZE),
transforms.CenterCrop(TRAIN_IMAGE_SIZE),
# transforms.Grayscale(num_output_channels=3),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
train_dataset = datasets.ImageFolder(DATASET_PATH, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True)
# Load networks
TransformerNetwork = transformer.TransformerNetwork().to(device)
if USE_LATEST_CHECKPOINT is True:
files = glob.glob(
"/home/clng/github/fast-neural-style-pytorch/models/checkpoint*")
if len(files) == 0:
print("use latest checkpoint but no checkpoint found")
else:
files.sort(key=os.path.getmtime, reverse=True)
latest_checkpoint_path = files[0]
print("using latest checkpoint %s" % (latest_checkpoint_path))
params = torch.load(latest_checkpoint_path, map_location=device)
TransformerNetwork.load_state_dict(params)
VGG = vgg.VGG19().to(device)
# Get Style Features
imagenet_neg_mean = torch.tensor([-103.939, -116.779, -123.68],
dtype=torch.float32).reshape(1, 3, 1,
1).to(device)
style_image = utils.load_image(STYLE_IMAGE_PATH)
if ADJUST_BRIGHTNESS == "1":
style_image = cv2.cvtColor(style_image, cv2.COLOR_BGR2GRAY)
style_image = utils.hist_norm(style_image,
[0, 64, 96, 128, 160, 192, 255],
[0, 0.05, 0.15, 0.5, 0.85, 0.95, 1],
inplace=True)
elif ADJUST_BRIGHTNESS == "2":
style_image = cv2.cvtColor(style_image, cv2.COLOR_BGR2GRAY)
style_image = cv2.equalizeHist(style_image)
elif ADJUST_BRIGHTNESS == "3":
a = 1
# hsv = cv2.cvtColor(style_image, cv2.COLOR_BGR2HSV)
# hsv = utils.auto_brightness(hsv)
# style_image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
style_image = ensure_three_channels(style_image)
sname = os.path.splitext(os.path.basename(STYLE_IMAGE_PATH))[0] + "_train"
cv2.imwrite(
"/home/clng/datasets/bytenow/neural_styles/{s}.jpg".format(s=sname),
style_image)
style_tensor = utils.itot(style_image,
max_size=TRAIN_STYLE_SIZE).to(device)
style_tensor = style_tensor.add(imagenet_neg_mean)
B, C, H, W = style_tensor.shape
style_features = VGG(style_tensor.expand([BATCH_SIZE, C, H, W]))
style_gram = {}
for key, value in style_features.items():
style_gram[key] = utils.gram(value)
# Optimizer settings
optimizer = optim.Adam(TransformerNetwork.parameters(), lr=ADAM_LR)
# Loss trackers
content_loss_history = []
style_loss_history = []
total_loss_history = []
batch_content_loss_sum = 0
batch_style_loss_sum = 0
batch_total_loss_sum = 0
# Optimization/Training Loop
batch_count = 1
start_time = time.time()
for epoch in range(NUM_EPOCHS):
print("========Epoch {}/{}========".format(epoch + 1, NUM_EPOCHS))
for content_batch, _ in train_loader:
# Get current batch size in case of odd batch sizes
curr_batch_size = content_batch.shape[0]
# Free-up unneeded cuda memory
# torch.cuda.empty_cache()
# Zero-out Gradients
optimizer.zero_grad()
# Generate images and get features
content_batch = content_batch[:, [2, 1, 0]].to(device)
generated_batch = TransformerNetwork(content_batch)
content_features = VGG(content_batch.add(imagenet_neg_mean))
generated_features = VGG(generated_batch.add(imagenet_neg_mean))
# Content Loss
MSELoss = nn.MSELoss().to(device)
content_loss = CONTENT_WEIGHT * \
MSELoss(generated_features['relu3_4'],
content_features['relu3_4'])
batch_content_loss_sum += content_loss
# Style Loss
style_loss = 0
for key, value in generated_features.items():
s_loss = MSELoss(utils.gram(value),
style_gram[key][:curr_batch_size])
style_loss += s_loss
style_loss *= STYLE_WEIGHT
batch_style_loss_sum += style_loss.item()
# Total Loss
total_loss = content_loss + style_loss
batch_total_loss_sum += total_loss.item()
# Backprop and Weight Update
total_loss.backward()
optimizer.step()
# Save Model and Print Losses
if (((batch_count - 1) % SAVE_MODEL_EVERY == 0)
or (batch_count == NUM_EPOCHS * len(train_loader))):
# Print Losses
print("========Iteration {}/{}========".format(
batch_count, NUM_EPOCHS * len(train_loader)))
print("\tContent Loss:\t{:.2f}".format(batch_content_loss_sum /
batch_count))
print("\tStyle Loss:\t{:.2f}".format(batch_style_loss_sum /
batch_count))
print("\tTotal Loss:\t{:.2f}".format(batch_total_loss_sum /
batch_count))
print("Time elapsed:\t{} seconds".format(time.time() -
start_time))
# Save Model
checkpoint_path = SAVE_MODEL_PATH + "checkpoint_" + str(
batch_count - 1) + ".pth"
torch.save(TransformerNetwork.state_dict(), checkpoint_path)
print("Saved TransformerNetwork checkpoint file at {}".format(
checkpoint_path))
# Save sample generated image
sample_tensor = generated_batch[0].clone().detach().unsqueeze(
dim=0)
sample_image = utils.ttoi(sample_tensor.clone().detach())
sample_image_path = SAVE_IMAGE_PATH + "sample0_" + str(
batch_count - 1) + ".png"
utils.saveimg(sample_image, sample_image_path)
print("Saved sample tranformed image at {}".format(
sample_image_path))
# Save loss histories
content_loss_history.append(batch_total_loss_sum / batch_count)
style_loss_history.append(batch_style_loss_sum / batch_count)
total_loss_history.append(batch_total_loss_sum / batch_count)
# Iterate Batch Counter
batch_count += 1
stop_time = time.time()
# Print loss histories
print("Done Training the Transformer Network!")
print("Training Time: {} seconds".format(stop_time - start_time))
print("========Content Loss========")
print(content_loss_history)
print("========Style Loss========")
print(style_loss_history)
print("========Total Loss========")
print(total_loss_history)
# Save TransformerNetwork weights
TransformerNetwork.eval()
TransformerNetwork.cpu()
final_path = SAVE_MODEL_PATH + STYLE_NAME + ".pth"
print("Saving TransformerNetwork weights at {}".format(final_path))
torch.save(TransformerNetwork.state_dict(), final_path)
print("Done saving final model")
# Plot Loss Histories
if (PLOT_LOSS):
utils.plot_loss_hist(content_loss_history, style_loss_history,
total_loss_history)