Spaces:
Paused
Paused
| import os | |
| import shutil | |
| import torch | |
| import torchvision.transforms as tform | |
| from PIL import Image | |
| def cpu(xs): | |
| if not isinstance(xs, (list, tuple)): | |
| return xs.cpu() | |
| else: | |
| return [x.cpu() for x in xs] | |
| def cuda(xs): | |
| if torch.cuda.is_available(): | |
| if not isinstance(xs, (list, tuple)): | |
| return xs.cuda() | |
| else: | |
| return [x.cuda() for x in xs] | |
| else: | |
| return xs | |
| def load_checkpoint(ckpt_dir_or_file, map_location=None, load_best=False): | |
| if os.path.isdir(ckpt_dir_or_file): | |
| if load_best: | |
| ckpt_path = os.path.join(ckpt_dir_or_file, 'best_model.ckpt') | |
| else: | |
| with open(os.path.join(ckpt_dir_or_file, 'latest_checkpoint')) as f: | |
| ckpt_path = os.path.join(ckpt_dir_or_file, f.readline()[:-1]) # -1去掉换行符 | |
| else: | |
| ckpt_path = ckpt_dir_or_file | |
| ckpt = torch.load(ckpt_path, map_location=map_location) | |
| print(' [*] Loading checkpoint succeeds! Copy variables from % s!' % ckpt_path) | |
| return ckpt | |
| def save_checkpoint(obj, save_path, is_best=False, max_keep=None): | |
| # save checkpoint | |
| torch.save(obj, save_path) | |
| # deal with max_keep | |
| save_dir = os.path.dirname(save_path) | |
| list_path = os.path.join(save_dir, 'latest_checkpoint') | |
| save_path = os.path.basename(save_path) | |
| if os.path.exists(list_path): | |
| with open(list_path) as f: | |
| ckpt_list = f.readlines() | |
| ckpt_list = [save_path + '\n'] + ckpt_list | |
| else: | |
| ckpt_list = [save_path + '\n'] | |
| if max_keep is not None: | |
| for ckpt in ckpt_list[max_keep:]: | |
| ckpt = os.path.join(save_dir, ckpt[:-1]) | |
| if os.path.exists(ckpt): | |
| os.remove(ckpt) | |
| ckpt_list[max_keep:] = [] | |
| with open(list_path, 'w') as f: | |
| f.writelines(ckpt_list) | |
| # copy best | |
| if is_best: | |
| shutil.copyfile(save_path, os.path.join(save_dir, 'best_model.ckpt')) | |
| def get_img_from_file(file_name, target_device, transform=False): | |
| img = tform.ToTensor()(Image.open(file_name)) # [0, 1.0] tensor | |
| if transform: | |
| val_test_img_transform = tform.Compose([ | |
| # crop face area 190 * 178 | |
| # tform.Lambda(lambda x: x[:, 14:204, :]), | |
| # center crop and resize on PLI image | |
| tform.ToPILImage(), | |
| tform.CenterCrop(170), # origin # elegant crop | |
| # tform.CenterCrop(178), | |
| tform.Resize(256, Image.BICUBIC), # elegant crop | |
| # back to tensor | |
| tform.ToTensor(), | |
| ]) | |
| img = val_test_img_transform(img) | |
| return (img * 2 - 1).to(target_device) | |