# ============================================ # ESRGAN for Artwork Super Resolution # Author: Minimal Training Script # ============================================ import os import glob import random import torch import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms from torch.utils.data import Dataset, DataLoader from PIL import Image from torchvision.models import vgg19 import torchvision.utils as vutils device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ============================================ # Dataset # ============================================ class ArtworkDataset(Dataset): def __init__(self, hr_dir, scale=4): self.hr_images = glob.glob(hr_dir + "/*.png") self.scale = scale self.hr_transform = transforms.Compose([ transforms.RandomCrop(128), transforms.ToTensor() ]) def __len__(self): return len(self.hr_images) def __getitem__(self, idx): img = Image.open(self.hr_images[idx]).convert("RGB") hr = self.hr_transform(img) lr = transforms.Resize( hr.shape[1]//self.scale, interpolation=Image.BICUBIC )(transforms.ToPILImage()(hr)) lr = transforms.Resize( hr.shape[1], interpolation=Image.BICUBIC )(lr) lr = transforms.ToTensor()(lr) return lr, hr # ============================================ # RRDB Block # ============================================ class DenseBlock(nn.Module): def __init__(self, nf=64, gc=32): super().__init__() self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1) self.conv2 = nn.Conv2d(nf+gc, gc, 3, 1, 1) self.conv3 = nn.Conv2d(nf+2*gc, gc, 3, 1, 1) self.conv4 = nn.Conv2d(nf+3*gc, gc, 3, 1, 1) self.conv5 = nn.Conv2d(nf+4*gc, nf, 3, 1, 1) self.lrelu = nn.LeakyReLU(0.2) def forward(self,x): x1 = self.lrelu(self.conv1(x)) x2 = self.lrelu(self.conv2(torch.cat([x,x1],1))) x3 = self.lrelu(self.conv3(torch.cat([x,x1,x2],1))) x4 = self.lrelu(self.conv4(torch.cat([x,x1,x2,x3],1))) x5 = self.conv5(torch.cat([x,x1,x2,x3,x4],1)) return x5*0.2 + x class RRDB(nn.Module): def __init__(self, nf): super().__init__() self.db1 = DenseBlock(nf) self.db2 = DenseBlock(nf) self.db3 = DenseBlock(nf) def forward(self,x): return self.db3(self.db2(self.db1(x))) * 0.2 + x # ============================================ # Generator # ============================================ class RRDBNet(nn.Module): def __init__(self, nf=64, nb=23): super().__init__() self.conv_first = nn.Conv2d(3, nf, 3,1,1) self.body = nn.Sequential(*[RRDB(nf) for _ in range(nb)]) self.conv_body = nn.Conv2d(nf,nf,3,1,1) self.up1 = nn.Conv2d(nf,nf,3,1,1) self.up2 = nn.Conv2d(nf,nf,3,1,1) self.conv_hr = nn.Conv2d(nf,nf,3,1,1) self.conv_last = nn.Conv2d(nf,3,3,1,1) self.lrelu = nn.LeakyReLU(0.2) def forward(self,x): fea = self.conv_first(x) trunk = self.conv_body(self.body(fea)) fea = fea + trunk fea = self.lrelu( self.up1( nn.functional.interpolate(fea, scale_factor=2) ) ) fea = self.lrelu( self.up2( nn.functional.interpolate(fea, scale_factor=2) ) ) out = self.conv_last(self.lrelu(self.conv_hr(fea))) return out # ============================================ # Discriminator # ============================================ class Discriminator(nn.Module): def __init__(self): super().__init__() def block(in_c,out_c,stride): return nn.Sequential( nn.Conv2d(in_c,out_c,3,stride,1), nn.BatchNorm2d(out_c), nn.LeakyReLU(0.2) ) self.net = nn.Sequential( block(3,64,1), block(64,64,2), block(64,128,1), block(128,128,2), block(128,256,1), block(256,256,2), block(256,512,1), block(512,512,2), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(512,100), nn.LeakyReLU(0.2), nn.Linear(100,1) ) def forward(self,x): return self.net(x) # ============================================ # VGG Perceptual Loss # ============================================ class VGGFeatureExtractor(nn.Module): def __init__(self): super().__init__() vgg = vgg19(pretrained=True).features[:35] self.features = vgg.eval() for p in self.features.parameters(): p.requires_grad=False def forward(self,x): return self.features(x) # ============================================ # Training Setup # ============================================ dataset = ArtworkDataset("data/artwork_hr") loader = DataLoader( dataset, batch_size=8, shuffle=True ) G = RRDBNet().to(device) D = Discriminator().to(device) vgg = VGGFeatureExtractor().to(device) criterion_gan = nn.BCEWithLogitsLoss() criterion_l1 = nn.L1Loss() opt_g = optim.Adam(G.parameters(), lr=1e-4) opt_d = optim.Adam(D.parameters(), lr=1e-4) # ============================================ # Training Loop # ============================================ epochs = 100 for epoch in range(epochs): for i,(lr,hr) in enumerate(loader): lr = lr.to(device) hr = hr.to(device) valid = torch.ones(hr.size(0),1).to(device) fake = torch.zeros(hr.size(0),1).to(device) # ----------------- # Train Generator # ----------------- sr = G(lr) pred_fake = D(sr) loss_gan = criterion_gan(pred_fake,valid) loss_l1 = criterion_l1(sr,hr) loss_perc = criterion_l1( vgg(sr), vgg(hr) ) loss_g = loss_l1 + 0.01*loss_gan + 0.1*loss_perc opt_g.zero_grad() loss_g.backward() opt_g.step() # ----------------- # Train Discriminator # ----------------- pred_real = D(hr) loss_real = criterion_gan(pred_real,valid) pred_fake = D(sr.detach()) loss_fake = criterion_gan(pred_fake,fake) loss_d = (loss_real + loss_fake)/2 opt_d.zero_grad() loss_d.backward() opt_d.step() if i % 50 == 0: print( f"Epoch {epoch} Batch {i}", f"G Loss {loss_g.item():.4f}", f"D Loss {loss_d.item():.4f}" ) # Save sample vutils.save_image(sr,"sample_sr.png") # ============================================ # Save model # ============================================ torch.save(G.state_dict(),"esrgan_artwork.pth")