import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from torch.utils.data import DataLoader
from data.hyperkvasir import KvasirInpaintingDataset
from models.inpainters import SegGenerator, SegDiscriminator
from perturbation.polyp_inpainter import Inpainter
# TODO refactor
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find("BatchNorm2d") != -1:
torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
torch.nn.init.constant_(m.bias.data, 0.0)
def train_new_inpainter():
# Loss function
adversarial_loss = torch.nn.BCELoss()
pixelwise_loss = torch.nn.L1Loss()
# Initialize generator and discriminator
# generator = Generator(channels=3)
# discriminator = Discriminator(channels=3)
generator = SegGenerator()
discriminator = SegDiscriminator()
generator.load_state_dict(torch.load("Predictors/Inpainters/no-pretrain-deeplab-generator-940"))
discriminator.load_state_dict(torch.load("Predictors/Inpainters/no-pretrain-deeplab-discriminator-940"))
cuda = True
if cuda:
generator.cuda()
discriminator.cuda()
adversarial_loss.cuda()
pixelwise_loss.cuda()
# Dataset loader TODO refactor w/ albumentation library
transforms_ = [
transforms.Resize((400, 400), Image.BICUBIC),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
dataloader = DataLoader(
KvasirInpaintingDataset("Datasets/HyperKvasir"),
batch_size=8,
shuffle=False,
num_workers=1,
)
# test_dataloader = DataLoader(
# EtisDataset("Datasets/ETIS-LaribPolypDB"),
# batch_size=12,
# shuffle=True,
# num_workers=1,
# )
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=0.0001)
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=0.00001)
scheduler_G = CosineAnnealingWarmRestarts(optimizer_G, T_0=100, T_mult=2)
scheduler_D = CosineAnnealingWarmRestarts(optimizer_D, T_0=100, T_mult=2)
# Initialize weights
# generator.apply(weights_init_normal)
# discriminator.apply(weights_init_normal)
# patch_h, patch_w = int(50 / 2 ** 3), int(50 / 2 ** 3)
# patch = (1, patch_h, patch_w)
# print(patch)
for epoch in range(990, 5000):
printed = False
d_losses = []
g_advs = []
g_pixels = []
for i, (imgs, mask, masked_imgs, masked_parts, filename) in enumerate(dataloader):
imgs = imgs.cuda()
mask = mask.cuda()
masked_imgs = masked_imgs.cuda()
masked_parts = masked_parts.cuda()
mask_bool = mask == 1
# Adversarial ground truths (boxes)
# valid = Variable(torch.Tensor(imgs.shape[0], *patch).fill_(1.0), requires_grad=False)
# fake = Variable(torch.Tensor(imgs.shape[0], *patch).fill_(0.0), requires_grad=False)
valid = torch.masked_select(torch.ones_like(mask), mask_bool)
fake = torch.masked_select(torch.zeros_like(mask), mask_bool)
# Configure input
imgs = Variable(imgs)
masked_imgs = Variable(masked_imgs)
masked_parts = Variable(masked_parts)
# -----------------
# Train Generator
# -----------------
optimizer_G.zero_grad()
# Generate a batch of images
gen_parts = generator(masked_imgs)
# Adversarial and pixelwise loss
disc = discriminator(gen_parts)
# print(disc)
g_adv = adversarial_loss(torch.masked_select(disc, mask_bool), valid)
g_pixel = pixelwise_loss(torch.masked_select(gen_parts, mask_bool), torch.masked_select(imgs, mask_bool))
g_advs.append(g_adv.item())
g_pixels.append(g_pixel.item())
# Total loss
g_loss = 0.001 * g_adv + 0.999 * g_pixel
g_loss.backward()
optimizer_G.step()
scheduler_G.step(epoch)
# ---------------------
# Train Discriminator
# ---------------------
optimizer_D.zero_grad()
# Measure discriminator's ability to classify real from generated samples
real_loss = adversarial_loss(torch.masked_select(discriminator(masked_parts), mask_bool), valid)
fake_loss = adversarial_loss(torch.masked_select(discriminator(gen_parts.detach()), mask_bool), fake)
d_loss = 0.5 * (real_loss + fake_loss)
d_losses.append(d_loss.item())
# wasserstein critic loss
# d_loss = -torch.mean(discriminator(masked_parts)) + torch.mean(discriminator(gen_parts.detach()))
d_loss.backward()
optimizer_D.step()
scheduler_D.step(epoch)
if not printed and epoch % 10 == 0:
torch.save(generator.state_dict(), f"Predictors/Inpainters/no-pretrain-deeplab-generator-{epoch}")
torch.save(discriminator.state_dict(),
f"Predictors/Inpainters/no-pretrain-deeplab-discriminator-{epoch}")
plt.title("Part")
plt.imshow((gen_parts[0].detach().cpu().numpy().T))
plt.show()
# plt.title("Superimposed")
# plt.imshow((gen_parts[0].detach().cpu().numpy().T))
# plt.imshow(masked_imgs[0].detach().cpu().numpy().T)
# plt.show()
# plt.title("Real")
# plt.imshow(masked_parts[0].detach().cpu().numpy().T)
# plt.show()
try:
test = Inpainter(f"Predictors/Inpainters/no-pretrain-deeplab-generator-{epoch}")
test.get_test()
except FileNotFoundError:
print("Weird...")
printed = True
print(
f"[Epoch {epoch}] [D loss: {np.mean(d_losses)}] [G adv: {np.mean(g_advs)}, pixel: {np.mean(g_pixels)}]"
)
if __name__ == '__main__':
train_new_inpainter()
Datasets
Models
