luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

slowly update dataset (untested)

Browse Source
This commit is contained in:
luni
2023-02-07 11:58:28 +00:00
parent 4c884d52c4
commit bedfa9661a
3 changed files with 259 additions and 267 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/cloth3d_io.py
+30
View File
@@ -1,3 +1,4 @@
import os
import numpy as np import numpy as np
import scipy.io as sio import scipy.io as sio
from struct import pack, unpack from struct import pack, unpack
@@ -220,6 +221,35 @@ def _todict(matobj):
return dict 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 os
# import sys # import sys
# import numpy as np # import numpy as np
data/dataset.py
+50 -112
View File
@@ -3,58 +3,19 @@ import sys
import time import time
import cv2 import cv2
import random import random
import pickle
import numpy as np import numpy as np
import tensorflow as tf 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 from custom_utils.custom_utils import nest_conditions, batch_to_dict
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo
# 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, )
# )
class StaticDataset(): 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.store = dict()
self.signature = signature self.signature = signature
self.folder = folder self.folder = folder
@@ -69,8 +30,6 @@ class StaticDataset():
self.resolutions = resolutions self.resolutions = resolutions
self.filled = '' self.filled = ''
self.is_cape = is_cape self.is_cape = is_cape
if use_filled:
self.filled = '_filled'
if is_cape: if is_cape:
self.uv_masks = dict() self.uv_masks = dict()
@@ -82,10 +41,20 @@ class StaticDataset():
self.uv_masks[res] = down self.uv_masks[res] = down
def read(self, id, garment):
def read(self, id):
data = dict() 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: for res in self.resolutions:
res_string = '' res_string = ''
if res != 512: if res != 512:
@@ -99,19 +68,12 @@ class StaticDataset():
else: 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 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(): 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) data['mesh_static'] = read_garment_template(self.folder, id, garment)
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
return data return data
def __getitem__(self, id) -> dict: def __getitem__(self, id) -> dict:
if id not in self.store: if id not in self.store:
item = self.read(id) item = self.read(id)
@@ -122,11 +84,12 @@ class StaticDataset():
return self.store[id] return self.store[id]
def expand_batch(self, batch): def expand_batch(self, batch):
for key in self.keys: for key in self.keys:
batch[key] = tf.zeros(((0,) + self.signature[key])) batch[key] = tf.zeros(((0,) + self.signature[key]))
for i in range(batch['id'].shape[0]): 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: for key in self.keys:
batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0) batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
return batch return batch
@@ -136,20 +99,19 @@ class StaticDataset():
class DatasetIdx(tf.data.Dataset): class DatasetIdx(tf.data.Dataset):
def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None): def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
if frames_file is None: # if frames_file is None:
uv_static_folder = os.path.join(folder, 'pose_shape') # uv_static_folder = os.path.join(folder, 'pose_shape')
frames = sorted(os.listdir(uv_static_folder)) # frames = sorted(os.listdir(uv_static_folder))
if is_cape: # if is_cape:
frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] # frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
else: # else:
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] # frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
if testing: # if testing:
frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
frames = [frame.split('.')[0] + '_4' for frame in frames] # frames = [frame.split('.')[0] + '_4' for frame in frames]
if frames_file is not None:
frames = [] frames = []
with open(os.path.join(folder, frames_file), 'r') as f: with open(os.path.join(folder, frames_file), 'r') as f:
lines = f.readlines() lines = f.readlines()
@@ -162,7 +124,6 @@ class DatasetIdx(tf.data.Dataset):
return tf.data.Dataset.from_tensor_slices(frames) return tf.data.Dataset.from_tensor_slices(frames)
# class DatasetFactory(tf.data.Dataset):
class DatasetFactory(): class DatasetFactory():
def __init__(self): def __init__(self):
@@ -175,6 +136,8 @@ class DatasetFactory():
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'frame_nr': elif key == 'frame_nr':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) 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': elif key == 'mesh_curvature':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
else: else:
@@ -191,7 +154,6 @@ class DatasetFactory():
in_memory=False, in_memory=False,
do_repeat=True, do_repeat=True,
is_cape=False, is_cape=False,
use_one_hot=False,
use_no_orient=True, use_no_orient=True,
testing=False, testing=False,
frames_file=None): frames_file=None):
@@ -203,31 +165,33 @@ class DatasetFactory():
if is_cape: if is_cape:
separator = '.' 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: if use_no_orient == False:
assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!' 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): def read_item(frame):
frame_split = tf.strings.split(frame, sep=separator) 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, garment = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8"), frame_split[2].numpy().decode("utf-8")
# id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy()
# static_info = static_data[id]
item = dict() item = dict()
if 'id' in signature.keys(): if 'id' in signature.keys():
item['id'] = id item['id'] = id
if 'frame_nr' in signature.keys(): if 'frame_nr' in signature.keys():
item['frame_nr'] = frame_nr item['frame_nr'] = frame_nr
if 'garment' in signature.keys():
item['garment'] = garment
if 'pose_shape' in signature.keys(): if 'pose_shape' in signature.keys():
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy')) pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
if use_no_orient: if use_no_orient:
@@ -247,15 +211,8 @@ class DatasetFactory():
else: 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) 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 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(): if 'mesh_unposed' in signature.keys():
# print('one_hot') 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']
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']
item['mesh_unposed'] = mesh_unposed_no_shape item['mesh_unposed'] = mesh_unposed_no_shape
@@ -279,7 +236,7 @@ class DatasetFactory():
item['uv_curvature' + res_string] = uv_curvature item['uv_curvature' + res_string] = uv_curvature
if 'uv_unposed' + res_string in signature.keys(): 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 item['uv_unposed' + res_string] = uv_unposed_no_shape
if 'uv_body_posed' + res_string in signature.keys(): 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'] 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 item['uv_posed' + res_string] = uv_posed
# print('here')
# item = { **static_info, **item }
as_list = [] as_list = []
for key in signature.keys(): for key in signature.keys():
if key in item: if key in item:
@@ -309,10 +263,6 @@ class DatasetFactory():
return as_list 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: if in_memory:
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
@@ -328,20 +278,8 @@ class DatasetFactory():
if do_repeat: if do_repeat:
dataset = dataset.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 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__': if __name__ == '__main__':
@@ -360,7 +298,7 @@ if __name__ == '__main__':
# print(d) # print(d)
factory = DatasetFactory() factory = DatasetFactory()
signature_dynamic = { 'id': (), 'frame_nr': (), signature_dynamic = { 'id': (), 'frame_nr': (), 'garment': (),
# 'pose_shape': (72 + 11,), # 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,), # 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (512, 256, 3), # 'uv_body_posed': (512, 256, 3),
data/preprocessing.py
+25 -1
View File
@@ -75,7 +75,31 @@ def split_data(folder, ratio):
f.writelines(valid_frames) 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)