from transformers import PreTrainedModel
from .configuration_vatrpp import VATrPPConfig
import json
import os
import random
import shutil
from collections import defaultdict
import time
from datetime import timedelta
from pathlib import Path
import cv2
import numpy as np
import torch
from data.dataset import FolderDataset
from models.model import VATr
from util.loading import load_checkpoint, load_generator
from util.misc import FakeArgs
from util.text import TextGenerator
from util.vision import detect_text_bounds
from torchvision.transforms.functional import to_pil_image
def get_long_tail_chars():
with open(f"files/longtail.txt", 'r') as f:
chars = [c.rstrip() for c in f]
chars.remove('')
return chars
class VATrPP(PreTrainedModel):
config_class = VATrPPConfig
def __init__(self, config: VATrPPConfig) -> None:
super().__init__(config)
self.model = VATr(config)
self.model.eval()
def set_style_folder(self, style_folder, num_examples=15):
word_lengths = None
if os.path.exists(os.path.join(style_folder, "word_lengths.txt")):
word_lengths = {}
with open(os.path.join(style_folder, "word_lengths.txt"), 'r') as f:
for line in f:
word, length = line.rstrip().split(",")
word_lengths[word] = int(length)
self.style_dataset = FolderDataset(style_folder, num_examples=num_examples, word_lengths=word_lengths)
@torch.no_grad()
def generate(self, gen_text, style_imgs, align_words: bool = False, at_once: bool = False):
style_images = style_imgs.unsqueeze(0).to(self.model.args.device)
fake = self.create_fake_sentence(style_images, gen_text, align_words, at_once)
return to_pil_image(fake)
# @torch.no_grad()
# def generate(self, texts, align_words: bool = False, at_once: bool = False):
# if isinstance(texts, str):
# texts = [texts]
# if self.style_dataset is None:
# raise Exception('Style is not set')
# fakes = []
# for i, text in enumerate(texts, 1):
# print(f'[{i}/{len(texts)}] Generating for text: {text}')
# style = self.style_dataset.sample_style()
# style_images = style['simg'].unsqueeze(0).to(self.model.args.device)
# fake = self.create_fake_sentence(style_images, text, align_words, at_once)
# fakes.append(fake)
# return fakes
@torch.no_grad()
def create_fake_sentence(self, style_images, text, align_words=False, at_once=False):
text = "".join([c for c in text if c in self.model.args.alphabet])
text = text.split() if not at_once else [text]
gap = np.ones((32, 16))
text_encode, len_text, encode_pos = self.model.netconverter.encode(text)
text_encode = text_encode.to(self.model.args.device).unsqueeze(0)
fake = self.model._generate_fakes(style_images, text_encode, len_text)
if not at_once:
if align_words:
fake = self.stitch_words(fake, show_lines=False)
else:
fake = np.concatenate(sum([[img, gap] for img in fake], []), axis=1)[:, :-16]
else:
fake = fake[0]
fake = (fake * 255).astype(np.uint8)
return fake
@torch.no_grad()
def generate_authors(self, text, dataset, align_words: bool = False, at_once: bool = False):
fakes = []
author_ids = []
style = []
for item in dataset:
print(f"Generating author {item['wcl']}")
style_images = item['simg'].to(self.model.args.device).unsqueeze(0)
generated_lines = [self.create_fake_sentence(style_images, line, align_words, at_once) for line in text]
fakes.append(generated_lines)
author_ids.append(item['author_id'])
style.append((((item['simg'].numpy() + 1.0) / 2.0) * 255).astype(np.uint8))
return fakes, author_ids, style
@torch.no_grad()
def generate_characters(self, dataset, characters: str):
"""
Generate each of the given characters for each of the authors in the dataset.
"""
fakes = []
text_encode, len_text, encode_pos = self.model.netconverter.encode([c for c in characters])
text_encode = text_encode.to(self.model.args.device).unsqueeze(0)
for item in dataset:
print(f"Generating author {item['wcl']}")
style_images = item['simg'].to(self.model.args.device).unsqueeze(0)
fake = self.model.netG.evaluate(style_images, text_encode)
fakes.append(fake)
return fakes
@torch.no_grad()
def generate_batch(self, style_imgs, text):
"""
Given a batch of style images and text, generate images using the model
"""
device = self.model.args.device
text_encode, _, _ = self.model.netconverter.encode(text)
fakes, _ = self.model.netG(style_imgs.to(device), text_encode.to(device))
return fakes
@torch.no_grad()
def generate_ocr(self, dataset, number: int, output_folder: str = 'saved_images/ocr', interpolate_style: bool = False, text_generator: TextGenerator = None, long_tail: bool = False):
def create_and_write(style, text, interpolated=False):
nonlocal image_counter, annotations
text_encode, len_text, encode_pos = self.model.netconverter.encode([text])
text_encode = text_encode.to(self.model.args.device)
fake = self.model.netG.generate(style, text_encode)
fake = (fake + 1) / 2
fake = fake.cpu().numpy()
fake = np.squeeze((fake * 255).astype(np.uint8))
image_filename = f"{image_counter}.png" if not interpolated else f"{image_counter}_i.png"
cv2.imwrite(os.path.join(output_folder, "generated", image_filename), fake)
annotations.append((image_filename, text))
image_counter += 1
image_counter = 0
annotations = []
previous_style = None
long_tail_chars = get_long_tail_chars()
os.mkdir(os.path.join(output_folder, "generated"))
if text_generator is None:
os.mkdir(os.path.join(output_folder, "reference"))
while image_counter < number:
author_index = random.randint(0, len(dataset) - 1)
item = dataset[author_index]
style_images = item['simg'].to(self.model.args.device).unsqueeze(0)
style = self.model.netG.compute_style(style_images)
if interpolate_style and previous_style is not None:
factor = float(np.clip(random.gauss(0.5, 0.15), 0.0, 1.0))
intermediate_style = torch.lerp(previous_style, style, factor)
text = text_generator.generate()
create_and_write(intermediate_style, text, interpolated=True)
if text_generator is not None:
text = text_generator.generate()
else:
text = str(item['label'].decode())
if long_tail and not any(c in long_tail_chars for c in text):
continue
fake = (item['img'] + 1) / 2
fake = fake.cpu().numpy()
fake = np.squeeze((fake * 255).astype(np.uint8))
image_filename = f"{image_counter}.png"
cv2.imwrite(os.path.join(output_folder, "reference", image_filename), fake)
create_and_write(style, text)
previous_style = style
if text_generator is None:
with open(os.path.join(output_folder, "reference", "labels.csv"), 'w') as fr:
fr.write(f"filename,words\n")
for annotation in annotations:
fr.write(f"{annotation[0]},{annotation[1]}\n")
with open(os.path.join(output_folder, "generated", "labels.csv"), 'w') as fg:
fg.write(f"filename,words\n")
for annotation in annotations:
fg.write(f"{annotation[0]},{annotation[1]}\n")
@staticmethod
def stitch_words(words: list, show_lines: bool = False, scale_words: bool = False):
gap_width = 16
bottom_lines = []
top_lines = []
for i in range(len(words)):
b, t = detect_text_bounds(words[i])
bottom_lines.append(b)
top_lines.append(t)
if show_lines:
words[i] = cv2.line(words[i], (0, b), (words[i].shape[1], b), (0, 0, 1.0))
words[i] = cv2.line(words[i], (0, t), (words[i].shape[1], t), (1.0, 0, 0))
bottom_lines = np.array(bottom_lines, dtype=float)
if scale_words:
top_lines = np.array(top_lines, dtype=float)
gaps = bottom_lines - top_lines
target_gap = np.mean(gaps)
scales = target_gap / gaps
bottom_lines *= scales
top_lines *= scales
words = [cv2.resize(word, None, fx=scale, fy=scale) for word, scale in zip(words, scales)]
highest = np.max(bottom_lines)
offsets = highest - bottom_lines
height = np.max(offsets + [word.shape[0] for word in words])
result = np.ones((int(height), gap_width * len(words) + sum([w.shape[1] for w in words])))
x_pos = 0
for bottom_line, word in zip(bottom_lines, words):
offset = int(highest - bottom_line)
result[offset:offset + word.shape[0], x_pos:x_pos+word.shape[1]] = word
x_pos += word.shape[1] + gap_width
return result
@torch.no_grad()
def generate_fid(self, path, loader, model_tag, split='train', fake_only=False, long_tail_only=False):
if not isinstance(path, Path):
path = Path(path)
path.mkdir(exist_ok=True, parents=True)
appendix = f"{split}" if not long_tail_only else f"{split}_lt"
real_base = path / f'real_{appendix}'
fake_base = path / model_tag / f'fake_{appendix}'
if real_base.exists() and not fake_only:
shutil.rmtree(real_base)
if fake_base.exists():
shutil.rmtree(fake_base)
real_base.mkdir(exist_ok=True)
fake_base.mkdir(exist_ok=True, parents=True)
print('Saving images...')
print(' Saving images on {}'.format(str(real_base)))
print(' Saving images on {}'.format(str(fake_base)))
long_tail_chars = get_long_tail_chars()
counter = 0
ann = defaultdict(lambda: {})
start_time = time.time()
for step, data in enumerate(loader):
style_images = data['simg'].to(self.model.args.device)
texts = [l.decode('utf-8') for l in data['label']]
texts = [t.encode('utf-8') for t in texts]
eval_text_encode, eval_len_text, _ = self.model.netconverter.encode(texts)
eval_text_encode = eval_text_encode.to(self.model.args.device).unsqueeze(1)
vis_style = np.vstack(style_images[0].detach().cpu().numpy())
vis_style = ((vis_style + 1) / 2) * 255
fakes = self.model.netG.evaluate(style_images, eval_text_encode)
fake_images = torch.cat(fakes, 1).detach().cpu().numpy()
real_images = data['img'].detach().cpu().numpy()
writer_ids = data['wcl'].int().tolist()
for i, (fake, real, wid, lb, img_id) in enumerate(zip(fake_images, real_images, writer_ids, data['label'], data['idx'])):
lb = lb.decode()
ann[f"{wid:03d}"][f'{img_id:05d}'] = lb
img_id = f'{img_id:05d}.png'
is_long_tail = any(c in long_tail_chars for c in lb)
if long_tail_only and not is_long_tail:
continue
fake_img_path = fake_base / f"{wid:03d}" / img_id
fake_img_path.parent.mkdir(exist_ok=True, parents=True)
cv2.imwrite(str(fake_img_path), 255 * ((fake.squeeze() + 1) / 2))
if not fake_only:
real_img_path = real_base / f"{wid:03d}" / img_id
real_img_path.parent.mkdir(exist_ok=True, parents=True)
cv2.imwrite(str(real_img_path), 255 * ((real.squeeze() + 1) / 2))
counter += 1
eta = (time.time() - start_time) / (step + 1) * (len(loader) - step - 1)
eta = str(timedelta(seconds=eta))
if step % 100 == 0:
print(f'[{(step + 1) / len(loader) * 100:.02f}%][{counter:05d}] ETA {eta}')
with open(path / 'ann.json', 'w') as f:
json.dump(ann, f)