slowly update dataset (untested)
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
from struct import pack, unpack
|
||||
@@ -220,6 +221,35 @@ def _todict(matobj):
|
||||
return dict
|
||||
|
||||
|
||||
# can i cache map?
|
||||
def read_posed_garment_frame(folder, id, garment, frame, mask=None, shape=None, shapedirs=None, unshape=False):
|
||||
pc16_path = os.path.join(folder, id, garment + '.pc16')
|
||||
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
|
||||
verts = readPC2Frame(pc16_path, frame, True)
|
||||
verts = verts[:, map, :]
|
||||
|
||||
if unshape:
|
||||
shape_offsets = np.tensordot(shape, shapedirs, axes=[[1], [2]])
|
||||
verts[mask, :] = verts[mask, :] - np.squeeze(shape_offsets)[mask, :]
|
||||
|
||||
return verts
|
||||
|
||||
|
||||
def read_garment_mask(folder, id, garment):
|
||||
raw_mask = sio.loadmat(os.path.join(folder, id, garment + '_mask.mat'))['mask'].reshape((-1))
|
||||
mask = raw_mask > 0
|
||||
return mask
|
||||
|
||||
|
||||
# can i cache mask?
|
||||
def read_garment_template(folder, id, garment, mask):
|
||||
offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3]
|
||||
body_static = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B']
|
||||
verts_static = (body_static + offsets_static) * mask[:, np.newaxis]
|
||||
return verts_static
|
||||
|
||||
|
||||
|
||||
# import os
|
||||
# import sys
|
||||
# import numpy as np
|
||||
|
||||
+50
-112
@@ -3,58 +3,19 @@ import sys
|
||||
import time
|
||||
import cv2
|
||||
import random
|
||||
import pickle
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
import scipy.io as sio
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
|
||||
from custom_utils.custom_utils import nest_conditions, batch_to_dict
|
||||
|
||||
|
||||
|
||||
# class DatasetSimple(tf.data.Dataset):
|
||||
|
||||
# def _generator(folder):
|
||||
# # Opening the file
|
||||
# folder = folder.decode("utf-8")
|
||||
# uv_folder = os.path.join(folder, 'uv_posed')
|
||||
# pose_hape_folder = os.path.join(folder, 'poses_shapes')
|
||||
# frames = os.listdir(uv_folder)
|
||||
# random.shuffle(frames)
|
||||
|
||||
# for frame in frames:
|
||||
# id, frame_nr = frame.split('.')[0].split('_')
|
||||
# uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1])
|
||||
# uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1])
|
||||
# pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
|
||||
# uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1])
|
||||
# mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']
|
||||
# mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
|
||||
# mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
|
||||
# uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
|
||||
# mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
|
||||
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
|
||||
|
||||
# # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask
|
||||
# yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed
|
||||
|
||||
# def __new__(cls, folder, res=256):
|
||||
# return tf.data.Dataset.from_generator(
|
||||
# cls._generator,
|
||||
# output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
|
||||
# (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
|
||||
# tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ),
|
||||
# args=(folder, )
|
||||
# )
|
||||
|
||||
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo
|
||||
|
||||
|
||||
class StaticDataset():
|
||||
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None:
|
||||
def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False) -> None:
|
||||
self.store = dict()
|
||||
self.signature = signature
|
||||
self.folder = folder
|
||||
@@ -69,8 +30,6 @@ class StaticDataset():
|
||||
self.resolutions = resolutions
|
||||
self.filled = ''
|
||||
self.is_cape = is_cape
|
||||
if use_filled:
|
||||
self.filled = '_filled'
|
||||
|
||||
if is_cape:
|
||||
self.uv_masks = dict()
|
||||
@@ -82,10 +41,20 @@ class StaticDataset():
|
||||
self.uv_masks[res] = down
|
||||
|
||||
|
||||
|
||||
def read(self, id):
|
||||
def read(self, id, garment):
|
||||
data = dict()
|
||||
|
||||
info = loadInfo(os.path.join(self.folder, id))
|
||||
|
||||
data['shape'] = info['shape']
|
||||
data['poses'] = info['poses']
|
||||
|
||||
if 'mesh_mask' in self.signature.keys():
|
||||
if self.is_cape:
|
||||
data['mesh_mask'] = np.ones((3638), np.float32)
|
||||
else:
|
||||
data['mesh_mask'] = read_garment_mask(self.folder, id, garment).astype(np.float32)
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
@@ -99,19 +68,12 @@ class StaticDataset():
|
||||
else:
|
||||
uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
|
||||
|
||||
|
||||
if 'mesh_static' in self.signature.keys():
|
||||
mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
|
||||
if 'mesh_mask' in self.signature.keys():
|
||||
if self.is_cape:
|
||||
data['mesh_mask'] = np.ones((3638), np.float32)
|
||||
else:
|
||||
mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
|
||||
if 'uv_shaped_body' in self.signature.keys():
|
||||
uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
|
||||
data['mesh_static'] = read_garment_template(self.folder, id, garment)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
if id not in self.store:
|
||||
item = self.read(id)
|
||||
@@ -122,11 +84,12 @@ class StaticDataset():
|
||||
|
||||
return self.store[id]
|
||||
|
||||
|
||||
def expand_batch(self, batch):
|
||||
for key in self.keys:
|
||||
batch[key] = tf.zeros(((0,) + self.signature[key]))
|
||||
for i in range(batch['id'].shape[0]):
|
||||
item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'))
|
||||
item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'), batch['garment'][i].numpy().decode('UTF-8'))
|
||||
for key in self.keys:
|
||||
batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
|
||||
return batch
|
||||
@@ -136,20 +99,19 @@ class StaticDataset():
|
||||
class DatasetIdx(tf.data.Dataset):
|
||||
|
||||
def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
|
||||
if frames_file is None:
|
||||
uv_static_folder = os.path.join(folder, 'pose_shape')
|
||||
frames = sorted(os.listdir(uv_static_folder))
|
||||
# if frames_file is None:
|
||||
# uv_static_folder = os.path.join(folder, 'pose_shape')
|
||||
# frames = sorted(os.listdir(uv_static_folder))
|
||||
|
||||
if is_cape:
|
||||
frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
|
||||
else:
|
||||
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
|
||||
# if is_cape:
|
||||
# frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
|
||||
# else:
|
||||
# frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
|
||||
|
||||
if testing:
|
||||
frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
|
||||
frames = [frame.split('.')[0] + '_4' for frame in frames]
|
||||
# if testing:
|
||||
# frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
|
||||
# frames = [frame.split('.')[0] + '_4' for frame in frames]
|
||||
|
||||
if frames_file is not None:
|
||||
frames = []
|
||||
with open(os.path.join(folder, frames_file), 'r') as f:
|
||||
lines = f.readlines()
|
||||
@@ -162,7 +124,6 @@ class DatasetIdx(tf.data.Dataset):
|
||||
return tf.data.Dataset.from_tensor_slices(frames)
|
||||
|
||||
|
||||
# class DatasetFactory(tf.data.Dataset):
|
||||
class DatasetFactory():
|
||||
|
||||
def __init__(self):
|
||||
@@ -175,6 +136,8 @@ class DatasetFactory():
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'frame_nr':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'garment':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'mesh_curvature':
|
||||
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
|
||||
else:
|
||||
@@ -191,7 +154,6 @@ class DatasetFactory():
|
||||
in_memory=False,
|
||||
do_repeat=True,
|
||||
is_cape=False,
|
||||
use_one_hot=False,
|
||||
use_no_orient=True,
|
||||
testing=False,
|
||||
frames_file=None):
|
||||
@@ -203,31 +165,33 @@ class DatasetFactory():
|
||||
if is_cape:
|
||||
separator = '.'
|
||||
|
||||
one_hot = ''
|
||||
one_hot_2 = ''
|
||||
if use_one_hot:
|
||||
one_hot = '_one_hot'
|
||||
one_hot_2 = '_one_hot_2'
|
||||
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
|
||||
if is_cape:
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
|
||||
|
||||
if use_no_orient == False:
|
||||
assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!'
|
||||
|
||||
to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
|
||||
shapedirs = []
|
||||
with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
|
||||
female = pickle.load(f, encoding="latin1")
|
||||
shapedirs.append(female['shapedirs'][to_keep, :])
|
||||
with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
|
||||
male = pickle.load(f, encoding="latin1")
|
||||
shapedirs.append(male['shapedirs'][to_keep, :])
|
||||
|
||||
|
||||
|
||||
def read_item(frame):
|
||||
frame_split = tf.strings.split(frame, sep=separator)
|
||||
id, frame_nr = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8")
|
||||
# id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy()
|
||||
# static_info = static_data[id]
|
||||
id, frame_nr, garment = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8"), frame_split[2].numpy().decode("utf-8")
|
||||
item = dict()
|
||||
|
||||
if 'id' in signature.keys():
|
||||
item['id'] = id
|
||||
if 'frame_nr' in signature.keys():
|
||||
item['frame_nr'] = frame_nr
|
||||
if 'garment' in signature.keys():
|
||||
item['garment'] = garment
|
||||
if 'pose_shape' in signature.keys():
|
||||
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
|
||||
if use_no_orient:
|
||||
@@ -247,15 +211,8 @@ class DatasetFactory():
|
||||
else:
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
|
||||
item['mesh_posed'] = mesh_posed
|
||||
# if 'mesh_unposed' in signature.keys():
|
||||
# mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
# item['mesh_unposed'] = mesh_unposed
|
||||
# if 'uv_unposed' in signature.keys():
|
||||
# uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
|
||||
# item['uv_unposed'] = uv_unposed
|
||||
if 'mesh_unposed' in signature.keys():
|
||||
# print('one_hot')
|
||||
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
item['mesh_unposed'] = mesh_unposed_no_shape
|
||||
|
||||
|
||||
@@ -279,7 +236,7 @@ class DatasetFactory():
|
||||
item['uv_curvature' + res_string] = uv_curvature
|
||||
|
||||
if 'uv_unposed' + res_string in signature.keys():
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in signature.keys():
|
||||
@@ -299,9 +256,6 @@ class DatasetFactory():
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
|
||||
# print('here')
|
||||
|
||||
# item = { **static_info, **item }
|
||||
as_list = []
|
||||
for key in signature.keys():
|
||||
if key in item:
|
||||
@@ -309,10 +263,6 @@ class DatasetFactory():
|
||||
|
||||
return as_list
|
||||
|
||||
# dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
# inp=[i],
|
||||
# Tout=self._build_tf_signature(signature)
|
||||
# ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
|
||||
|
||||
if in_memory:
|
||||
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
@@ -328,20 +278,8 @@ class DatasetFactory():
|
||||
if do_repeat:
|
||||
dataset = dataset.repeat()
|
||||
|
||||
# dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
|
||||
|
||||
return dataset
|
||||
|
||||
# def __new__(cls, folder):
|
||||
# return tf.data.Dataset.from_generator(
|
||||
# cls._generator,
|
||||
# # output_signature = (tf.TensorSpec(shape = (), dtype = tf.string), ),
|
||||
# output_types=(tf.string),
|
||||
# args=(folder, )
|
||||
# )
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
@@ -360,7 +298,7 @@ if __name__ == '__main__':
|
||||
# print(d)
|
||||
|
||||
factory = DatasetFactory()
|
||||
signature_dynamic = { 'id': (), 'frame_nr': (),
|
||||
signature_dynamic = { 'id': (), 'frame_nr': (), 'garment': (),
|
||||
# 'pose_shape': (72 + 11,),
|
||||
# 'uv_posed': (3, 512, 256,),
|
||||
# 'uv_body_posed': (512, 256, 3),
|
||||
|
||||
+25
-1
@@ -75,7 +75,31 @@ def split_data(folder, ratio):
|
||||
f.writelines(valid_frames)
|
||||
|
||||
|
||||
pass
|
||||
def genereated_registered_and_mask(folder, is_skirt = False):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
for id in tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')]
|
||||
|
||||
assert len(garments) == 1
|
||||
|
||||
if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')):
|
||||
continue
|
||||
|
||||
garment = garments[0]
|
||||
|
||||
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
|
||||
map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1))
|
||||
verts = verts[:, map, :] # ---------------> This is the garment after registration
|
||||
raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1))
|
||||
mask = raw_mask > 0
|
||||
# boundaries = raw_mask == 2
|
||||
|
||||
writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True)
|
||||
np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask)
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user