This repository provides code for Multi-Scale Discriminative Feature (MDF) loss for image reconstruction algorithms.
Central to the application of neural networks in image restoration problems, such as single image super resolution, is the choice of a loss function that encourages natural and perceptually pleasing results. We provide a lightweight feature extractor that outperforms state-of-the-art loss functions in single image super resolution, denoising, and JPEG artefact removal. We propose a novel Multi-Scale Discriminative Feature (MDF) loss comprising a series of discriminators, trained to penalize errors introduced by a generator. For further information please refer to the project webpage.
The code runs in Python3 and Pytorch.
First install the dependencies by running:
pip3 install -r requirements.txt
To run a simple example, optimizing image pixels:
import torch as pt
import torch.optim as optim
import imageio
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable
from mdfloss import MDFLoss
# Set parameters
cuda_available = False
epochs = 25
application = 'Denoising'
image_path = './misc/i10.png'
if application =='SISR':
path_disc = "./weights/Ds_SISR.pth"
elif application == 'Denoising':
path_disc = "./weights/Ds_Denoising.pth"
elif application == 'JPEG':
path_disc = "./weights/Ds_JPEG.pth"
# Read reference images
imgr = imageio.imread(image_path)
imgr = pt.from_numpy(imageio.core.asarray(imgr/255.0))
imgr = imgr.type(dtype=pt.float64)
imgr = imgr.permute(2,0,1)
imgr = imgr.unsqueeze(0).type(pt.FloatTensor)
# Create a noisy image
imgd = pt.rand(imgr.size())
if cuda_available:
imgr = imgr.cuda()
imgd = imgd.cuda()
# Convert images to variables to support gradients
imgrb = Variable( imgr, requires_grad = False)
imgdb = Variable( imgd, requires_grad = True)
optimizer = optim.Adam([imgdb], lr=0.1)
# Initialise the loss
criterion = MDFLoss(path_disc, cuda_available=cuda_available)
#Â Iterate over the epochs optimizing for the noisy image
for ii in range(0,epochs):
optimizer.zero_grad()
loss = criterion(imgrb,imgdb)
print("Epoch: ",ii," loss: ", loss.item())
loss.backward()
optimizer.step()If using, please cite:
@inproceedings{mustafa2022training,
title={Training a Task-Specific Image Reconstruction Loss},
author={Mustafa, Aamir and Mikhailiuk, Aliaksei and Iliescu, Dan Andrei and Babbar, Varun and Mantiuk, Rafa{\l} K},
booktitle={Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision},
pages={2319--2328},
year={2022}
}
This project has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement N◦ 725253–EyeCode).