Update.

[pytorch.git] / hallu.py

diff --git a/hallu.py b/hallu.py
index 6b0b303..9d2706d 100755 (executable)
--- a/hallu.py
+++ b/hallu.py
@@ -1,5 +1,10 @@
 #!/usr/bin/env python
 
+# Any copyright is dedicated to the Public Domain.
+# https://creativecommons.org/publicdomain/zero/1.0/
+
+# Written by Francois Fleuret <francois@fleuret.org>
+
 # ImageMagick's montage to make the mosaic
 #
 # montage hallu-*.png -tile 5x6 -geometry +1+1 result.png
@@ -7,30 +12,32 @@
 import PIL, torch, torchvision
 from torch.nn import functional as F
 
+
 class MultiScaleEdgeEnergy(torch.nn.Module):
     def __init__(self):
-        super(MultiScaleEdgeEnergy, self).__init__()
-        k = torch.exp(- torch.tensor([[-2., -1., 0., 1., 2.]])**2 / 2)
+        super().__init__()
+        k = torch.exp(-torch.tensor([[-2.0, -1.0, 0.0, 1.0, 2.0]]) ** 2 / 2)
         k = (k.t() @ k).view(1, 1, 5, 5)
-        self.register_buffer('gaussian_5x5', k / k.sum())
+        self.gaussian_5x5 = torch.nn.Parameter(k / k.sum()).requires_grad_(False)
 
     def forward(self, x):
         u = x.view(-1, 1, x.size(2), x.size(3))
         result = 0.0
         while min(u.size(2), u.size(3)) > 5:
-            blurry  = F.conv2d(u, self.gaussian_5x5, padding = 2)
+            blurry = F.conv2d(u, self.gaussian_5x5, padding=2)
             result += (u - blurry).view(u.size(0), -1).pow(2).sum(1)
-            u = F.avg_pool2d(u, kernel_size = 2, padding = 1)
+            u = F.avg_pool2d(u, kernel_size=2, padding=1)
         return result.view(x.size(0), -1).sum(1)
 
-img = torchvision.transforms.ToTensor()(PIL.Image.open('blacklab.jpg'))
+
+img = torchvision.transforms.ToTensor()(PIL.Image.open("blacklab.jpg"))
 img = img.view((1,) + img.size())
 ref_input = 0.5 + 0.5 * (img - img.mean()) / img.std()
 
 mse_loss = torch.nn.MSELoss()
 edge_energy = MultiScaleEdgeEnergy()
 
-layers = torchvision.models.vgg16(pretrained = True).features
+layers = torchvision.models.vgg16(pretrained=True).features
 layers.eval()
 
 if torch.cuda.is_available():
@@ -38,13 +45,13 @@ if torch.cuda.is_available():
     ref_input = ref_input.cuda()
     layers.cuda()
 
-for l in [ 5, 7, 12, 17, 21, 28 ]:
+for l in [5, 7, 12, 17, 21, 28]:
     model = torch.nn.Sequential(layers[:l])
     ref_output = model(ref_input).detach()
 
     for n in range(5):
-        input = ref_input.new_empty(ref_input.size()).uniform_(-0.01, 0.01).requires_grad_()
-        optimizer = torch.optim.Adam( [ input ], lr = 1e-2)
+        input = torch.empty_like(ref_input).uniform_(-0.01, 0.01).requires_grad_()
+        optimizer = torch.optim.Adam([input], lr=1e-2)
         for k in range(1000):
             output = model(input)
             loss = mse_loss(output, ref_output) + 1e-3 * edge_energy(input)
@@ -53,7 +60,7 @@ for l in [ 5, 7, 12, 17, 21, 28 ]:
             optimizer.step()
 
         img = 0.5 + 0.2 * (input - input.mean()) / input.std()
-        result_name = 'hallu-l%02d-n%02d.png' % (l, n)
+        result_name = "hallu-l%02d-n%02d.png" % (l, n)
         torchvision.utils.save_image(img, result_name)
 
-        print('Wrote ' + result_name)
+        print("Wrote " + result_name)