diff --git a/GFPGAN.py b/GFPGAN.py index 8ffa70a..473b9f4 100644 --- a/GFPGAN.py +++ b/GFPGAN.py @@ -10,27 +10,28 @@ import numpy as np import torch.nn.functional as F +# 将Tensor张量转换为Numpy数组 def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): - """Convert torch Tensors into image numpy arrays. - - After clamping to [min, max], values will be normalized to [0, 1]. + """ + 会将值归一化到[0,1]之间，输入的Tensor向量为RGB格式 Args: - tensor (Tensor or list[Tensor]): Accept shapes: - 1) 4D mini-batch Tensor of shape (B x 3/1 x H x W); - 2) 3D Tensor of shape (3/1 x H x W); - 3) 2D Tensor of shape (H x W). - Tensor channel should be in RGB order. - rgb2bgr (bool): Whether to change rgb to bgr. - out_type (numpy type): output types. If ``np.uint8``, transform outputs - to uint8 type with range [0, 255]; otherwise, float type with - range [0, 1]. Default: ``np.uint8``. - min_max (tuple[int]): min and max values for clamp. - + tensor (Tensor or list[Tensor]): Accept shapes:(B:批次大小, H:图像高, W:图像宽度) + 1.4D mini-batch张量，形状为(B x 3/1 x H x W)； + 2.3D张量，形状为(3/1 x H x W)； + 3.2D张量，形状为(H x W)。 + rgb2bgr (bool): 是否将RGB通道顺序转换为BGR + out_type (numpy type): 输出的NumPy数组类型。 + 如果为np.uint8，则将输出转换为uint8类型，范围为[0, 255]。 + 否则，输出将是浮点数类型，范围为[0, 1]。 + 默认值为np.uint8。 + min_max (tuple[int])：clamp（夹取）操作的最小和最大值 + 形式为元组，包含两个整数值 Returns: - (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of - shape (H x W). The channel order is BGR. + 三维阵列(H×W×C)或二维阵列，形状(H × W)，通道顺序为BGR + 这里的C指的是通道数 """ + # 检查输入参数tensor是否为PyTorch张量或张量列表，如果不是则引发TypeError if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}') @@ -38,37 +39,38 @@ def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)): _tensor = tensor import time start = time.time() + # 把批处理维度去掉同时把原值夹成[0,1] _tensor = _tensor.squeeze(0).float().detach().clamp_(*min_max) end = time.time() _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0]) _tensor = (_tensor.permute(1, 2, 0)) - + print("tensor:") + print(_tensor) img_np = (_tensor * 255.0).round().cpu().numpy()[:, :, ::-1] img_np = img_np.astype(out_type) result.append(img_np) + # 消除列表嵌套 if len(result) == 1: result = result[0] end = time.time() return result - +# 创建用于图像恢复的类 class GFPGANer(): - """Helper for restoration with GFPGAN. + """使用GFPGAN进行恢复 - It will detect and crop faces, and then resize the faces to 512x512. - GFPGAN is used to restored the resized faces. - The background is upsampled with the bg_upsampler. - Finally, the faces will be pasted back to the upsample background image. + 检测裁剪，将尺寸调为512*512 + 背景使用bg_upsampler进行上采样. Args: - model_path (str): The path to the GFPGAN model. It can be urls (will first download it automatically). - upscale (float): The upscale of the final output. Default: 2. - arch (str): The GFPGAN architecture. Option: clean | original. Default: clean. - channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2. - bg_upsampler (nn.Module): The upsampler for the background. Default: None. + model_path (str): GFPGAN的模型路径. It can be urls (will first download it automatically). + upscale (float): 最终输出的缩放比例. Default: 2. + arch (str): GFPGAN的架构. Option: clean | original. Default: clean. + channel_multiplier (int): StyleGAN2大型网络的通道乘法器. Default: 2. + bg_upsampler (nn.Module): 背景的上采样器. Default: None. """ def __init__(self, device,model_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=None): @@ -111,10 +113,14 @@ def __init__(self, device,model_path, upscale=2, arch='clean', channel_multiplie print('GFPGAN model loaded') @torch.no_grad() + # 对整个图像进行加强，适用于多张人脸的图像 def enhance_allimg(self, img, has_aligned=False, only_center_face=False, paste_back=True): + # 清除 self.face_helper.clean_all() import time start = time.time() + # 如果图像已对齐，将其调整大小为(512, 512)，作为单个人脸存储； + # 如果图像未对齐，获取每张脸的关键点，对每张脸进行对齐 if has_aligned: # the inputs are already aligned img = cv2.resize(img, (512, 512)) self.face_helper.cropped_faces = [img] @@ -130,6 +136,7 @@ def enhance_allimg(self, img, has_aligned=False, only_center_face=False, paste_b # print('got face: ', (end - start)*1000) # face restoration start = time.time() + # 迭代处理每一张人脸 for cropped_face in self.face_helper.cropped_faces: # prepare data start = time.time() @@ -183,10 +190,15 @@ def enhance_allimg(self, img, has_aligned=False, only_center_face=False, paste_b @torch.no_grad() def enhance(self, img, has_aligned=False, only_center_face=False, paste_back=True): + # has_aligned: 输入图像是否已经对齐，默认为False + # only_center_face: 仅处理中心脸，默认为False + # paste_back: 是否将处理后的脸粘贴回原图，默认为True self.face_helper.clean_all() + # 如果输入图像已经对齐，则直接将图像添加到self.face_helper.cropped_faces列表 if has_aligned: # the inputs are already aligned self.face_helper.cropped_faces = [img] + # 否则，使用FaceRestoreHelper类对输入图像进行人脸检测、对齐和裁剪 else: self.face_helper.read_image(img) # get face landmarks for each face @@ -197,8 +209,7 @@ def enhance(self, img, has_aligned=False, only_center_face=False, paste_back=Tru self.face_helper.align_warp_face() #print('got face: ', (end - start)*1000) - # face restoration - + # 对每个裁剪的人脸进行GFPGAN模型的推理和图像的还原 for cropped_face in self.face_helper.cropped_faces: # prepare data @@ -209,8 +220,6 @@ def enhance(self, img, has_aligned=False, only_center_face=False, paste_back=Tru cropped_face_t = F.interpolate(cropped_face_t, (512, 512), mode='bilinear', align_corners=True) - - try: output = self.gfpgan(cropped_face_t, return_rgb=False)[0] # #NCHW 33ms @@ -224,8 +233,7 @@ def enhance(self, img, has_aligned=False, only_center_face=False, paste_back=Tru # restored_face = restored_face.astype('uint8') self.face_helper.add_restored_face(restored_face) - - + # 如果图像没有对齐，并且需要将脸粘贴回原图，则进行背景的上采样 if not has_aligned and paste_back: # upsample the background if self.bg_upsampler is not None: @@ -233,7 +241,7 @@ def enhance(self, img, has_aligned=False, only_center_face=False, paste_back=Tru bg_img = self.bg_upsampler.enhance(img, outscale=self.upscale)[0] else: bg_img = None - + # 获取逆仿射变换参数，将还原的脸粘贴回原图 self.face_helper.get_inverse_affine(None) # paste each restored face to the input image restored_img = self.face_helper.paste_faces_to_input_image(upsample_img=bg_img) @@ -268,8 +276,12 @@ def GFPGANInit(device,face_enhancement_path): return restorer def GFPGANInfer(img, restorer, aligned): + # restorer：GFPGANer类的实例，用于对图像进行增强或人脸修复 + # aligned：一个布尔值，指示输入的图像是否已经对齐 only_center_face = True start = time.time() + # 如果输入图像已对齐，那么调用 restorer.enhance 方法对单个人脸图像进行增强，并返回还原后的人脸 + # 如果输入图像未对齐，那么调用 restorer.enhance_allimg 方法对整个图像进行增强，并返回整个合成图像 if aligned: cropped_faces, restored_faces, restored_img = restorer.enhance( img, has_aligned=aligned, only_center_face=only_center_face, paste_back=True) @@ -285,3 +297,4 @@ def GFPGANInfer(img, restorer, aligned): return restored_faces[0] +