slowly update dataset (untested)
This commit is contained in:
+176
-146
@@ -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('<i', f.read(4))[0]
|
||||
# Num points
|
||||
# data['nPoints'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nPoints'] = unpack('<i', f.read(4))[0]
|
||||
# Start frame
|
||||
data['startFrame'] = unpack('f', f.read(4))
|
||||
# Sample rate
|
||||
data['sampleRate'] = unpack('f', f.read(4))
|
||||
# Number of samples
|
||||
# data['nSamples'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nSamples'] = unpack('<i', f.read(4))[0]
|
||||
# Animation data
|
||||
size = data['nPoints']*data['nSamples']*3*bytes
|
||||
data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
|
||||
data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
|
||||
|
||||
return data
|
||||
|
||||
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('<i', f.read(4))[0]
|
||||
# Num points
|
||||
# data['nPoints'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nPoints'] = unpack('<i', f.read(4))[0]
|
||||
# Start frame
|
||||
data['startFrame'] = unpack('f', f.read(4))
|
||||
# Sample rate
|
||||
data['sampleRate'] = unpack('f', f.read(4))
|
||||
# Number of samples
|
||||
# data['nSamples'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nSamples'] = unpack('<i', f.read(4))[0]
|
||||
# Animation data
|
||||
size = data['nPoints']*data['nSamples']*3*bytes
|
||||
data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
|
||||
data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
|
||||
|
||||
return data
|
||||
|
||||
"""
|
||||
Reads an specific frame of PC2/PC16 files
|
||||
Inputs:
|
||||
@@ -50,29 +51,29 @@ Output:
|
||||
- T: mesh vertex data at specified frame
|
||||
"""
|
||||
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 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('<i', f.read(4))[0]
|
||||
# Number of samples
|
||||
f.seek(28)
|
||||
# nSamples = int.from_bytes(f.read(4), 'little')
|
||||
nSamples = unpack('<i', f.read(4))[0]
|
||||
if frame > 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('<i', f.read(4))[0]
|
||||
# Number of samples
|
||||
f.seek(28)
|
||||
# nSamples = int.from_bytes(f.read(4), 'little')
|
||||
nSamples = unpack('<i', f.read(4))[0]
|
||||
if frame > 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
|
||||
|
||||
+58
-120
@@ -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),
|
||||
|
||||
+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