dev(narugo): add noisy models

docs/source/api_doc/restore/restore_demo.plot.py

@@ -0,0 +1,17 @@
+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),
+    )

imgutils/restore/__init__.py

@@ -1,6 +1,15 @@
 """
 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

test/restore/conftest.py

@@ -0,0 +1,41 @@
+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)

test/restore/test_nafnet.py

@@ -0,0 +1,22 @@
+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

test/restore/test_scunet.py

@@ -0,0 +1,25 @@
+import pytest
+
+from imgutils.metrics import psnr
+from imgutils.restore import restore_with_scunet
+from imgutils.restore.scunet import _open_scunet_model
+
+
+@pytest.fixture(autouse=True, scope='module')
+def _clear_cache():
+    try:
+        yield
+    finally:
+        _open_scunet_model.cache_clear()
+
+
+@pytest.mark.unittest
+class TestRestoreSCUNet:
+    def test_restore_with_scunet_original(self, sample_image, clear_image):
+        assert psnr(restore_with_scunet(sample_image), clear_image) >= 34.5
+
+    def test_restore_with_scunet_q45(self, q45_image, clear_image):
+        assert psnr(restore_with_scunet(q45_image), clear_image) >= 34.5
+
+    def test_restore_with_scunet_gnoise(self, gaussian_noise_image, clear_image):
+        assert psnr(restore_with_scunet(gaussian_noise_image), clear_image) >= 33