app.py

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

'''
@paper: GAN Prior Embedded Network for Blind Face Restoration in the Wild (CVPR2021)
@author: yangxy (yangtao9009@gmail.com)
'''
import os
import cv2
import glob
import time
import math
import imutils
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 image_resize(image, width = None, height = None, inter = cv2.INTER_AREA):
    # initialize the dimensions of the image to be resized and
    # grab the image size
    dim = None
    (h, w) = image.shape[:2]

    # if both the width and height are None, then return the
    # original image
    if width is None and height is None:
        return image

    # check to see if the width is None
    if width is None:
        # calculate the ratio of the height and construct the
        # dimensions
        r = height / float(h)
        dim = (int(w * r), height)

    # otherwise, the height is None
    else:
        # calculate the ratio of the width and construct the
        # dimensions
        r = width / float(w)
        dim = (width, int(h * r))

    # resize the image
    resized = cv2.resize(image, dim, interpolation = inter)

    # return the resized image
    return resized

def inference(file, mode):

    im = cv2.imread(file, cv2.IMREAD_COLOR)
    im = cv2.resize(im, (0,0), fx=3, fy=3)
    faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=False)
    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
        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":
        im1 = cv2.imread(file, cv2.IMREAD_COLOR)
        im2 = image_resize(im1, width = 1024)
        model = {'name':'GPEN-Inpainting-1024', 'size':1024}
        faceinpainter = FaceInpainting(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
        im3 = np.asarray(brush_stroke_mask(Image.fromarray(im2)))
        inpaint = faceinpainter.process(im3)

        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)
        cv2.imwrite(os.path.join("output.png"), img)
        return os.path.join("output.png")
        
        
title = "GPEN"
description = "Gradio demo for GAN Prior Embedded Network for Blind Face Restoration in the Wild. This version of gradio demo includes face colorization from GPEN. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center;'><a href='https://arxiv.org/abs/2105.06070' target='_blank'>GAN Prior Embedded Network for Blind Face Restoration in the Wild</a> | <a href='https://github.com/yangxy/GPEN' target='_blank'>Github Repo</a></p><p style='text-align: center;'><img src='https://img.shields.io/badge/Hugging%20Face-Original%20demo-blue' alt='https://huggingface.co/spaces/akhaliq/GPEN' width='172' height='20' /></p>"


gr.Interface(
    inference, 
    [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', 'colorize'],
    ['inpainting.png', 'inpainting'],
    ['selfie.png', 'selfie']
    ],
    enable_queue=True
    ).launch()