jorge-pessoa / pytorch-msssim

PyTorch differentiable Multi-Scale Structural Similarity (MS-SSIM) loss
Other
453 stars 70 forks source link

Artifacts on color images #4

Closed trougnouf closed 4 years ago

trougnouf commented 5 years ago

Running the test script generates corrupted results whereas pytorch-ssim yields no such issue. Tested with

totensor = torchvision.transforms.ToTensor()
img1 = totensor(Image.open('kate.png'))
img1 = img1.reshape([1]+list(img1.shape))
...
torchvision.utils.save_image(img2, 'res.jpg')

kate res

jorge-pessoa commented 5 years ago

Would you mind providing a complete example of the issue you are facing?

I simplified the test script to accommodate your example and the output did not display the artifacts contained in your initial post. The simplified optimization script is as follows:

from pytorch_msssim import msssim
import torch
from torch import optim
import numpy as np
import torchvision
from PIL import Image

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

totensor = torchvision.transforms.ToTensor()
img1 = totensor(Image.open('kate.png'))
img1 = img1.reshape([1]+list(img1.shape))
img2 = torch.rand(img1.size())
img2 = torch.nn.functional.sigmoid(img2)

img1 = img1.to(device)
img2 = img2.to(device)

img1.requires_grad = False
img2.requires_grad = True

value = msssim(img1, img2)
optimizer = optim.Adam([img2], lr=0.01)
print("Initial MS-SSIM: %.5f" % (value.item()))

while value < 0.999:
    optimizer.zero_grad()
    msssim_out = -msssim(img1, img2)
    value = -msssim_out.item()
    print('Current MS-SSIM = %.5f' % value)
    msssim_out.backward()
    optimizer.step()

torchvision.utils.save_image(img2, 'out.png')
trougnouf commented 5 years ago

I don't get the artifacts when I use the code you provided here, but they come back when I comment out the sigmoid function. It doesn't seem correct to use it because torch.rand returns uniformly distributed numbers between 0 and 1 which limits sigmoid to >=0.5

>>> img1.min(),img1.max()
(tensor(0.), tensor(0.6549))
>>> img2.min(),img2.max()
(tensor(1.2696e-05), tensor(1.0000))
>>> img2 = torch.nn.functional.sigmoid(img2)
>>> img2.min(),img2.max()
(tensor(0.5000), tensor(0.7311))
jorge-pessoa commented 5 years ago

While that would be true before optimization, the sigmoid application guarantees that the output of optimization is always in the interval [0,1], thus describing a valid image (without artifacts). The minimum value of 0.5 will decrease during optimization to reach the intended value. You could just shift the initial uniform sampling in order to begin the optimization with a [0,1] interval after the sigmoid if you prefer.

If you look at the minimum and maximum values during optimization for your version (without sigmoid) you should see that the minimum value drops below 0, creating the reported artifacts in the initial issue. This is an expected consequence of the optimization procedure used.

Let me know if this solved your concerns, or otherwise close the issue. Thank you for your feedback!

trougnouf commented 5 years ago

The output values are indeed outside the range [0,1] and clamping the output doesn't remove the artifacts. I don't think that using sigmoid on the random image is a viable option because the loss function is meant to be used with the output of a neural network, which initially outputs random garbage before it learns to generate realistic images. This loss function doesn't seem to discriminate against these strong artifacts, I was only able to get rid of them using value < 0.99997 (which is far higher than the score my denoising application can achieve), so I suspect that the network would never learn to get rid of these artifacts once they appear in the first few itterations. Do you think not being able to detect these artifacts is an inherent issue with MS-SSIM or a bug?

jorge-pessoa commented 4 years ago

Hello, I apologize for the delay since my last answer. I cannot guarantee that this problem is an inherent MS-SSIM issue and not an implementation introduced error but I was not able to find or reproduce your issues. If your problem still persists, please feel free to re-open a new issue where we can actively discuss this. I will be closing this one due to age.

Regarding stability for the optimization, I will introduce a check in the value for the ssim and mc calculations in each scale iteration using a relu to ensure both values are always > 0, based on the implementation by https://github.com/VainF.