v/GFPGAN
1
0
Fork 0
mirror of https://github.com/TencentARC/GFPGAN.git synced 2025-05-21 01:30:13 -07:00
Code Issues Projects Releases Wiki Activity

support GFPGAN clean arch

Browse Source
This commit is contained in:
Xintao 2021-08-06 15:45:54 +08:00
parent cc3c881f85
commit 7f67e12999
1 changed files with 17 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
inference_gfpgan_full.py
View File

@ -19,7 +19,8 @@ def restoration(gfpgan,
has_aligned=False, has_aligned=False,
only_center_face=True, only_center_face=True,
suffix=None, suffix=None,
paste_back=False): paste_back=False,
device='cuda'):
# read image # read image
img_name = os.path.basename(img_path) img_name = os.path.basename(img_path)
print(f'Processing {img_name} ...') print(f'Processing {img_name} ...')
@ -43,7 +44,7 @@ def restoration(gfpgan,
# prepare data # prepare data
cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True) cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
cropped_face_t = cropped_face_t.unsqueeze(0).to('cuda') cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
try: try:
with torch.no_grad(): with torch.no_grad():
@ -77,17 +78,18 @@ def restoration(gfpgan,
if __name__ == '__main__': if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--upscale_factor', type=int, default=1) parser.add_argument('--upscale_factor', type=int, default=2)
parser.add_argument('--arch', type=str, default='clean') parser.add_argument('--arch', type=str, default='clean')
parser.add_argument('--channel', type=int, default=2) parser.add_argument('--channel', type=int, default=2)
parser.add_argument('--model_path', type=str, default='experiments/pretrained_models/GFPGANv1.pth') parser.add_argument('--model_path', type=str, default='experiments/pretrained_models/GFPGANCleanv1-NoCE-C2.pth')
parser.add_argument('--test_path', type=str, default='inputs/whole_imgs') parser.add_argument('--test_path', type=str, default='inputs/whole_imgs')
parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces') parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces')
parser.add_argument('--only_center_face', action='store_true') parser.add_argument('--only_center_face', action='store_true')
parser.add_argument('--aligned', action='store_true') parser.add_argument('--aligned', action='store_true')
parser.add_argument('--paste_back', action='store_true') parser.add_argument('--paste_back', action='store_false')
parser.add_argument('--save_root', type=str, default='results') parser.add_argument('--save_root', type=str, default='results')
args = parser.parse_args() args = parser.parse_args()
@ -123,14 +125,17 @@ if __name__ == '__main__':
narrow=1, narrow=1,
sft_half=True) sft_half=True)
gfpgan.to(device) gfpgan.load_state_dict(torch.load(args.model_path, map_location=lambda storage, loc: storage)['params_ema'])
checkpoint = torch.load(args.model_path, map_location=lambda storage, loc: storage) gfpgan.to(device).eval()
gfpgan.load_state_dict(checkpoint['params_ema'])
gfpgan.eval()
# initialize face helper # initialize face helper
face_helper = FaceRestoreHelper( face_helper = FaceRestoreHelper(
args.upscale_factor, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png') args.upscale_factor,
face_size=512,
crop_ratio=(1, 1),
det_model='retinaface_resnet50',
save_ext='png',
device=device)
img_list = sorted(glob.glob(os.path.join(args.test_path, '*'))) img_list = sorted(glob.glob(os.path.join(args.test_path, '*')))
for img_path in img_list: for img_path in img_list:
@ -142,6 +147,7 @@ if __name__ == '__main__':
has_aligned=args.aligned, has_aligned=args.aligned,
only_center_face=args.only_center_face, only_center_face=args.only_center_face,
suffix=args.suffix, suffix=args.suffix,
paste_back=args.paste_back) paste_back=args.paste_back,
device=device)
print(f'Results are in the [{args.save_root}] folder.') print(f'Results are in the [{args.save_root}] folder.')