def __init__(self,
prefix,
topo_loader: TopologyLoader,
device,
is_train=True,
fk=None):
super(MultiGarmentDataset, self).__init__(device)
self.prefix = prefix
self.smpl_hires = SMPL_Layer(highRes=True).to(device)
self.smpl = SMPL_Layer().to(device)
self.parents = self.smpl.kintree_parents
self.faces_hires = self.smpl_hires.th_faces
self.faces = self.smpl.th_faces
self.bone_num = len(self.parents)
lst = [
f for f in os.listdir(prefix) if os.path.isdir(pjoin(prefix, f))
]
lst.sort()
lst = lst[:80] if is_train else lst[80:]
self.t_pose_list = []
self.offset_list = []
self.weight_hires = self.smpl_hires.th_weights.to(device)
self.weight = self.smpl.th_weights.to(device)
self.cloth_all = np.load(pjoin(prefix, 'all_cloths.npy'))
self.cloth_all = torch.tensor(self.cloth_all, device=device)
for name in lst:
prefix2 = pjoin(prefix, name)
t_pose = np.load(pjoin(prefix2, 't-pose.npy'))
offset = np.load(pjoin(prefix2, 'offset.npy'))
t_pose = torch.tensor(t_pose, device=device)
offset = torch.tensor(offset, device=device)
self.t_pose_list.append(t_pose.unsqueeze(0))
self.offset_list.append(offset.unsqueeze(0))
high2o_mask = np.array(
[True] * self.smpl.num_verts + [False] *
(self.smpl_hires.num_verts - self.smpl.num_verts))
self.topo_id_hires = topo_loader.load_from_obj(
pjoin(prefix, 'high_res.obj'))
self.topo_id = topo_loader.load_from_obj(pjoin(prefix, 'original.obj'))
self.t_pose_list = torch.cat(self.t_pose_list, dim=0)
self.offset_list = torch.cat(self.offset_list, dim=0)
if fk is None:
fk = ForwardKinematics(self.parents)
self.fk = fk
def __init__(self, filenames, topo_loader: TopologyLoader, weight_gt=None):
self.t_poses = []
self.topo_id = []
self.faces = []
for filename in filenames:
self.topo_id.append(topo_loader.load_from_obj(filename))
self.t_poses.append(topo_loader.t_poses[-1])
self.faces.append(topo_loader.faces[-1])
if weight_gt is None:
weight_gt = torch.tensor([
0.,
])
self.weight = weight_gt
def prepare_dataset(device, args):
topo_loader = TopologyLoader(device=device, debug=args.debug)
# Prepare SMPL dataset and MultiGarmentDataset
dataset_smpl = SMPLDataset(device=device)
dataset_garment = MultiGarmentDataset('./dataset/Meshes/MultiGarment',
topo_loader, device)
# Prepare topology augmentation
if args.topo_augment:
begin_aug_topo, len_topo = topo_loader.load_smpl_group(
'./dataset/Meshes/SMPL/topology/', is_train=True)
else:
begin_aug_topo = topo_loader.load_from_obj(
'./dataset/eval_constant/meshes/smpl_std.obj')
len_topo = 1
return topo_loader, dataset_smpl, dataset_garment, begin_aug_topo, len_topo