| """This script defines the custom dataset for Deep3DFaceRecon_pytorch |
| """ |
|
|
| import os.path |
| from data.base_dataset import BaseDataset, get_transform, get_affine_mat, apply_img_affine, apply_lm_affine |
| from data.image_folder import make_dataset |
| from PIL import Image |
| import random |
| import util.util as util |
| import numpy as np |
| import json |
| import torch |
| from scipy.io import loadmat, savemat |
| import pickle |
| from util.preprocess import align_img, estimate_norm |
| from util.load_mats import load_lm3d |
|
|
|
|
| def default_flist_reader(flist): |
| """ |
| flist format: impath label\nimpath label\n ...(same to caffe's filelist) |
| """ |
| imlist = [] |
| with open(flist, 'r') as rf: |
| for line in rf.readlines(): |
| impath = line.strip() |
| imlist.append(impath) |
|
|
| return imlist |
|
|
| def jason_flist_reader(flist): |
| with open(flist, 'r') as fp: |
| info = json.load(fp) |
| return info |
|
|
| def parse_label(label): |
| return torch.tensor(np.array(label).astype(np.float32)) |
|
|
|
|
| class FlistDataset(BaseDataset): |
| """ |
| It requires one directories to host training images '/path/to/data/train' |
| You can train the model with the dataset flag '--dataroot /path/to/data'. |
| """ |
|
|
| def __init__(self, opt): |
| """Initialize this dataset class. |
| |
| Parameters: |
| opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions |
| """ |
| BaseDataset.__init__(self, opt) |
| |
| self.lm3d_std = load_lm3d(opt.bfm_folder) |
| |
| msk_names = default_flist_reader(opt.flist) |
| self.msk_paths = [os.path.join(opt.data_root, i) for i in msk_names] |
|
|
| self.size = len(self.msk_paths) |
| self.opt = opt |
| |
| self.name = 'train' if opt.isTrain else 'val' |
| if '_' in opt.flist: |
| self.name += '_' + opt.flist.split(os.sep)[-1].split('_')[0] |
| |
|
|
| def __getitem__(self, index): |
| """Return a data point and its metadata information. |
| |
| Parameters: |
| index (int) -- a random integer for data indexing |
| |
| Returns a dictionary that contains A, B, A_paths and B_paths |
| img (tensor) -- an image in the input domain |
| msk (tensor) -- its corresponding attention mask |
| lm (tensor) -- its corresponding 3d landmarks |
| im_paths (str) -- image paths |
| aug_flag (bool) -- a flag used to tell whether its raw or augmented |
| """ |
| msk_path = self.msk_paths[index % self.size] |
| img_path = msk_path.replace('mask/', '') |
| lm_path = '.'.join(msk_path.replace('mask', 'landmarks').split('.')[:-1]) + '.txt' |
|
|
| raw_img = Image.open(img_path).convert('RGB') |
| raw_msk = Image.open(msk_path).convert('RGB') |
| raw_lm = np.loadtxt(lm_path).astype(np.float32) |
|
|
| _, img, lm, msk = align_img(raw_img, raw_lm, self.lm3d_std, raw_msk) |
| |
| aug_flag = self.opt.use_aug and self.opt.isTrain |
| if aug_flag: |
| img, lm, msk = self._augmentation(img, lm, self.opt, msk) |
| |
| _, H = img.size |
| M = estimate_norm(lm, H) |
| transform = get_transform() |
| img_tensor = transform(img) |
| msk_tensor = transform(msk)[:1, ...] |
| lm_tensor = parse_label(lm) |
| M_tensor = parse_label(M) |
|
|
|
|
| return {'imgs': img_tensor, |
| 'lms': lm_tensor, |
| 'msks': msk_tensor, |
| 'M': M_tensor, |
| 'im_paths': img_path, |
| 'aug_flag': aug_flag, |
| 'dataset': self.name} |
|
|
| def _augmentation(self, img, lm, opt, msk=None): |
| affine, affine_inv, flip = get_affine_mat(opt, img.size) |
| img = apply_img_affine(img, affine_inv) |
| lm = apply_lm_affine(lm, affine, flip, img.size) |
| if msk is not None: |
| msk = apply_img_affine(msk, affine_inv, method=Image.BILINEAR) |
| return img, lm, msk |
| |
|
|
|
|
|
|
| def __len__(self): |
| """Return the total number of images in the dataset. |
| """ |
| return self.size |
|
|
