import sys import torch import torch.nn as nn import numpy as np from torch.nn import CTCLoss import os import cv2 from tqdm import tqdm from params import * from .BigGAN_networks import Discriminator from util.util import ( loss_hinge_dis, loss_hinge_gen, padding, ) from data.dataset import TextDataset, TextDatasetval import shutil from .recognizer import ViT_OCR from .Generator import Generator from .Writer import Writer, strLabelConverter class WriteViT(nn.Module): def __init__(self, batch_size=batch_size,backbone="resnet18"): super(WriteViT, self).__init__() self.batch_size = batch_size self.epsilon = 1e-7 self.netG = Generator().to(DEVICE) self.netD = nn.DataParallel(Discriminator()).to(DEVICE) self.netW = Writer().to(DEVICE) self.netconverter = strLabelConverter(ALPHABET) self.netOCR = ViT_OCR(backbone=backbone).to(DEVICE) self.OCR_criterion = CTCLoss(zero_infinity=True, reduction="none") self.optimizer_G = torch.optim.Adam( self.netG.parameters(), lr=G_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8, ) self.optimizer_OCR = torch.optim.Adam( self.netOCR.parameters(), lr=OCR_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8, ) self.optimizer_D = torch.optim.Adam( self.netD.parameters(), lr=D_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8, ) self.optimizer_wl = torch.optim.Adam( self.netW.parameters(), lr=W_LR, betas=(0.0, 0.999), weight_decay=0, eps=1e-8, ) self.optimizers = [ self.optimizer_G, self.optimizer_OCR, self.optimizer_D, self.optimizer_wl, ] self.optimizer_G.zero_grad() self.optimizer_OCR.zero_grad() self.optimizer_D.zero_grad() self.optimizer_wl.zero_grad() self.loss_G = 0 self.loss_D = 0 self.loss_Dfake = 0 self.loss_Dreal = 0 self.loss_OCR_fake = 0 self.loss_OCR_real = 0 self.loss_w_fake = 0 self.loss_w_real = 0 self.Lcycle1 = 0 self.Lcycle2 = 0 self.lda1 = 0 self.lda2 = 0 self.KLD = 0 self.loss_patch_real = 0 self.loss_patch_fake = 0 self.loss_patch = 0 with open(WORDS_PATH, "rb") as f: self.lex = f.read().splitlines() lex = [] lex_upper_number = [] for word in self.lex: try: word = word.decode("utf-8") except: continue if len(word) < 20: if word.isupper() or word.isdigit(): lex_upper_number.append(word) else: lex.append(word) self.lex = lex self.lex_upper_number = lex_upper_number self.fake_y_dist = torch.distributions.Categorical( torch.tensor([1.0 / len(self.lex)] * len(self.lex)) ) my_string = MY_STRING self.text = [j.encode() for j in my_string.split(" ")] self.eval_text_encode, self.eval_len_text = self.netconverter.encode(self.text) self.eval_text_encode = self.eval_text_encode.to(DEVICE).repeat( self.batch_size, 1, 1 ) def save_images_for_fid_calculation(self, dataloader, epoch, mode="train"): self.real_base = os.path.join("saved_images", EXP_NAME, "Real") self.fake_base = os.path.join("saved_images", EXP_NAME, "Fake") if os.path.isdir(self.real_base): shutil.rmtree(self.real_base) if os.path.isdir(self.fake_base): shutil.rmtree(self.fake_base) os.makedirs(self.real_base, exist_ok=True) os.makedirs(self.fake_base, exist_ok=True) # ========================= # Save fake images # ========================= with torch.no_grad(): for step, data in enumerate(tqdm(dataloader)): self.sdata = data["img"].to(DEVICE) self.label = data["label"] writer_ids = data["wcl"] # nhãn id người viết, shape thường là [B] self.text_encode_fake, self.len_text_fake = self.netconverter.encode(self.label) self.text_encode_fake = self.text_encode_fake.to(DEVICE) feat_w, _ = self.netW(self.sdata.detach(), writer_ids.to(DEVICE)) self.fakes = self.netG(feat_w, self.text_encode_fake) fake_images = self.fakes.detach().cpu().numpy() # fake_images: thường là [B, C, H, W] hoặc [B, N, H, W] for i in range(fake_images.shape[0]): writer_id = writer_ids[i].item() if torch.is_tensor(writer_ids[i]) else int(writer_ids[i]) if mode == "train": writer_fake_dir = os.path.join(self.fake_base, str(writer_id)) os.makedirs(writer_fake_dir, exist_ok=True) else: writer_fake_dir = self.fake_base for j in range(fake_images.shape[1]): img = 255 * (((fake_images[i, j]) + 1) / 2) img = padding(img) filename = f"{step}_{i}_{j}.png" cv2.imwrite( os.path.join(writer_fake_dir, filename), img, ) # ========================= # Save real images # ========================= for step, data in enumerate(tqdm(dataloader)): real_images = data["img"].numpy() writer_ids = data["wcl"] for i in range(real_images.shape[0]): writer_id = writer_ids[i].item() if torch.is_tensor(writer_ids[i]) else int(writer_ids[i]) if mode == "train": writer_real_dir = os.path.join(self.real_base, str(writer_id)) os.makedirs(writer_real_dir, exist_ok=True) else: writer_real_dir = self.real_base for j in range(real_images.shape[1]): img = 255 * ((real_images[i, j] + 1) / 2) img = padding(img) filename = f"{step}_{i}_{j}.png" cv2.imwrite( os.path.join(writer_real_dir, filename), img, ) return self.real_base, self.fake_base def _generate_page( self, img, ST, wcl,SLEN, eval_text_encode=None, eval_len_text=None ): if eval_text_encode == None: eval_text_encode = self.eval_text_encode if eval_len_text == None: eval_len_text = self.eval_len_text feat_w = self.netW(img.detach(), wcl,training=False) self.fakes = self.netG.Eval(feat_w, eval_text_encode) page1s = [] page2s = [] for batch_idx in range(self.batch_size): word_t = [] word_l = [] gap = np.ones([IMG_HEIGHT, 16]) line_wids = [] for idx, fake_ in enumerate(self.fakes): word_t.append( ( fake_[batch_idx, 0, :, : eval_len_text[idx] * resolution] .cpu() .numpy() + 1 ) / 2 ) word_t.append(gap) if len(word_t) == 16 or idx == len(self.fakes) - 1: line_ = np.concatenate(word_t, -1) word_l.append(line_) line_wids.append(line_.shape[1]) word_t = [] gap_h = np.ones([16, max(line_wids)]) page_ = [] for l in word_l: pad_ = np.ones([IMG_HEIGHT, max(line_wids) - l.shape[1]]) page_.append(np.concatenate([l, pad_], 1)) page_.append(gap_h) page1 = np.concatenate(page_, 0) word_t = [] word_l = [] gap = np.ones([IMG_HEIGHT, 16]) line_wids = [] sdata_ = [i.unsqueeze(1) for i in torch.unbind(ST, 1)] for idx, st in enumerate((sdata_)): word_t.append( ( st[batch_idx, 0, :, : int(SLEN.cpu().numpy()[batch_idx][idx])] .cpu() .numpy() + 1 ) / 2 ) word_t.append(gap) if len(word_t) == 16 or idx == len(sdata_) - 1: line_ = np.concatenate(word_t, -1) word_l.append(line_) line_wids.append(line_.shape[1]) word_t = [] gap_h = np.ones([16, max(line_wids)]) page_ = [] for l in word_l: pad_ = np.ones([IMG_HEIGHT, max(line_wids) - l.shape[1]]) page_.append(np.concatenate([l, pad_], 1)) page_.append(gap_h) page2 = np.concatenate(page_, 0) merge_w_size = max(page1.shape[0], page2.shape[0]) if page1.shape[0] != merge_w_size: page1 = np.concatenate( [page1, np.ones([merge_w_size - page1.shape[0], page1.shape[1]])], 0 ) if page2.shape[0] != merge_w_size: page2 = np.concatenate( [page2, np.ones([merge_w_size - page2.shape[0], page2.shape[1]])], 0 ) page1s.append(page1) page2s.append(page2) # page = np.concatenate([page2, page1], 1) page1s_ = np.concatenate(page1s, 0) max_wid = max([i.shape[1] for i in page2s]) padded_page2s = [] for para in page2s: padded_page2s.append( np.concatenate( [para, np.ones([para.shape[0], max_wid - para.shape[1]])], 1 ) ) padded_page2s_ = np.concatenate(padded_page2s, 0) return np.concatenate([padded_page2s_, page1s_], 1) def get_current_losses(self): losses = {} losses["G"] = self.loss_G losses["D"] = self.loss_D losses["Dfake"] = self.loss_Dfake losses["Dreal"] = self.loss_Dreal losses["OCR_fake"] = self.loss_OCR_fake losses["OCR_real"] = self.loss_OCR_real losses["w_fake"] = self.loss_w_fake losses["w_real"] = self.loss_w_real losses["cycle1"] = self.Lcycle1 losses["cycle2"] = self.Lcycle2 losses["lda1"] = self.lda1 losses["lda2"] = self.lda2 losses["KLD"] = self.KLD losses["patch_real"] = self.loss_patch_real losses["patch_fake"] = self.loss_patch_fake losses["patch"] = self.loss_patch return losses def load_networks(self, epoch): BaseModel.load_networks(self, epoch) if self.opt.single_writer: load_filename = "%s_z.pkl" % (epoch) load_path = os.path.join(self.save_dir, load_filename) self.z = torch.load(load_path) def _set_input(self, input): self.input = input def set_requires_grad(self, nets, requires_grad=False): """Set requies_grad=Fasle for all the networks to avoid unnecessary computations Parameters: nets (network list) -- a list of networks requires_grad (bool) -- whether the networks require gradients or not """ if not isinstance(nets, list): nets = [nets] for net in nets: if net is not None: for param in net.parameters(): param.requires_grad = requires_grad def forward(self): self.real = self.input["img"].to(DEVICE) self.label = self.input["label"] self.sdata = self.input["img"].to(DEVICE) self.ST_LEN = self.input["swids"] self.text_encode, self.len_text = self.netconverter.encode(self.label) self.text_encode = self.text_encode.to(DEVICE).detach() self.len_text = self.len_text.detach() sample_lex_idx = self.fake_y_dist.sample([self.batch_size]) fake_y = [self.lex[i].encode("utf-8") for i in sample_lex_idx] self.text_encode_fake, self.len_text_fake = self.netconverter.encode(fake_y) self.text_encode_fake = self.text_encode_fake.to(DEVICE) def backward_D_OCR_W(self): feat_w, self.loss_w_real = self.netW( self.real.detach(), self.input["wcl"].to(DEVICE) ) _, self.pred_real_OCR = self.netOCR(self.real.detach()) self.loss_w_real = self.loss_w_real.mean() self.fake = self.netG(feat_w, self.text_encode_fake) pred_real = self.netD(self.real.detach()) pred_fake = self.netD(**{"x": self.fake.detach()}) self.loss_Dreal, self.loss_Dfake = loss_hinge_dis( pred_fake, pred_real, self.len_text_fake.detach(), self.len_text.detach(), True, ) self.loss_D = self.loss_Dreal + self.loss_Dfake self.pred_real_OCR = self.pred_real_OCR.float() preds_size = torch.IntTensor( [self.pred_real_OCR.size(1)] * self.batch_size ).detach() self.pred_real_OCR = self.pred_real_OCR.permute(1, 0, 2).log_softmax(2) loss_OCR_real = self.OCR_criterion( self.pred_real_OCR, self.text_encode.detach(), preds_size, self.len_text.detach(), ) self.loss_OCR_real = torch.mean(loss_OCR_real[~torch.isnan(loss_OCR_real)]) loss_total = ( self.loss_D * 2 + self.loss_OCR_real + self.loss_w_real ) # backward loss_total.backward() return loss_total def backward_G_only(self): self.gb_alpha = 0.7 self.gb_beta = 0.7 feat_w, _ = self.netW(self.real.detach(), self.input["wcl"].to(DEVICE)) self.fake = self.netG(feat_w, self.text_encode_fake) pred_fake = self.netD(**{"x": self.fake}) self.loss_G = loss_hinge_gen( pred_fake, self.len_text_fake.detach(), True ).mean() _, pred_fake_OCR = self.netOCR(self.fake) pred_fake_OCR = pred_fake_OCR.float() preds_size = torch.IntTensor([pred_fake_OCR.size(1)] * self.batch_size).detach() pred_fake_OCR = pred_fake_OCR.permute(1, 0, 2).log_softmax(2) loss_OCR_fake = self.OCR_criterion( pred_fake_OCR, self.text_encode_fake.detach(), preds_size, self.len_text_fake.detach(), ) self.loss_OCR_fake = torch.mean(loss_OCR_fake[~torch.isnan(loss_OCR_fake)]) _, self.loss_w_fake = self.netW(self.fake, self.input["wcl"].to(DEVICE)) self.loss_w_fake = self.loss_w_fake.mean() self.loss_G = self.loss_G self.loss_T = ( self.loss_G + self.loss_OCR_fake + self.loss_w_fake ) grad_fake_OCR = torch.autograd.grad( self.loss_OCR_fake, self.fake, retain_graph=True )[0] grad_fake_WL = torch.autograd.grad( self.loss_w_fake, self.fake, retain_graph=True )[0] grad_fake_adv = torch.autograd.grad(self.loss_G, self.fake, retain_graph=True)[ 0 ] self.loss_T.backward(retain_graph=True) if True: grad_fake_OCR = torch.autograd.grad( self.loss_OCR_fake, self.fake, create_graph=True, retain_graph=True )[0] grad_fake_adv = torch.autograd.grad( self.loss_G, self.fake, create_graph=True, retain_graph=True )[0] grad_fake_WL = torch.autograd.grad( self.loss_w_fake, self.fake, create_graph=True, retain_graph=True )[0] gp_ocr = self.gb_alpha * torch.div( torch.std(grad_fake_adv), self.epsilon + torch.std(grad_fake_OCR) ) gp_wl = self.gb_beta * torch.div( torch.std(grad_fake_adv), self.epsilon + torch.std(grad_fake_WL) ) self.loss_OCR_fake = gp_ocr.detach() * self.loss_OCR_fake self.loss_w_fake = gp_wl.detach() * self.loss_w_fake self.loss_T = ( self.loss_G * 2 + self.loss_OCR_fake + self.loss_w_fake ) self.loss_T.backward(retain_graph=True) with torch.no_grad(): self.loss_T.backward() def optimize_D_OCR_W(self): self.forward() self.set_requires_grad([self.netD], True) self.set_requires_grad([self.netOCR], True) self.set_requires_grad([self.netW], True) self.optimizer_D.zero_grad() self.optimizer_OCR.zero_grad() self.optimizer_wl.zero_grad() self.backward_D_OCR_W() def optimize_D_OCR_W_step(self): self.optimizer_D.step() self.optimizer_wl.step() self.optimizer_OCR.step() self.optimizer_D.zero_grad() self.optimizer_OCR.zero_grad() self.optimizer_wl.zero_grad() def optimize_G_only(self): self.forward() self.set_requires_grad([self.netD], False) self.set_requires_grad([self.netOCR], False) self.set_requires_grad([self.netW], False) self.backward_G_only() def optimize_G_step(self): self.optimizer_G.step() self.optimizer_G.zero_grad()