Face Inpainting and Selfie

.gitattributes

@@ -25,3 +25,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+inpainting.png filter=lfs diff=lfs merge=lfs -text

GPEN.ipynb

@@ -1,343 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "7TLDy5mi1rid"
-      },
-      "source": [
-        "# Run once only!"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "nG64GBsoyoMH"
-      },
-      "outputs": [],
-      "source": []
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "--fxPES99Cp0",
-        "outputId": "aa291a7e-68b8-4115-f177-fad3ce43a8fd"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Cloning into 'GPEN'...\n",
-            "remote: Enumerating objects: 142, done.\u001b[K\n",
-            "remote: Counting objects: 100% (142/142), done.\u001b[K\n",
-            "remote: Compressing objects: 100% (123/123), done.\u001b[K\n",
-            "remote: Total 142 (delta 21), reused 127 (delta 12), pack-reused 0\u001b[K\n",
-            "Receiving objects: 100% (142/142), 22.05 MiB | 39.74 MiB/s, done.\n",
-            "Resolving deltas: 100% (21/21), done.\n"
-          ]
-        }
-      ],
-      "source": [
-        "! git clone https://github.com/KenjieDec/GPEN GPEN"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "t8AU9f7U-vKz",
-        "outputId": "b2489822-858c-4e26-82c4-9da16573501c"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "--2021-06-13 05:26:36--  https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth\n",
-            "Resolving public-vigen-video.oss-cn-shanghai.aliyuncs.com (public-vigen-video.oss-cn-shanghai.aliyuncs.com)... 47.101.88.25\n",
-            "Connecting to public-vigen-video.oss-cn-shanghai.aliyuncs.com (public-vigen-video.oss-cn-shanghai.aliyuncs.com)|47.101.88.25|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 109497761 (104M) [application/octet-stream]\n",
-            "Saving to: ‘RetinaFace-R50.pth’\n",
-            "\n",
-            "RetinaFace-R50.pth  100%[===================>] 104.42M  9.53MB/s    in 17s     \n",
-            "\n",
-            "2021-06-13 05:26:55 (6.17 MB/s) - ‘RetinaFace-R50.pth’ saved [109497761/109497761]\n",
-            "\n",
-            "--2021-06-13 05:26:55--  https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-512.pth\n",
-            "Resolving public-vigen-video.oss-cn-shanghai.aliyuncs.com (public-vigen-video.oss-cn-shanghai.aliyuncs.com)... 47.101.88.25\n",
-            "Connecting to public-vigen-video.oss-cn-shanghai.aliyuncs.com (public-vigen-video.oss-cn-shanghai.aliyuncs.com)|47.101.88.25|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 284085738 (271M) [application/octet-stream]\n",
-            "Saving to: ‘GPEN-512.pth’\n",
-            "\n",
-            "GPEN-512.pth        100%[===================>] 270.92M  8.48MB/s    in 43s     \n",
-            "\n",
-            "2021-06-13 05:27:40 (6.30 MB/s) - ‘GPEN-512.pth’ saved [284085738/284085738]\n",
-            "\n",
-            "--2021-06-13 05:27:40--  https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-1024-Color.pth\n",
-            "Resolving public-vigen-video.oss-cn-shanghai.aliyuncs.com (public-vigen-video.oss-cn-shanghai.aliyuncs.com)... 47.101.88.25\n",
-            "Connecting to public-vigen-video.oss-cn-shanghai.aliyuncs.com (public-vigen-video.oss-cn-shanghai.aliyuncs.com)|47.101.88.25|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 284914645 (272M) [application/octet-stream]\n",
-            "Saving to: ‘GPEN-1024-Color.pth’\n",
-            "\n",
-            "GPEN-1024-Color.pth 100%[===================>] 271.71M  9.09MB/s    in 42s     \n",
-            "\n",
-            "2021-06-13 05:28:25 (6.43 MB/s) - ‘GPEN-1024-Color.pth’ saved [284914645/284914645]\n",
-            "\n"
-          ]
-        }
-      ],
-      "source": [
-        "!wget \"https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116085&Signature=GlUNW6%2B8FxvxWmE9jKIZYOOciKQ%3D\" -O GPEN/weights/RetinaFace-R50.pth\n",
-        "!wget \"https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116208&Signature=hBgvVvKVSNGeXqT8glG%2Bd2t2OKc%3D\" -O GPEN/weights/GPEN-512.pth\n",
-        "!wget \"https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116315&Signature=9tPavW2h%2F1LhIKiXj73sTQoWqcc%3D\" -O GPEN/weights/GPEN-1024-Color.pth "
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "Zfc5z0CyAVKY",
-        "outputId": "98dc0eab-167c-4415-96f5-9e8af5bc2a6e"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Collecting torch==1.7.1\n",
-            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/90/5d/095ddddc91c8a769a68c791c019c5793f9c4456a688ddd235d6670924ecb/torch-1.7.1-cp37-cp37m-manylinux1_x86_64.whl (776.8MB)\n",
-            "\u001b[K     |████████████████████████████████| 776.8MB 24kB/s \n",
-            "\u001b[?25hCollecting torchvision==0.8.2\n",
-            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/94/df/969e69a94cff1c8911acb0688117f95e1915becc1e01c73e7960a2c76ec8/torchvision-0.8.2-cp37-cp37m-manylinux1_x86_64.whl (12.8MB)\n",
-            "\u001b[K     |████████████████████████████████| 12.8MB 264kB/s \n",
-            "\u001b[?25hCollecting torchaudio==0.7.2\n",
-            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/37/16/ecdb9eb09ec6b8133d6c9536ea9e49cd13c9b5873c8488b8b765a39028da/torchaudio-0.7.2-cp37-cp37m-manylinux1_x86_64.whl (7.6MB)\n",
-            "\u001b[K     |████████████████████████████████| 7.6MB 18.2MB/s \n",
-            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch==1.7.1) (1.19.5)\n",
-            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.7.1) (3.7.4.3)\n",
-            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.8.2) (7.1.2)\n",
-            "\u001b[31mERROR: torchtext 0.9.1 has requirement torch==1.8.1, but you'll have torch 1.7.1 which==incompatible.\u001b[0m\n",
-            "Installing collected packages: torch, torchvision, torchaudio\n",
-            "  Found existing installation: torch 1.8.1+cu101\n",
-            "    Uninstalling torch-1.8.1+cu101:\n",
-            "      Successfully uninstalled torch-1.8.1+cu101\n",
-            "  Found existing installation: torchvision 0.9.1+cu101\n",
-            "    Uninstalling torchvision-0.9.1+cu101:\n",
-            "      Successfully uninstalled torchvision-0.9.1+cu101\n",
-            "Successfully installed torch-1.7.1 torchaudio-0.7.2 torchvision-0.8.2\n",
-            "Requirement already satisfied: pip in /usr/local/lib/python3.7/dist-packages (19.3.1)\n",
-            "Requirement already satisfied: install in /usr/local/lib/python3.7/dist-packages (1.3.4)\n",
-            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (4.1.2.30)\n",
-            "Requirement already satisfied: numpy>=1.14.5 in /usr/local/lib/python3.7/dist-packages (from opencv-python) (1.19.5)\n"
-          ]
-        }
-      ],
-      "source": [
-        "!pip install torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2\n",
-        "!pip install pip install opencv-python"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "BlSfEFXgnAHu",
-        "outputId": "6502bfec-2efd-4e36-fef6-4fc354132ea2"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "/content/GPEN\n"
-          ]
-        }
-      ],
-      "source": [
-        "%cd GPEN"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "dt6hIEjyx7Nt",
-        "outputId": "13de37fd-6dca-437f-c892-be89560a83d6"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "--2021-06-13 05:31:09--  https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
-            "Resolving github.com (github.com)... 140.82.113.3\n",
-            "Connecting to github.com (github.com)|140.82.113.3|:443... connected.\n",
-            "HTTP request sent, awaiting response... 302 Found\n",
-            "Location: https://github-releases.githubusercontent.com/1335132/d2f252e2-9801-11e7-9fbf-bc7b4e4b5c83?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210613%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210613T053109Z&X-Amz-Expires=300&X-Amz-Signature=0297051ab9ba7cd75b20f001e6392f931bd0fbde6cea59a28c6b3bb7fdb97b78&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=1335132&response-content-disposition=attachment%3B%20filename%3Dninja-linux.zip&response-content-type=application%2Foctet-stream [following]\n",
-            "--2021-06-13 05:31:09--  https://github-releases.githubusercontent.com/1335132/d2f252e2-9801-11e7-9fbf-bc7b4e4b5c83?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20210613%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20210613T053109Z&X-Amz-Expires=300&X-Amz-Signature=0297051ab9ba7cd75b20f001e6392f931bd0fbde6cea59a28c6b3bb7fdb97b78&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=1335132&response-content-disposition=attachment%3B%20filename%3Dninja-linux.zip&response-content-type=application%2Foctet-stream\n",
-            "Resolving github-releases.githubusercontent.com (github-releases.githubusercontent.com)... 185.199.108.154, 185.199.109.154, 185.199.110.154, ...\n",
-            "Connecting to github-releases.githubusercontent.com (github-releases.githubusercontent.com)|185.199.108.154|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 77854 (76K) [application/octet-stream]\n",
-            "Saving to: ‘ninja-linux.zip’\n",
-            "\n",
-            "ninja-linux.zip     100%[===================>]  76.03K  --.-KB/s    in 0.002s  \n",
-            "\n",
-            "2021-06-13 05:31:09 (32.1 MB/s) - ‘ninja-linux.zip’ saved [77854/77854]\n",
-            "\n",
-            "Archive:  ninja-linux.zip\n",
-            "  inflating: /usr/local/bin/ninja    \n",
-            "update-alternatives: using /usr/local/bin/ninja to provide /usr/bin/ninja (ninja) in auto mode\n"
-          ]
-        }
-      ],
-      "source": [
-        "!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip\n",
-        "!sudo unzip ninja-linux.zip -d /usr/local/bin/\n",
-        "!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "MHS0GTSK1hnf"
-      },
-      "source": [
-        "# Main Code\n",
-        "\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "NOo9qIPb1Ivn",
-        "outputId": "268e81d3-bb69-42c0-a5ca-704cce5b37a8"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "0 Solvay_conference_1927.png\n"
-          ]
-        }
-      ],
-      "source": [
-        "!python face_enhancement.py"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "XeAs7bPg1KpZ",
-        "outputId": "a9e0b3de-d00b-4329-b079-2ed52a22f58a"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "0 examples/grays/106000_gray.png\n"
-          ]
-        }
-      ],
-      "source": [
-        "!python face_colorization.py"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "WDCm7xJee4c2"
-      },
-      "source": [
-        "# Download multiple images all at once"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6zC7BIgHe-Z2"
-      },
-      "source": [
-        "download on outs"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "4tAqGLPpe0Ez"
-      },
-      "outputs": [],
-      "source": [
-        "!zip -r /content/GPEN/examples/outs.zip /content/GPEN/examples/outs"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0W7jMJlqfD2R"
-      },
-      "source": [
-        "download on couts"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "B3-JMXM3e0wT"
-      },
-      "outputs": [],
-      "source": [
-        "!zip -r /content/GPEN/examples/couts.zip /content/GPEN/examples/couts"
-      ]
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "name": "GPEN.ipynb",
-      "provenance": []
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}

