hallu.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 # ImageMagick's montage to make the mosaic
   9 #
  10 # montage hallu-*.png -tile 5x6 -geometry +1+1 result.png
  11
  12 import PIL, torch, torchvision
  13 from torch.nn import functional as F
  14
  15
  16 class MultiScaleEdgeEnergy(torch.nn.Module):
  17     def __init__(self):
  18         super().__init__()
  19         k = torch.exp(-torch.tensor([[-2.0, -1.0, 0.0, 1.0, 2.0]]) ** 2 / 2)
  20         k = (k.t() @ k).view(1, 1, 5, 5)
  21         self.gaussian_5x5 = torch.nn.Parameter(k / k.sum()).requires_grad_(False)
  22
  23     def forward(self, x):
  24         u = x.view(-1, 1, x.size(2), x.size(3))
  25         result = 0.0
  26         while min(u.size(2), u.size(3)) > 5:
  27             blurry = F.conv2d(u, self.gaussian_5x5, padding=2)
  28             result += (u - blurry).view(u.size(0), -1).pow(2).sum(1)
  29             u = F.avg_pool2d(u, kernel_size=2, padding=1)
  30         return result.view(x.size(0), -1).sum(1)
  31
  32
  33 img = torchvision.transforms.ToTensor()(PIL.Image.open("blacklab.jpg"))
  34 img = img.view((1,) + img.size())
  35 ref_input = 0.5 + 0.5 * (img - img.mean()) / img.std()
  36
  37 mse_loss = torch.nn.MSELoss()
  38 edge_energy = MultiScaleEdgeEnergy()
  39
  40 layers = torchvision.models.vgg16(pretrained=True).features
  41 layers.eval()
  42
  43 if torch.cuda.is_available():
  44     edge_energy.cuda()
  45     ref_input = ref_input.cuda()
  46     layers.cuda()
  47
  48 for l in [5, 7, 12, 17, 21, 28]:
  49     model = torch.nn.Sequential(layers[:l])
  50     ref_output = model(ref_input).detach()
  51
  52     for n in range(5):
  53         input = torch.empty_like(ref_input).uniform_(-0.01, 0.01).requires_grad_()
  54         optimizer = torch.optim.Adam([input], lr=1e-2)
  55         for k in range(1000):
  56             output = model(input)
  57             loss = mse_loss(output, ref_output) + 1e-3 * edge_energy(input)
  58             optimizer.zero_grad()
  59             loss.backward()
  60             optimizer.step()
  61
  62         img = 0.5 + 0.2 * (input - input.mean()) / input.std()
  63         result_name = "hallu-l%02d-n%02d.png" % (l, n)
  64         torchvision.utils.save_image(img, result_name)
  65
  66         print("Wrote " + result_name)