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
+176 -146
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
@@ -11,35 +12,35 @@ Inputs:
Output: Output:
- data: dictionary with .pc2/.pc16 file data - data: dictionary with .pc2/.pc16 file data
NOTE: 16-bit floats lose precision with high values (positive or negative), 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): 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' assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
data = {} data = {}
bytes = 2 if float16 else 4 bytes = 2 if float16 else 4
dtype = np.float16 if float16 else np.float32 dtype = np.float16 if float16 else np.float32
with open(file, 'rb') as f: with open(file, 'rb') as f:
# Header # Header
data['sign'] = f.read(12) data['sign'] = f.read(12)
# data['version'] = int.from_bytes(f.read(4), 'little') # data['version'] = int.from_bytes(f.read(4), 'little')
data['version'] = unpack('<i', f.read(4))[0] data['version'] = unpack('<i', f.read(4))[0]
# Num points # Num points
# data['nPoints'] = int.from_bytes(f.read(4), 'little') # data['nPoints'] = int.from_bytes(f.read(4), 'little')
data['nPoints'] = unpack('<i', f.read(4))[0] data['nPoints'] = unpack('<i', f.read(4))[0]
# Start frame # Start frame
data['startFrame'] = unpack('f', f.read(4)) data['startFrame'] = unpack('f', f.read(4))
# Sample rate # Sample rate
data['sampleRate'] = unpack('f', f.read(4)) data['sampleRate'] = unpack('f', f.read(4))
# Number of samples # Number of samples
# data['nSamples'] = int.from_bytes(f.read(4), 'little') # data['nSamples'] = int.from_bytes(f.read(4), 'little')
data['nSamples'] = unpack('<i', f.read(4))[0] data['nSamples'] = unpack('<i', f.read(4))[0]
# Animation data # Animation data
size = data['nPoints']*data['nSamples']*3*bytes size = data['nPoints']*data['nSamples']*3*bytes
data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32) data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3) data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
return data return data
""" """
Reads an specific frame of PC2/PC16 files Reads an specific frame of PC2/PC16 files
Inputs: Inputs:
@@ -50,29 +51,29 @@ Output:
- T: mesh vertex data at specified frame - T: mesh vertex data at specified frame
""" """
def readPC2Frame(file, frame, float16=False): def readPC2Frame(file, frame, float16=False):
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' 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' assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer'
bytes = 2 if float16 else 4 bytes = 2 if float16 else 4
dtype = np.float16 if float16 else np.float32 dtype = np.float16 if float16 else np.float32
with open(file,'rb') as f: with open(file,'rb') as f:
# Num points # Num points
f.seek(16) f.seek(16)
# nPoints = int.from_bytes(f.read(4), 'little') # nPoints = int.from_bytes(f.read(4), 'little')
nPoints = unpack('<i', f.read(4))[0] nPoints = unpack('<i', f.read(4))[0]
# Number of samples # Number of samples
f.seek(28) f.seek(28)
# nSamples = int.from_bytes(f.read(4), 'little') # nSamples = int.from_bytes(f.read(4), 'little')
nSamples = unpack('<i', f.read(4))[0] nSamples = unpack('<i', f.read(4))[0]
if frame > nSamples: if frame > nSamples:
print("Frame index outside size") print("Frame index outside size")
print("\tN. frame: " + str(frame)) print("\tN. frame: " + str(frame))
print("\tN. samples: " + str(nSamples)) print("\tN. samples: " + str(nSamples))
return return
# Read frame # Read frame
size = nPoints * 3 * bytes size = nPoints * 3 * bytes
f.seek(size * frame, 1) # offset from current '1' f.seek(size * frame, 1) # offset from current '1'
T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32) T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
return T.reshape(nPoints, 3) return T.reshape(nPoints, 3)
""" """
Writes PC2 and PC16 files Writes PC2 and PC16 files
@@ -82,20 +83,20 @@ Inputs:
- float16: False for writing as PC2 file, True for PC16 - float16: False for writing as PC2 file, True for PC16
This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1 This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1
NOTE: 16-bit floats lose precision with high values (positive or negative), 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): 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' 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) if float16: V = V.astype(np.float16)
else: V = V.astype(np.float32) else: V = V.astype(np.float32)
with open(file, 'wb') as f: with open(file, 'wb') as f:
# Create the header # Create the header
headerFormat='<12siiffi' headerFormat='<12siiffi'
headerStr = pack(headerFormat, b'POINTCACHE2\0', headerStr = pack(headerFormat, b'POINTCACHE2\0',
1, V.shape[1], 0, 1, V.shape[0]) 1, V.shape[1], 0, 1, V.shape[0])
f.write(headerStr) f.write(headerStr)
# Write vertices # Write vertices
f.write(V.tobytes()) f.write(V.tobytes())
""" """
Reads OBJ files 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 If no UV map data in .obj file, it shall return Vt=None and Ft=None
""" """
def readOBJ(file): def readOBJ(file):
V, Vt, F, Ft = [], [], [], [] V, Vt, F, Ft = [], [], [], []
with open(file, 'r') as f: with open(file, 'r') as f:
T = f.readlines() T = f.readlines()
for t in T: for t in T:
# 3D vertex # 3D vertex
if t.startswith('v '): if t.startswith('v '):
v = [float(n) for n in t.replace('v ','').split(' ')] v = [float(n) for n in t.replace('v ','').split(' ')]
V += [v] V += [v]
# UV vertex # UV vertex
elif t.startswith('vt '): elif t.startswith('vt '):
v = [float(n) for n in t.replace('vt ','').split(' ')] v = [float(n) for n in t.replace('vt ','').split(' ')]
Vt += [v] Vt += [v]
# Face # Face
elif t.startswith('f '): elif t.startswith('f '):
idx = [n.split('/') for n in t.replace('f ','').split(' ')] idx = [n.split('/') for n in t.replace('f ','').split(' ')]
f = [int(n[0]) - 1 for n in idx] f = [int(n[0]) - 1 for n in idx]
F += [f] F += [f]
# UV face # UV face
if '/' in t: if '/' in t:
f = [int(n[1]) - 1 for n in idx] f = [int(n[1]) - 1 for n in idx]
Ft += [f] Ft += [f]
V = np.array(V, np.float32) V = np.array(V, np.float32)
Vt = np.array(Vt, 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' 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 else: Vt, Ft = None, None
return V, F, Vt, Ft 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 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): def writeOBJ(file, V, F, Vt=None, Ft=None):
if not Vt is 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' 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: with open(file, 'w') as file:
# Vertices # Vertices
for v in V: for v in V:
line = 'v ' + ' '.join([str(_) for _ in v]) + '\n' line = 'v ' + ' '.join([str(_) for _ in v]) + '\n'
file.write(line) file.write(line)
# UV verts # UV verts
if not Vt is None: if not Vt is None:
for v in Vt: for v in Vt:
line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n' line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n'
file.write(line) file.write(line)
# 3D Faces / UV faces # 3D Faces / UV faces
if Ft: if Ft:
F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)] F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)]
else: else:
F = [[str(i + 1) for i in f] for f in F] F = [[str(i + 1) for i in f] for f in F]
for f in F: for f in F:
line = 'f ' + ' '.join(f) + '\n' line = 'f ' + ' '.join(f) + '\n'
file.write(line) file.write(line)
def loadInfo(filename): def loadInfo(filename):
''' '''
this function should be called instead of direct sio.loadmat this function should be called instead of direct sio.loadmat
as it cures the problem of not properly recovering python dictionaries as it cures the problem of not properly recovering python dictionaries
from mat files. It calls the function check keys to cure all entries from mat files. It calls the function check keys to cure all entries
which are still mat-objects which are still mat-objects
''' '''
data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True) data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True)
del data['__globals__'] del data['__globals__']
del data['__header__'] del data['__header__']
del data['__version__'] del data['__version__']
return _check_keys(data) return _check_keys(data)
def _check_keys(dict): def _check_keys(dict):
''' '''
checks if entries in dictionary are mat-objects. If yes checks if entries in dictionary are mat-objects. If yes
todict is called to change them to nested dictionaries todict is called to change them to nested dictionaries
''' '''
for key in dict: for key in dict:
if isinstance(dict[key], sio.matlab.mio5_params.mat_struct): if isinstance(dict[key], sio.matlab.mio5_params.mat_struct):
dict[key] = _todict(dict[key]) dict[key] = _todict(dict[key])
return dict return dict
def _todict(matobj): def _todict(matobj):
''' '''
A recursive function which constructs from matobjects nested dictionaries A recursive function which constructs from matobjects nested dictionaries
''' '''
dict = {} dict = {}
for strg in matobj._fieldnames: for strg in matobj._fieldnames:
elem = matobj.__dict__[strg] elem = matobj.__dict__[strg]
if isinstance(elem, sio.matlab.mio5_params.mat_struct): if isinstance(elem, sio.matlab.mio5_params.mat_struct):
dict[strg] = _todict(elem) dict[strg] = _todict(elem)
elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in 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) dict[strg] = [None] * len(elem)
for i,item in enumerate(elem): for i,item in enumerate(elem):
if isinstance(item, sio.matlab.mio5_params.mat_struct): if isinstance(item, sio.matlab.mio5_params.mat_struct):
dict[strg][i] = _todict(item) dict[strg][i] = _todict(item)
else: else:
dict[strg][i] = item dict[strg][i] = item
else: else:
dict[strg] = elem dict[strg] = elem
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
data/dataset.py
+58 -120
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,25 +99,24 @@ 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() lines = [line.strip() for line in lines]
lines = [line.strip() for line in lines] frames = lines
frames = lines
if shuffle: if shuffle:
random.shuffle(frames) random.shuffle(frames)
@@ -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:
@@ -190,8 +153,7 @@ class DatasetFactory():
resolutions=[512], resolutions=[512],
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 = '' uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
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: 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,11 +263,7 @@ 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,
inp=[i], inp=[i],
@@ -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)