Merge pull request #46 from deepghs/dev/restore

dev(narugo): add image restore models (NafNet and SCUNet)

.github/workflows/waifu2x.yml

@@ -0,0 +1,44 @@
+name: Sync Waifu2x Models
+
+on:
+  #  push:
+  workflow_dispatch:
+  schedule:
+    - cron: '30 16 * * *'
+
+jobs:
+  sync:
+    name: Sync Waifu2x ONNX
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+          - 'ubuntu-latest'
+        python-version:
+          - '3.8'
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 20
+      - name: Set up python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Set up python dependences
+        run: |
+          pip install --upgrade pip
+          pip install --upgrade flake8 setuptools wheel twine
+          if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+          if [ -f requirements-build.txt ]; then pip install -r requirements-build.txt; fi
+          if [ -f requirements-test.txt ]; then pip install -r requirements-test.txt; fi
+          if [ -f requirements-test.txt ]; then pip install -r requirements-zoo.txt; fi
+          pip install --upgrade build
+      - name: Sync Models
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          GH_ACCESS_TOKEN: ${{ secrets.GH_ACCESS_TOKEN }}
+        run: |
+          python -m zoo.waifu2x sync

docs/source/api_doc/restore/index.rst

@@ -0,0 +1,14 @@
+imgutils.restore
+========================
+
+.. currentmodule:: imgutils.restore
+
+.. automodule:: imgutils.restore
+
+
+.. toctree::
+    :maxdepth: 3
+
+    nafnet
+    scunet
+

docs/source/api_doc/restore/nafnet.rst

@@ -0,0 +1,14 @@
+imgutils.restore.nafnet
+====================================
+
+.. currentmodule:: imgutils.restore.nafnet
+
+.. automodule:: imgutils.restore.nafnet
+
+
+restore_with_nafnet
+------------------------
+
+.. autofunction:: restore_with_nafnet
+
+

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

@@ -0,0 +1,35 @@
+import random
+
+from benchmark import BaseBenchmark, create_plot_cli
+from imgutils.restore.nafnet import NafNetModelTyping, restore_with_nafnet
+
+
+class NafNetBenchmark(BaseBenchmark):
+    def __init__(self, model: NafNetModelTyping):
+        BaseBenchmark.__init__(self)
+        self.model = model
+
+    def load(self):
+        from imgutils.restore.nafnet import _open_nafnet_model
+        _open_nafnet_model(self.model)
+
+    def unload(self):
+        from imgutils.restore.nafnet import _open_nafnet_model
+        _open_nafnet_model.cache_clear()
+
+    def run(self):
+        image_file = random.choice(self.all_images)
+        _ = restore_with_nafnet(image_file, model=self.model)
+
+
+if __name__ == '__main__':
+    create_plot_cli(
+        [
+            ('NafNet-REDS', NafNetBenchmark('REDS')),
+            ('NafNet-GoPro', NafNetBenchmark('GoPro')),
+            ('NafNet-SIDD', NafNetBenchmark('SIDD')),
+        ],
+        title='Benchmark for NafNet Models',
+        run_times=5,
+        try_times=10,
+    )()