def _valid_args(self, image: Path, copy_style_image: Path, _output,
_resize, progress_dir: Path, _max_size) -> (bool, str):
if not image.exists():
raise InvalidCopyStyleArgs(
"Given image path does not exist:\n\t{}".format(str(image)))
if not copy_style_image.exists():
raise InvalidCopyStyleArgs(
"Given image path does not exist:\n\t{}".format(
str(copy_style_image)))
if not progress_dir.exists():
eprint("Created new directory:\n\t" + str(progress_dir))
progress_dir.mkdir(parents=True, exist_ok=True)
return (True, "")
def closure():
# correct the values of updated input image
input_img.data.clamp_(0, 1)
optimizer.zero_grad()
model(input_img)
style_score = 0
content_score = 0
for sl in style_losses:
style_score += sl.loss
for cl in content_losses:
content_score += cl.loss
style_score *= style_weight
content_score *= content_weight
loss = style_score + content_score
loss.backward()
run[0] += 1
if run[0] % 50 == 0:
eprint("run {}:".format(run))
eprint("Style Loss : {:4f} Content Loss: {:4f}".format(
style_score.item(), content_score.item()))
eprint()
if progress_dir:
filename = Path(progress_dir,
'progress-{}.png'.format(run[0]))
save_image(input_img, filename)
return style_score + content_score
def main():
start = time.time()
parser = argparse.ArgumentParser(description="Copy art styles to images")
parser.add_argument("IMAGE", help="the image you want transformed")
parser.add_argument("-s",
"--style",
help="the piece where the style is copied from",
required=True)
parser.add_argument(
"-o",
"--output",
help=
"the output file. If this is not provided then [IMAGE].style-copied is used",
required=True,
)
parser.add_argument("-r",
"--resize",
help="resize the image to the copy image",
action="store_true")
parser.add_argument("-p",
"--progress",
help="directory to output progress")
parser.add_argument("-m",
"--max-size",
help="use IMSIZE 512 regarldess of GPU availibilyt",
action="store_true")
args = parser.parse_args()
image = args.IMAGE
resize = args.resize
style_image = args.style
progress_dir = args.progress
max_size = args.max_size
error = 0
if not image:
eprint("Image cannot be empty!")
sys.exit(1)
if not style_image:
eprint("Copy cannot be empty!")
sys.exit(1)
try:
args = CopyStyleArgs(image,
style_image,
output=args.output,
resize=resize,
progress_dir=progress_dir,
max_size=max_size)
run(args)
except Exception as e:
error = 1
eprint(str(e))
traceback.print_exc()
if error == 1:
sys.exit(error)
end = time.time()
print("Time elapsed: " + str(end - start) + "seconds")
def run_style_transfer(
cnn,
normalization_mean,
normalization_std,
content_img,
style_img,
input_img,
num_steps=300,
style_weight=1000000,
content_weight=1,
progress_dir: Path = None,
) -> FloatTensor:
"""Run the style transfer."""
eprint("Building the style transfer model..")
model, style_losses, content_losses = get_style_model_and_losses(
cnn, normalization_mean, normalization_std, style_img, content_img)
optimizer = get_input_optimizer(input_img)
eprint("Optimizing..")
run = [0]
while run[0] <= num_steps:
def closure():
# correct the values of updated input image
input_img.data.clamp_(0, 1)
optimizer.zero_grad()
model(input_img)
style_score = 0
content_score = 0
for sl in style_losses:
style_score += sl.loss
for cl in content_losses:
content_score += cl.loss
style_score *= style_weight
content_score *= content_weight
loss = style_score + content_score
loss.backward()
run[0] += 1
if run[0] % 50 == 0:
eprint("run {}:".format(run))
eprint("Style Loss : {:4f} Content Loss: {:4f}".format(
style_score.item(), content_score.item()))
eprint()
if progress_dir:
filename = Path(progress_dir,
'progress-{}.png'.format(run[0]))
save_image(input_img, filename)
return style_score + content_score
optimizer.step(closure)
# a last correction...
input_img.data.clamp_(0, 1)
return input_img
def _imsize(use_large=False):
if use_large:
eprint("Using image size of {0} instead of {1}".format(_LARGE, _SMALL))
return _LARGE
has_gpu = torch.cuda.is_available()
return _LARGE if has_gpu else _SMALL