Flips channels using array slicing instead of using OpenCV

This commit is contained in:
psychedelicious 2022-10-20 11:42:01 +08:00 committed by Lincoln Stein
parent 4101127011
commit 87469a5fdd
2 changed files with 13 additions and 12 deletions

View File

@ -3,7 +3,6 @@ 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'
@ -41,12 +40,13 @@ class CodeFormerRestoration():
cf.load_state_dict(checkpoint) cf.load_state_dict(checkpoint)
cf.eval() cf.eval()
# Codeformer expects BGR image data image = image.convert('RGB')
bgrImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # Codeformer expects a BGR np array; make array and flip channels
bgr_image_array = np.array(image, dtype=np.uint8)[...,::-1]
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(bgrImage) face_helper.read_image(bgr_image_array)
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()
@ -73,8 +73,8 @@ class CodeFormerRestoration():
restored_img = face_helper.paste_faces_to_input_image() restored_img = face_helper.paste_faces_to_input_image()
# Convert back to RGB for PIL # Flip the channels back to RGB
res = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)) res = Image.fromarray(restored_img[...,::-1])
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

View File

@ -3,7 +3,6 @@ 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
@ -54,18 +53,20 @@ 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.'
) )
# GFPGAN expects BGR image data image = image.convert('RGB')
bgrImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
# GFPGAN expects a BGR np array; make array and flip channels
bgr_image_array = np.array(image, dtype=np.uint8)[...,::-1]
_, _, restored_img = self.gfpgan.enhance( _, _, restored_img = self.gfpgan.enhance(
bgrImage, bgr_image_array,
has_aligned=False, has_aligned=False,
only_center_face=False, only_center_face=False,
paste_back=True, paste_back=True,
) )
# Convert back to RGB for PIL # Flip the channels back to RGB
res = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)) res = Image.fromarray(restored_img[...,::-1])
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