Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Corrects color channels in face restoration; Fixes #1167

Browse Source
This commit is contained in:
psychedelicious 2022-10-20 10:12:20 +08:00 committed by Lincoln Stein
parent f6191a4f12
commit 4101127011
2 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
ldm/invoke/restoration/codeformer.py
View File

@ -3,6 +3,7 @@ import torch
import numpy as np import numpy as np
import warnings import warnings
import sys import sys
import cv2
pretrained_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth' pretrained_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
@ -40,11 +41,12 @@ class CodeFormerRestoration():
cf.load_state_dict(checkpoint) cf.load_state_dict(checkpoint)
cf.eval() cf.eval()
image = image.convert('RGB') # Codeformer expects BGR image data
bgrImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
face_helper = FaceRestoreHelper(upscale_factor=1, use_parse=True, device=device) face_helper = FaceRestoreHelper(upscale_factor=1, use_parse=True, device=device)
face_helper.clean_all() face_helper.clean_all()
face_helper.read_image(np.array(image, dtype=np.uint8)) face_helper.read_image(bgrImage)
face_helper.get_face_landmarks_5(resize=640, eye_dist_threshold=5) face_helper.get_face_landmarks_5(resize=640, eye_dist_threshold=5)
face_helper.align_warp_face() face_helper.align_warp_face()
@ -71,7 +73,8 @@ class CodeFormerRestoration():
restored_img = face_helper.paste_faces_to_input_image() restored_img = face_helper.paste_faces_to_input_image()
res = Image.fromarray(restored_img) # Convert back to RGB for PIL
res = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
if strength < 1.0: if strength < 1.0:
# Resize the image to the new image if the sizes have changed # Resize the image to the new image if the sizes have changed

10
ldm/invoke/restoration/gfpgan.py
View File

@ -3,6 +3,7 @@ import warnings
import os import os
import sys import sys
import numpy as np import numpy as np
import cv2
from PIL import Image from PIL import Image
@ -53,15 +54,18 @@ class GFPGAN():
f'>> Download https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth to {self.model_path}, \nor change GFPGAN directory with --gfpgan_dir.' f'>> Download https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth to {self.model_path}, \nor change GFPGAN directory with --gfpgan_dir.'
) )
image = image.convert('RGB') # GFPGAN expects BGR image data
bgrImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
_, _, restored_img = self.gfpgan.enhance( _, _, restored_img = self.gfpgan.enhance(
np.array(image, dtype=np.uint8), bgrImage,
has_aligned=False, has_aligned=False,
only_center_face=False, only_center_face=False,
paste_back=True, paste_back=True,
) )
res = Image.fromarray(restored_img)
# Convert back to RGB for PIL
res = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
if strength < 1.0: if strength < 1.0:
# Resize the image to the new image if the sizes have changed # Resize the image to the new image if the sizes have changed