diff --git a/data/cloth3d_io.py b/data/cloth3d_io.py index 838b794..057b7c3 100644 --- a/data/cloth3d_io.py +++ b/data/cloth3d_io.py @@ -1,3 +1,4 @@ +import os import numpy as np import scipy.io as sio from struct import pack, unpack @@ -11,35 +12,35 @@ Inputs: Output: - data: dictionary with .pc2/.pc16 file data NOTE: 16-bit floats lose precision with high values (positive or negative), - we do not recommend using this format for data outside range [-2, 2] + we do not recommend using this format for data outside range [-2, 2] """ def readPC2(file, float16=False): - assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' - data = {} - bytes = 2 if float16 else 4 - dtype = np.float16 if float16 else np.float32 - with open(file, 'rb') as f: - # Header - data['sign'] = f.read(12) - # data['version'] = int.from_bytes(f.read(4), 'little') - data['version'] = unpack('= 0 and isinstance(frame,int), 'Frame must be a positive integer' - bytes = 2 if float16 else 4 - dtype = np.float16 if float16 else np.float32 - with open(file,'rb') as f: - # Num points - f.seek(16) - # nPoints = int.from_bytes(f.read(4), 'little') - nPoints = unpack(' nSamples: - print("Frame index outside size") - print("\tN. frame: " + str(frame)) - print("\tN. samples: " + str(nSamples)) - return - # Read frame - size = nPoints * 3 * bytes - f.seek(size * frame, 1) # offset from current '1' - T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32) - return T.reshape(nPoints, 3) + assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' + assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer' + bytes = 2 if float16 else 4 + dtype = np.float16 if float16 else np.float32 + with open(file,'rb') as f: + # Num points + f.seek(16) + # nPoints = int.from_bytes(f.read(4), 'little') + nPoints = unpack(' nSamples: + print("Frame index outside size") + print("\tN. frame: " + str(frame)) + print("\tN. samples: " + str(nSamples)) + return + # Read frame + size = nPoints * 3 * bytes + f.seek(size * frame, 1) # offset from current '1' + T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32) + return T.reshape(nPoints, 3) """ Writes PC2 and PC16 files @@ -82,20 +83,20 @@ Inputs: - float16: False for writing as PC2 file, True for PC16 This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1 NOTE: 16-bit floats lose precision with high values (positive or negative), - we do not recommend using this format for data outside range [-2, 2] + we do not recommend using this format for data outside range [-2, 2] """ def writePC2(file, V, float16=False): - assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' - if float16: V = V.astype(np.float16) - else: V = V.astype(np.float32) - with open(file, 'wb') as f: - # Create the header - headerFormat='<12siiffi' - headerStr = pack(headerFormat, b'POINTCACHE2\0', - 1, V.shape[1], 0, 1, V.shape[0]) - f.write(headerStr) - # Write vertices - f.write(V.tobytes()) + assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' + if float16: V = V.astype(np.float16) + else: V = V.astype(np.float32) + with open(file, 'wb') as f: + # Create the header + headerFormat='<12siiffi' + headerStr = pack(headerFormat, b'POINTCACHE2\0', + 1, V.shape[1], 0, 1, V.shape[0]) + f.write(headerStr) + # Write vertices + f.write(V.tobytes()) """ Reads OBJ files @@ -111,32 +112,32 @@ Correspondence between mesh and UV map is implicit in F to Ft correspondences If no UV map data in .obj file, it shall return Vt=None and Ft=None """ def readOBJ(file): - V, Vt, F, Ft = [], [], [], [] - with open(file, 'r') as f: - T = f.readlines() - for t in T: - # 3D vertex - if t.startswith('v '): - v = [float(n) for n in t.replace('v ','').split(' ')] - V += [v] - # UV vertex - elif t.startswith('vt '): - v = [float(n) for n in t.replace('vt ','').split(' ')] - Vt += [v] - # Face - elif t.startswith('f '): - idx = [n.split('/') for n in t.replace('f ','').split(' ')] - f = [int(n[0]) - 1 for n in idx] - F += [f] - # UV face - if '/' in t: - f = [int(n[1]) - 1 for n in idx] - Ft += [f] - V = np.array(V, np.float32) - Vt = np.array(Vt, np.float32) - if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces' - else: Vt, Ft = None, None - return V, F, Vt, Ft + V, Vt, F, Ft = [], [], [], [] + with open(file, 'r') as f: + T = f.readlines() + for t in T: + # 3D vertex + if t.startswith('v '): + v = [float(n) for n in t.replace('v ','').split(' ')] + V += [v] + # UV vertex + elif t.startswith('vt '): + v = [float(n) for n in t.replace('vt ','').split(' ')] + Vt += [v] + # Face + elif t.startswith('f '): + idx = [n.split('/') for n in t.replace('f ','').split(' ')] + f = [int(n[0]) - 1 for n in idx] + F += [f] + # UV face + if '/' in t: + f = [int(n[1]) - 1 for n in idx] + Ft += [f] + V = np.array(V, np.float32) + Vt = np.array(Vt, np.float32) + if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces' + else: Vt, Ft = None, None + return V, F, Vt, Ft @@ -153,71 +154,100 @@ Correspondence between mesh and UV map is implicit in F to Ft correspondences If no UV map data as input, it will write only 3D data in .obj file """ def writeOBJ(file, V, F, Vt=None, Ft=None): - if not Vt is None: - assert len(F) == len(Ft), 'Inconsistent data, mesh and UV map do not have the same number of faces' - - with open(file, 'w') as file: - # Vertices - for v in V: - line = 'v ' + ' '.join([str(_) for _ in v]) + '\n' - file.write(line) - # UV verts - if not Vt is None: - for v in Vt: - line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n' - file.write(line) - # 3D Faces / UV faces - if Ft: - F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)] - else: - F = [[str(i + 1) for i in f] for f in F] - for f in F: - line = 'f ' + ' '.join(f) + '\n' - file.write(line) + if not Vt is None: + assert len(F) == len(Ft), 'Inconsistent data, mesh and UV map do not have the same number of faces' + + with open(file, 'w') as file: + # Vertices + for v in V: + line = 'v ' + ' '.join([str(_) for _ in v]) + '\n' + file.write(line) + # UV verts + if not Vt is None: + for v in Vt: + line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n' + file.write(line) + # 3D Faces / UV faces + if Ft: + F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)] + else: + F = [[str(i + 1) for i in f] for f in F] + for f in F: + line = 'f ' + ' '.join(f) + '\n' + file.write(line) def loadInfo(filename): - ''' - this function should be called instead of direct sio.loadmat - as it cures the problem of not properly recovering python dictionaries - from mat files. It calls the function check keys to cure all entries - which are still mat-objects - ''' - data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True) - del data['__globals__'] - del data['__header__'] - del data['__version__'] - return _check_keys(data) + ''' + this function should be called instead of direct sio.loadmat + as it cures the problem of not properly recovering python dictionaries + from mat files. It calls the function check keys to cure all entries + which are still mat-objects + ''' + data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True) + del data['__globals__'] + del data['__header__'] + del data['__version__'] + return _check_keys(data) def _check_keys(dict): - ''' - checks if entries in dictionary are mat-objects. If yes - todict is called to change them to nested dictionaries - ''' - for key in dict: - if isinstance(dict[key], sio.matlab.mio5_params.mat_struct): - dict[key] = _todict(dict[key]) - return dict + ''' + checks if entries in dictionary are mat-objects. If yes + todict is called to change them to nested dictionaries + ''' + for key in dict: + if isinstance(dict[key], sio.matlab.mio5_params.mat_struct): + dict[key] = _todict(dict[key]) + return dict def _todict(matobj): - ''' - A recursive function which constructs from matobjects nested dictionaries - ''' - dict = {} - for strg in matobj._fieldnames: - elem = matobj.__dict__[strg] - if isinstance(elem, sio.matlab.mio5_params.mat_struct): - dict[strg] = _todict(elem) - elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in elem]): - dict[strg] = [None] * len(elem) - for i,item in enumerate(elem): - if isinstance(item, sio.matlab.mio5_params.mat_struct): - dict[strg][i] = _todict(item) - else: - dict[strg][i] = item - else: - dict[strg] = elem - return dict + ''' + A recursive function which constructs from matobjects nested dictionaries + ''' + dict = {} + for strg in matobj._fieldnames: + elem = matobj.__dict__[strg] + if isinstance(elem, sio.matlab.mio5_params.mat_struct): + dict[strg] = _todict(elem) + elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in elem]): + dict[strg] = [None] * len(elem) + for i,item in enumerate(elem): + if isinstance(item, sio.matlab.mio5_params.mat_struct): + dict[strg][i] = _todict(item) + else: + dict[strg][i] = item + else: + dict[strg] = elem + 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 diff --git a/data/dataset.py b/data/dataset.py index d2ee42d..d583bce 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -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,25 +99,24 @@ 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() - lines = [line.strip() for line in lines] - frames = lines + frames = [] + with open(os.path.join(folder, frames_file), 'r') as f: + lines = f.readlines() + lines = [line.strip() for line in lines] + frames = lines if shuffle: random.shuffle(frames) @@ -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: @@ -190,8 +153,7 @@ class DatasetFactory(): resolutions=[512], in_memory=False, do_repeat=True, - is_cape=False, - use_one_hot=False, + is_cape=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' + 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,11 +263,7 @@ 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, inp=[i], @@ -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), diff --git a/data/preprocessing.py b/data/preprocessing.py index f48b859..fc489e8 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -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)