Add Docker environment & web demo (#67)

* enable cog

* Update README.md

* Update README.md

* refactor

* fix temp input dir bug

Co-authored-by: CJWBW <[email protected]>
Co-authored-by: Chenxi <[email protected]>
Co-authored-by: Xintao <[email protected]>

README.md

@@ -26,7 +26,6 @@ It leverages rich and diverse priors encapsulated in a pretrained face GAN (*e.g
 :question: Frequently Asked Questions can be found in [FAQ.md](FAQ.md).
 
 :triangular_flag_on_post: **Updates**
-
 - :fire::fire::white_check_mark: Add **[V1.3 model](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth)**, which produces **more natural** restoration results, and better results on *very low-quality* / *high-quality* inputs. See more in [Model zoo](#european_castle-model-zoo), [Comparisons.md](Comparisons.md)
 - :white_check_mark: Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/GFPGAN).
 - :white_check_mark: Support enhancing non-face regions (background) with [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN).

cog.yaml

@@ -0,0 +1,25 @@
+build:
+  gpu: true
+  python_version: "3.8"
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+  python_packages:
+    - "torch==1.8.0"
+    - "torchvision==0.9.0"
+    - "numpy==1.21.1"
+    - "ipython==7.21.0"
+    - "lmdb==1.2.1"
+    - "opencv-python==4.5.3.56"
+    - "PyYAML==5.4.1"
+    - "tqdm==4.62.2"
+    - "yapf==0.31.0"
+    - "tb-nightly==2.7.0a20210825"
+  run:
+    - pip install facexlib==0.2.0.2
+    - pip install realesrgan
+
+predict: "predict.py:Predictor"
+
+
+

download-weights

@@ -0,0 +1,3 @@
+#!/bin/sh
+
+wget https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth -P experiments/pretrained_models

predict.py

@@ -0,0 +1,144 @@
+import subprocess
+
+subprocess.call(["sh", "./run_setup.sh"])
+
+import warnings
+import tempfile
+import os
+from pathlib import Path
+import argparse
+import glob
+
+import shutil
+from basicsr.utils import imwrite
+import torch
+import cv2
+import cog
+from realesrgan import RealESRGANer
+from gfpgan import GFPGANer
+
+
+class Predictor(cog.Predictor):
+    def setup(self):
+        parser = argparse.ArgumentParser()
+
+        parser.add_argument("--upscale", type=int, default=2)
+        parser.add_argument("--arch", type=str, default="clean")
+        parser.add_argument("--channel", type=int, default=2)
+        parser.add_argument(
+            "--model_path",
+            type=str,
+            default="experiments/pretrained_models/GFPGANCleanv1-NoCE-C2.pth",
+        )
+        parser.add_argument("--bg_upsampler", type=str, default="realesrgan")
+        parser.add_argument("--bg_tile", type=int, default=400)
+        parser.add_argument("--test_path", type=str, default="inputs/whole_imgs")
+        parser.add_argument(
+            "--suffix", type=str, default=None, help="Suffix of the restored faces"
+        )
+        parser.add_argument("--only_center_face", action="store_true")
+        parser.add_argument("--aligned", action="store_true")
+        parser.add_argument("--paste_back", action="store_false")
+        parser.add_argument("--save_root", type=str, default="results")
+
+        self.args = parser.parse_args(
+            ["--upscale", "2", "--test_path", "cog_temp", "--save_root", "results"]
+        )
+        os.makedirs(self.args.test_path, exist_ok=True)
+        # background upsampler
+        if self.args.bg_upsampler == "realesrgan":
+            if not torch.cuda.is_available():  # CPU
+
+                warnings.warn(
+                    "The unoptimized RealESRGAN is very slow on CPU. We do not use it. "
+                    "If you really want to use it, please modify the corresponding codes."
+                )
+                bg_upsampler = None
+            else:
+                bg_upsampler = RealESRGANer(
+                    scale=2,
+                    model_path="https://github.com/xinntao/Real-ESRGAN/releases"
+                    "/download/v0.2.1/RealESRGAN_x2plus.pth",
+                    tile=self.args.bg_tile,
+                    tile_pad=10,
+                    pre_pad=0,
+                    half=True,
+                )  # need to set False in CPU mode
+        else:
+            bg_upsampler = None
+
+        # set up GFPGAN restorer
+        self.restorer = GFPGANer(
+            model_path=self.args.model_path,
+            upscale=self.args.upscale,
+            arch=self.args.arch,
+            channel_multiplier=self.args.channel,
+            bg_upsampler=bg_upsampler,
+        )
+
+    @cog.input("image", type=Path, help="input image")
+    def predict(self, image):
+        try:
+            input_dir = self.args.test_path
+
+            input_path = os.path.join(input_dir, os.path.basename(image))
+            shutil.copy(str(image), input_path)
+
+            os.makedirs(self.args.save_root, exist_ok=True)
+
+            img_list = sorted(glob.glob(os.path.join(input_dir, "*")))
+
+            out_path = Path(tempfile.mkdtemp()) / "output.png"
+
+            for img_path in img_list:
+                # read image
+                img_name = os.path.basename(img_path)
+                print(f"Processing {img_name} ...")
+                basename, ext = os.path.splitext(img_name)
+                input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+
+                cropped_faces, restored_faces, restored_img = self.restorer.enhance(
+                    input_img,
+                    has_aligned=self.args.aligned,
+                    only_center_face=self.args.only_center_face,
+                    paste_back=self.args.paste_back,
+                )
+
+                imwrite(restored_img, str(out_path))
+                clean_folder(self.args.test_path)
+
+                # save faces
+                for idx, (cropped_face, restored_face) in enumerate(
+                    zip(cropped_faces, restored_faces)
+                ):
+                    # save cropped face
+                    save_crop_path = os.path.join(
+                        self.args.save_root, "cropped_faces", f"{basename}_{idx:02d}.png"
+                    )
+                    imwrite(cropped_face, save_crop_path)
+                    # save restored face
+                    if self.args.suffix is not None:
+                        save_face_name = f"{basename}_{idx:02d}_{self.args.suffix}.png"
+                    else:
+                        save_face_name = f"{basename}_{idx:02d}.png"
+                    save_restore_path = os.path.join(
+                        self.args.save_root, "restored_faces", save_face_name
+                    )
+                    imwrite(restored_face, save_restore_path)
+                    imwrite(restored_img, str(out_path))
+        finally:
+            clean_folder(self.args.test_path)
+
+        return out_path
+
+
+def clean_folder(folder):
+    for filename in os.listdir(folder):
+        file_path = os.path.join(folder, filename)
+        try:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+                os.unlink(file_path)
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)
+        except Exception as e:
+            print("Failed to delete %s. Reason: %s" % (file_path, e))

run_setup.sh

@@ -0,0 +1,2 @@
+#!/bin/sh
+python setup.py develop