| import os | |
| import shutil | |
| import cv2 | |
| import numpy as np | |
| from data.dataset import CollectionTextDataset, TextDataset | |
| from generate.util import stack_lines | |
| from generate.writer import Writer | |
| def generate_authors(args): | |
| dataset = CollectionTextDataset( | |
| args.dataset, 'files', TextDataset, file_suffix=args.file_suffix, num_examples=args.num_examples, | |
| collator_resolution=args.resolution, validation=args.test_set | |
| ) | |
| args.num_writers = dataset.num_writers | |
| writer = Writer(args.checkpoint, args, only_generator=True) | |
| if args.text.endswith(".txt"): | |
| with open(args.text, 'r') as f: | |
| lines = [l.rstrip() for l in f] | |
| else: | |
| lines = [args.text] | |
| output_dir = "saved_images/author_samples/" | |
| if os.path.exists(output_dir): | |
| shutil.rmtree(output_dir) | |
| os.mkdir(output_dir) | |
| fakes, author_ids, style_images = writer.generate_authors(lines, dataset, args.align, args.at_once) | |
| for fake, author_id, style in zip(fakes, author_ids, style_images): | |
| author_dir = os.path.join(output_dir, str(author_id)) | |
| os.mkdir(author_dir) | |
| for i, line in enumerate(fake): | |
| cv2.imwrite(os.path.join(author_dir, f"line_{i}.png"), line) | |
| total = stack_lines(fake) | |
| cv2.imwrite(os.path.join(author_dir, "total.png"), total) | |
| if args.output_style: | |
| for i, image in enumerate(style): | |
| cv2.imwrite(os.path.join(author_dir, f"style_{i}.png"), image) | |