File size: 3,416 Bytes
1a41d53 | 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 | import os
import copy
import torch
import random
import argparse
import numpy as np
from PIL import Image
from pathlib import Path
from tqdm.auto import tqdm
from lib.style_transfer_utils import (
tensor2pil,
load_style_transfer_model,
run_style_transfer,
style_content_image_loader,
)
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument(
"--content-imgs", type=str, help="Path to the content images.", required=True
)
parser.add_argument(
"--style-imgs", type=str, help="Path to the style images.", required=True
)
parser.add_argument(
"--save-folder",
type=str,
help="Path to the save the generated images.",
required=True,
)
parser.add_argument(
"--vgg", type=str, help="Path to the pretrained VGG model.", required=True
)
parser.add_argument("--cuda", action="store_true", help="use cuda.")
parser.add_argument(
"--ext", type=str, default="stl", help="extension for generated image."
)
parser.add_argument(
"--min-step", type=int, default=100, help="minimum iteration steps"
)
parser.add_argument(
"--max-step", type=int, default=200, help="maximum iteration steps"
)
parser.add_argument(
"--style-weight", type=float, default=100000, help="weight for style loss"
)
parser.add_argument(
"--content-weight", type=float, default=2, help="weight for content loss"
)
return parser.parse_args()
def transfer_style(
cnn_path,
cimg,
simg,
min_step=100,
max_step=200,
style_weight=100000,
content_weight=2,
device="cpu",
):
cnn = load_style_transfer_model(pretrained=cnn_path)
content_img, style_img = style_content_image_loader(cimg, simg)
input_img = copy.deepcopy(content_img).to(device, torch.float)
output = run_style_transfer(
cnn,
content_img,
style_img,
input_img,
num_steps=random.randint(min_step, max_step),
style_weight=style_weight,
content_weight=content_weight,
device=device,
)
return tensor2pil(output[0].detach().cpu())
def main():
args = parse_arguments()
if args.cuda and torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
content_images = sorted(Path(args.content_imgs).glob("*"))
# with open(Path(args.content_imgs), "r") as f:
# lines = f.read()
# content_images = lines.split("\n")
# content_images = [Path("./content_images") / f for f in content_images]
style_images = sorted(Path(args.style_imgs).glob("*"))
save_folder = Path(args.save_folder)
if not os.path.exists(args.save_folder):
print(f"Creating {args.save_folder}")
os.makedirs(str(save_folder))
for i, cimg in enumerate(content_images):
name, extension = cimg.name.split(".")
simg = style_images[i % len(style_images)]
output_img = transfer_style(
cnn_path=args.vgg,
cimg=cimg,
simg=simg,
min_step=args.min_step,
max_step=args.max_step,
style_weight=args.style_weight,
content_weight=args.content_weight,
device=device,
)
output_img.save(save_folder / f"{name}.{extension}")
if __name__ == "__main__":
main()
|