app.py

import os

os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth" -O weights/RetinaFace-R50.pth')
os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth" -O weights/GPEN-512.pth')
os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth" -O weights/GPEN-1024-Color.pth ')
os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/realesrnet_x2.pth" -O weights/realesrnet_x2.pth ')
os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Inpainting-1024.pth" -O weights/GPEN-Inpainting-1024.pth ')
jksp= os.environ['SELFIE']
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 ([email protected])
'''
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 
from gradio_imageslider import ImageSlider

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 resize(image, width = 1024):
    aspect_ratio = float(image.shape[1])/float(image.shape[0])
    height = width/aspect_ratio
    image = cv2.resize(image, (int(height),int(width)))
    return image

def inference(file, mode, res_percentage, zoom, x_shift, y_shift):

    im = cv2.resize(file, None, fx = (res_percentage/100), fy = (res_percentage/100))
    
    if mode == "enhance":
        faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=False)
        img, orig_faces, enhanced_faces = faceenhancer.process(im)
        
    elif mode == "colorize":
        model = {'name':'GPEN-1024-Color', 'size':1024}
        if len(im.shape) == 3:
            if im.shape[2] == 1:
                grayf = im[:, :, 0]
            elif im.shape[2] == 3:
                grayf = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
            elif im.shape[2] == 4:
                grayf = cv2.cvtColor(im, cv2.COLOR_BGRA2GRAY)
            grayf = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
        
        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)
        img = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))
        
    elif mode == "inpainting":
        im2 = resize(im, 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)))
        img = faceinpainter.process(im3)
        
    elif mode == "selfie":
        model = {'name':'GPEN-BFR-2048', 'size':2048}
        faceenhancer = FaceEnhancement(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
        img, orig_faces, enhanced_faces = faceenhancer.process(im)
        
    else:
        faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cpu', u=True)
        img, orig_faces, enhanced_faces = faceenhancer.process(im)
    
    (in_img, out_img) = zoom_image(zoom, x_shift, y_shift, im, img)

    return img, (in_img, out_img)
        
        
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>"


def zoom_image(zoom, x_shift, y_shift, input_img, output_img = None):
    if output_img is None:
        return None
    
    img = Image.fromarray(input_img)
    out_img = Image.fromarray(output_img)
    
    img_w, img_h = img.size
    zoom_factor = (100 - zoom) / 100
    x_shift /= 100
    y_shift /= 100
    
    zoom_w, zoom_h = int(img_w * zoom_factor), int(img_h * zoom_factor)
    x_offset = int((img_w - zoom_w) * x_shift)
    y_offset = int((img_h - zoom_h) * y_shift)
    
    crop_box = (x_offset, y_offset, x_offset + zoom_w, y_offset + zoom_h)
    img = img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)
    out_img = out_img.resize((img_w, img_h), Image.BILINEAR).crop(crop_box)

    return (img, out_img)

with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image(label="Input Image")
        output_img = gr.Image(label="Result", interactive=False)
    
    restore_type = gr.Radio(["enhance", "colorize", "inpainting", "selfie", "enhanced+background"], value="enhance", type="value", label="Type")
    max_res = gr.Slider(1, 200, step=0.5, value=100, label="Output image Resolution Percentage (Higher% = longer processing time)")
    
    zoom = gr.Slider(0, 100, step=1, value=50, label="Zoom Percentage (0 = original size)")
    x_shift = gr.Slider(0, 100, step=1, label="Horizontal shift Percentage (Before/After)")
    y_shift = gr.Slider(0, 100, step=1, label="Vertical shift Percentage (Before/After)")
    
    run = gr.Button("Run")

    with gr.Row():
        before_after = ImageSlider(label="Before/After", type="pil", value=None)

    run.click(
        inference, 
        inputs=[input_img, restore_type, max_res, zoom, x_shift, y_shift], 
        outputs=[output_img, before_after]
    )
    
    gr.Examples([
        ['enhance.png', 'enhance', 100, 0, 0, 0],
        ['color.png', 'colorize', 100, 0, 0, 0],
        ['inpainting.png', 'inpainting', 100, 0, 0, 0],
        ['selfie.png', 'selfie', 100, 0, 0, 0]
    ], inputs=[input_img, restore_type, max_res, zoom, x_shift, y_shift])
    
    zoom.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img, output_img], outputs=[before_after])
    x_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img, output_img], outputs=[before_after])
    y_shift.release(zoom_image, inputs=[zoom, x_shift, y_shift, input_img, output_img], outputs=[before_after])

demo.launch()