app.py

@@ -1,9 +1,12 @@
 import os
-
+ 
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116085&Signature=GlUNW6%2B8FxvxWmE9jKIZYOOciKQ%3D" -O weights/RetinaFace-R50.pth')
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116208&Signature=hBgvVvKVSNGeXqT8glG%2Bd2t2OKc%3D" -O weights/GPEN-512.pth')
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116315&Signature=9tPavW2h%2F1LhIKiXj73sTQoWqcc%3D" -O weights/GPEN-1024-Color.pth ')
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x2.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1962694780&Signature=lI%2FolhA%2FyigiTRvoDIVbtMIyhjI%3D" -O weights/realesrnet_x2.pth ')
+os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Inpainting-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116338&Signature=tvYhdLaLgW7UdcUrApXp2jsek8w%3D" -O weights/GPEN-Inpainting-1024.pth ')
+jksp= os.environ['GPEN-BFR-2048']
+os.system(f'wget "{jksp}" -O weights/GPEN-BFR-2048.pth')
 
 import gradio as gr
 
@@ -13,9 +16,66 @@ import gradio as gr
 '''
 import os
 import cv2
+import glob
+import time
+import math
+import argparse
+import numpy as np
+from PIL import Image, ImageDraw
+import __init_paths
 from face_enhancement import FaceEnhancement 
 from face_colorization import FaceColorization 
+from face_inpainting import FaceInpainting 
   
+def brush_stroke_mask(img, color=(255,255,255)):
+    min_num_vertex = 8
+    max_num_vertex = 28
+    mean_angle = 2*math.pi / 5
+    angle_range = 2*math.pi / 15
+    min_width = 12
+    max_width = 80
+    def generate_mask(H, W, img=None):
+        average_radius = math.sqrt(H*H+W*W) / 8
+        mask = Image.new('RGB', (W, H), 0)
+        if img is not None: mask = img #Image.fromarray(img)
+
+        for _ in range(np.random.randint(1, 4)):
+            num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
+            angle_min = mean_angle - np.random.uniform(0, angle_range)
+            angle_max = mean_angle + np.random.uniform(0, angle_range)
+            angles = []
+            vertex = []
+            for i in range(num_vertex):
+                if i % 2 == 0:
+                    angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
+                else:
+                    angles.append(np.random.uniform(angle_min, angle_max))
+
+            h, w = mask.size
+            vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
+            for i in range(num_vertex):
+                r = np.clip(
+                    np.random.normal(loc=average_radius, scale=average_radius//2),
+                    0, 2*average_radius)
+                new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
+                new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
+                vertex.append((int(new_x), int(new_y)))
+
+            draw = ImageDraw.Draw(mask)
+            width = int(np.random.uniform(min_width, max_width))
+            draw.line(vertex, fill=color, width=width)
+            for v in vertex:
+                draw.ellipse((v[0] - width//2,
+                              v[1] - width//2,
+                              v[0] + width//2,
+                              v[1] + width//2),
+                             fill=color)
+
+        return mask
+
+    width, height = img.size
+    mask = generate_mask(height, width, img)
+    return mask
 
 def inference(file, mode):
 
@@ -25,18 +85,34 @@ def inference(file, mode):
     img, orig_faces, enhanced_faces = faceenhancer.process(im)
     cv2.imwrite(os.path.join("e.png"), img)
     
+    
     if mode == "enhance":
         return os.path.join("e.png")
     elif mode == "colorize":
         model = {'name':'GPEN-1024-Color', 'size':1024}
         grayf = cv2.imread("e.png", cv2.IMREAD_GRAYSCALE)
-        grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR) # channel: 1->3
+        grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR)
         facecolorizer = FaceColorization(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         colorf = facecolorizer.process(grayf)
 
         colorf = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
         cv2.imwrite(os.path.join("output.png"), colorf)
         return os.path.join("output.png")
+    elif mode == "inpainting":
+        model = {'name':'GPEN-Inpainting-1024', 'size':1024}
+        faceinpainter = FaceInpainting(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
+        im = np.asarray(brush_stroke_mask(Image.fromarray(im)))
+        inpaint = faceinpainter.process(im)
+
+        cv2.imwrite(os.path.join("output.png"), inpaint)
+        return os.path.join("output.png")
+    elif mode == "selfie":
+        model = {'name':'GPEN-BFR-2048', 'size':2048}
+        im = cv2.resize(im, (0,0), fx=4, fy=4)
+        faceenhancer = FaceEnhancement(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
+        img, orig_faces, enhanced_faces = faceenhancer.process(im)
+        cv2.imwrite(os.path.join("output.png"), img)
+        return os.path.join("output.png")
     else:
         faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=True)
         img, orig_faces, enhanced_faces = faceenhancer.process(im)
@@ -52,14 +128,16 @@ article = "<p style='text-align: center;'><a href='https://arxiv.org/abs/2105.06
 
 gr.Interface(
     inference, 
-    [gr.inputs.Image(type="filepath", label="Input"),gr.inputs.Radio(["enhance", "colorize", "enhanced+background"], type="value", default="enhance", label="Type")], 
+    [gr.inputs.Image(type="filepath", label="Input"),gr.inputs.Radio(["enhance", "colorize", "inpainting", "selfie", "enhanced+background"], type="value", default="enhance", label="Type")], 
     gr.outputs.Image(type="file", label="Output"),
     title=title,
     description=description,
     article=article,
     examples=[
-    ['enhance.png', 'Enhance'],
-    ['color.png', 'Colorization']
+    ['enhance.png', 'enhance'],
+    ['color.png', 'colorize'],
+    ['inpainting.png', 'inpainting'],
+    ['selfie.png', 'selfie']
     ],
     enable_queue=True
     ).launch()

face_colorization.py

@@ -21,31 +21,3 @@ class FaceColorization(object):
         out = self.facegan.process(gray)
 
         return out
-        
-
-if __name__=='__main__':
-    model = {'name':'GPEN-1024-Color', 'size':1024}
-    
-    indir = 'examples/grays'
-    outdir = 'examples/couts'
-    os.makedirs(outdir, exist_ok=True)
-
-    facecolorizer = FaceColorization(size=model['size'], model=model['name'])
-
-    files = sorted(glob.glob(os.path.join(indir, '*.*g')))
-    for n, file in enumerate(files[:]):
-        filename = os.path.basename(file)
-        
-        grayf = cv2.imread(file, cv2.IMREAD_GRAYSCALE)
-        grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR) # channel: 1->3
-
-        colorf = facecolorizer.process(grayf)
-        
-        colorf = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
-
-        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'.jpg'), grayf)
-        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_COMP.jpg'), np.hstack((grayf, colorf)))
-        cv2.imwrite(os.path.join(outdir, '.'.join(filename.split('.')[:-1])+'_GPEN.jpg'), colorf)
-        
-        if n%10==0: print(n, file)
-        

face_inpainting.py

@@ -0,0 +1,18 @@
+'''
+@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
+@author: yangxy (yangtao9009@gmail.com)
+'''
+from face_model.face_gan import FaceGAN
+
+class FaceInpainting(object):
+    def __init__(self, base_dir='./', size=1024, out_size=1024, model=None, channel_multiplier=2, narrow=1, key=None, device='cuda'):
+        self.facegan = FaceGAN(base_dir, size, out_size, model, channel_multiplier, narrow, key, device=device)
+
+    # make sure the face image is well aligned. Please refer to face_enhancement.py
+    def process(self, brokenf, aligned=True):
+        # complete the face
+        out = self.facegan.process(brokenf)
+
+        return out
+
+        

face_model/face_gan.py

@@ -19,7 +19,7 @@ class FaceGAN(object):
         self.device = device
         self.is_norm = is_norm
         self.in_resolution = size
-        self.out_resolution = size if out_size==None else out_size
+        self.out_resolution = size if out_size == None else out_size
         self.key = key
         self.load_model(channel_multiplier, narrow)
 
@@ -40,6 +40,7 @@ class FaceGAN(object):
 
         with torch.no_grad():
             out, __ = self.model(img_t)
+        del img_t
 
         out = self.tensor2img(out)
 

requirements.txt

@@ -4,4 +4,6 @@ torchvision
 opencv-python
 numpy
 scikit-image
-pillow
+pillow
+scikit-learn
+joblib