Spaces:
Running
Running
| 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() |