dev(narugo): add noisy models

from imgutils.data import load_image

from imgutils.restore.nafnet import restore_with_nafnet

from imgutils.restore.scunet import restore_with_scunet

from plot import image_plot

if __name__ == '__main__':

    img_q35 = load_image('sample/jpg-q35.jpg')

    img_gnoise = load_image('sample/gnoise.png')

    image_plot(

        (img_q35, 'JPEG Quality35'),

        (restore_with_nafnet(img_q35), 'JPEG Quality35\n(Fixed By NafNet)'),

        (img_gnoise, 'Gaussian Noise'),

        (restore_with_scunet(img_gnoise), 'Gaussian Noise\n(Fixed By SCUNet)'),

        columns=2,

        figsize=(6, 8),

    )

"""

Overview:

    Image restoration models.

    Utilities for restoring images, which may be jpeg, blurry or noisy.

    The following models are used:

    * `NafNet <https://github.com/megvii-research/NAFNet>`_

    * `SCUNet <https://github.com/cszn/SCUNet>`_

    .. image:: restore_demo.plot.py.svg

        :align: center

"""

from .nafnet import restore_with_nafnet

from .scunet import restore_with_scunet

import os.path

import numpy as np

import pytest

from PIL import Image

from hbutils.system import TemporaryDirectory

from imgutils.data import load_image

from test.testings import get_testfile

@pytest.fixture()

def sample_image():

    yield load_image(get_testfile('surtr_logo.png'), mode='RGB', force_background='white')

@pytest.fixture()

def clear_image():

    yield load_image(get_testfile('surtr_logo_clear.png'), mode='RGB', force_background='white')

def add_gaussian_noise(image, mean=0, std=25):

    img_array = np.array(image)

    noise = np.random.normal(mean, std, img_array.shape)

    noisy_image = img_array + noise

    noisy_image = np.clip(noisy_image, 0, 255)

    noisy_image = Image.fromarray(np.uint8(noisy_image))

    return noisy_image

@pytest.fixture()

def gaussian_noise_image(sample_image):

    yield add_gaussian_noise(sample_image)

@pytest.fixture()

def q45_image(sample_image):

    with TemporaryDirectory() as td:

        img_file = os.path.join(td, 'image.jpg')

        sample_image.save(img_file, quality=45)

        yield load_image(img_file)

import pytest

from imgutils.metrics import psnr

from imgutils.restore import restore_with_nafnet

from imgutils.restore.nafnet import _open_nafnet_model

@pytest.fixture(autouse=True, scope='module')

def _clear_cache():

    try:

        yield

    finally:

        _open_nafnet_model.cache_clear()

@pytest.mark.unittest

class TestRestoreNafNet:

    def test_restore_with_nafnet_original(self, sample_image, clear_image):

        assert psnr(restore_with_nafnet(sample_image), clear_image) >= 40.0

    def test_restore_with_nafnet_q45(self, q45_image, clear_image):

        assert psnr(restore_with_nafnet(q45_image), clear_image) >= 40.0