Spaces:
Running
on
T4
Running
on
T4
| # -*- coding: utf-8 -*- | |
| # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is | |
| # holder of all proprietary rights on this computer program. | |
| # You can only use this computer program if you have closed | |
| # a license agreement with MPG or you get the right to use the computer | |
| # program from someone who is authorized to grant you that right. | |
| # Any use of the computer program without a valid license is prohibited and | |
| # liable to prosecution. | |
| # | |
| # Copyright©2019 Max-Planck-Gesellschaft zur Förderung | |
| # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute | |
| # for Intelligent Systems. All rights reserved. | |
| # | |
| # Contact: ps-license@tuebingen.mpg.de | |
| import os.path as osp | |
| import numpy as np | |
| from PIL import Image | |
| import torchvision.transforms as transforms | |
| class NormalDataset(): | |
| def __init__(self, cfg, split='train'): | |
| self.split = split | |
| self.root = cfg.root | |
| self.overfit = cfg.overfit | |
| self.opt = cfg.dataset | |
| self.datasets = self.opt.types | |
| self.input_size = self.opt.input_size | |
| self.set_splits = self.opt.set_splits | |
| self.scales = self.opt.scales | |
| self.pifu = self.opt.pifu | |
| # input data types and dimensions | |
| self.in_nml = [item[0] for item in cfg.net.in_nml] | |
| self.in_nml_dim = [item[1] for item in cfg.net.in_nml] | |
| self.in_total = self.in_nml + ['normal_F', 'normal_B'] | |
| self.in_total_dim = self.in_nml_dim + [3, 3] | |
| if self.split != 'train': | |
| self.rotations = range(0, 360, 120) | |
| else: | |
| self.rotations = np.arange(0, 360, 360 / | |
| self.opt.rotation_num).astype(np.int) | |
| self.datasets_dict = {} | |
| for dataset_id, dataset in enumerate(self.datasets): | |
| dataset_dir = osp.join(self.root, dataset, "smplx") | |
| self.datasets_dict[dataset] = { | |
| "subjects": | |
| np.loadtxt(osp.join(self.root, dataset, "all.txt"), dtype=str), | |
| "path": | |
| dataset_dir, | |
| "scale": | |
| self.scales[dataset_id] | |
| } | |
| self.subject_list = self.get_subject_list(split) | |
| # PIL to tensor | |
| self.image_to_tensor = transforms.Compose([ | |
| transforms.Resize(self.input_size), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) | |
| ]) | |
| # PIL to tensor | |
| self.mask_to_tensor = transforms.Compose([ | |
| transforms.Resize(self.input_size), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.0, ), (1.0, )) | |
| ]) | |
| def get_subject_list(self, split): | |
| subject_list = [] | |
| for dataset in self.datasets: | |
| if self.pifu: | |
| txt = osp.join(self.root, dataset, f'{split}_pifu.txt') | |
| else: | |
| txt = osp.join(self.root, dataset, f'{split}.txt') | |
| if osp.exists(txt): | |
| print(f"load from {txt}") | |
| subject_list += sorted(np.loadtxt(txt, dtype=str).tolist()) | |
| if self.pifu: | |
| miss_pifu = sorted( | |
| np.loadtxt(osp.join(self.root, dataset, | |
| "miss_pifu.txt"), | |
| dtype=str).tolist()) | |
| subject_list = [ | |
| subject for subject in subject_list | |
| if subject not in miss_pifu | |
| ] | |
| subject_list = [ | |
| "renderpeople/" + subject for subject in subject_list | |
| ] | |
| else: | |
| train_txt = osp.join(self.root, dataset, 'train.txt') | |
| val_txt = osp.join(self.root, dataset, 'val.txt') | |
| test_txt = osp.join(self.root, dataset, 'test.txt') | |
| print( | |
| f"generate lists of [train, val, test] \n {train_txt} \n {val_txt} \n {test_txt} \n" | |
| ) | |
| split_txt = osp.join(self.root, dataset, f'{split}.txt') | |
| subjects = self.datasets_dict[dataset]['subjects'] | |
| train_split = int(len(subjects) * self.set_splits[0]) | |
| val_split = int( | |
| len(subjects) * self.set_splits[1]) + train_split | |
| with open(train_txt, "w") as f: | |
| f.write("\n".join(dataset + "/" + item | |
| for item in subjects[:train_split])) | |
| with open(val_txt, "w") as f: | |
| f.write("\n".join( | |
| dataset + "/" + item | |
| for item in subjects[train_split:val_split])) | |
| with open(test_txt, "w") as f: | |
| f.write("\n".join(dataset + "/" + item | |
| for item in subjects[val_split:])) | |
| subject_list += sorted( | |
| np.loadtxt(split_txt, dtype=str).tolist()) | |
| bug_list = sorted( | |
| np.loadtxt(osp.join(self.root, 'bug.txt'), dtype=str).tolist()) | |
| subject_list = [ | |
| subject for subject in subject_list if (subject not in bug_list) | |
| ] | |
| return subject_list | |
| def __len__(self): | |
| return len(self.subject_list) * len(self.rotations) | |
| def __getitem__(self, index): | |
| # only pick the first data if overfitting | |
| if self.overfit: | |
| index = 0 | |
| rid = index % len(self.rotations) | |
| mid = index // len(self.rotations) | |
| rotation = self.rotations[rid] | |
| # choose specific test sets | |
| subject = self.subject_list[mid] | |
| subject_render = "/".join( | |
| [subject.split("/")[0] + "_12views", | |
| subject.split("/")[1]]) | |
| # setup paths | |
| data_dict = { | |
| 'dataset': | |
| subject.split("/")[0], | |
| 'subject': | |
| subject, | |
| 'rotation': | |
| rotation, | |
| 'image_path': | |
| osp.join(self.root, subject_render, 'render', | |
| f'{rotation:03d}.png') | |
| } | |
| # image/normal/depth loader | |
| for name, channel in zip(self.in_total, self.in_total_dim): | |
| if name != 'image': | |
| data_dict.update({ | |
| f'{name}_path': | |
| osp.join(self.root, subject_render, name, | |
| f'{rotation:03d}.png') | |
| }) | |
| data_dict.update({ | |
| name: | |
| self.imagepath2tensor(data_dict[f'{name}_path'], | |
| channel, | |
| inv='depth_B' in name) | |
| }) | |
| path_keys = [ | |
| key for key in data_dict.keys() if '_path' in key or '_dir' in key | |
| ] | |
| for key in path_keys: | |
| del data_dict[key] | |
| return data_dict | |
| def imagepath2tensor(self, path, channel=3, inv=False): | |
| rgba = Image.open(path).convert('RGBA') | |
| mask = rgba.split()[-1] | |
| image = rgba.convert('RGB') | |
| image = self.image_to_tensor(image) | |
| mask = self.mask_to_tensor(mask) | |
| image = (image * mask)[:channel] | |
| return (image * (0.5 - inv) * 2.0).float() | |