save progress
This commit is contained in:
+111
-21
@@ -96,6 +96,73 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
|
||||
return result
|
||||
|
||||
|
||||
def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
|
||||
batch = verts.shape[0]
|
||||
verts_uv = np.zeros((batch, res * 2, res, 3))
|
||||
uv_values = np.zeros((batch, uv_map.shape[0], 3))
|
||||
uv_values[:, m2uv[:, 0]] = verts
|
||||
uv_values[:, m2uv[:, 1]] += verts
|
||||
uv_values = uv_values[:, uv_map_mask, :]
|
||||
uv_values /= uv_counter[np.newaxis, :, np.newaxis]
|
||||
return_garment_uv_mask = False
|
||||
# test
|
||||
verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
|
||||
|
||||
if mask is not None:
|
||||
garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
|
||||
garment_mask[m2uv[:, 0]] = mask
|
||||
garment_mask[m2uv[:, 1]] += mask
|
||||
garment_mask = garment_mask[uv_map_mask] # add garment mask
|
||||
else:
|
||||
garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
|
||||
|
||||
unique_uv_values = uv_values[:, garment_mask, :]
|
||||
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
|
||||
rbf_dest_current = rbf_dest
|
||||
|
||||
# add back mesh vertices that are not part of the garment
|
||||
uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
|
||||
rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
|
||||
rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
|
||||
|
||||
if garment_uv_mask is None and mask is not None:
|
||||
garment_uv_mask = np.zeros((res * 2, res))
|
||||
garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
|
||||
garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
|
||||
garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
|
||||
return_garment_uv_mask = True
|
||||
|
||||
if garment_uv_mask is not None:
|
||||
rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
|
||||
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
final = []
|
||||
for i in range(batch):
|
||||
print('processing frame: ' + str(i))
|
||||
if method == 'rbf':
|
||||
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear')
|
||||
elif method == 'nearest':
|
||||
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:])
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
result = np.zeros((res * 2, res, 3), np.float32)
|
||||
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
|
||||
result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[i,:,:]
|
||||
|
||||
if garment_uv_mask is not None:
|
||||
result *= garment_uv_mask[:,:,np.newaxis]
|
||||
|
||||
final.append(result)
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return final, garment_uv_mask
|
||||
else:
|
||||
return final
|
||||
|
||||
|
||||
def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None):
|
||||
batch = len(verts)
|
||||
|
||||
@@ -388,15 +455,25 @@ if __name__ == '__main__':
|
||||
with open(os.path.join(folder, frames_file), 'r') as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
frames = lines
|
||||
frames = lines[int(args.start):int(args.end)]
|
||||
ids = [frame.split()[0] for frame in frames]
|
||||
ids = sorted(list(set(ids)))[int(args.start):int(args.end)]
|
||||
ids = sorted(list(set(ids)))
|
||||
|
||||
shapedirs = get_shapedirs()
|
||||
|
||||
root = '/'.join(folder.split('/')[:-1])
|
||||
type = folder.split('/')[-1]
|
||||
h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
|
||||
h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
|
||||
# h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
|
||||
|
||||
extras_folder = os.path.join(root, f'{type}_extras')
|
||||
os.makedirs(extras_folder, exist_ok=True)
|
||||
os.makedirs(os.path.join(extras_folder, 'uv_posed'), exist_ok=True)
|
||||
os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
|
||||
os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
|
||||
os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
|
||||
os.makedirs(os.path.join(extras_folder, 'uv_body_posed'), exist_ok=True)
|
||||
os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
|
||||
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
@@ -416,14 +493,15 @@ if __name__ == '__main__':
|
||||
mask = read_garment_mask(folder, id, garment)
|
||||
|
||||
if 'mesh_posed' in data_to_convert:
|
||||
mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
|
||||
# mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
|
||||
mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
|
||||
if 'mesh_unposed' in data_to_convert:
|
||||
raise NotImplementedError()
|
||||
mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
|
||||
if 'mesh_static' in data_to_convert:
|
||||
mesh_static = read_garment_template(folder, id, garment, mask)
|
||||
if 'mesh_curvature' in data_to_convert:
|
||||
mesh_curvature = read_curvatures(h5file, id, frames, neighbourhood=None)
|
||||
mesh_curvature = read_curvatures(h5file_curv, id, garment, frames, neighbourhood=None)
|
||||
if 'body_posed' in data_to_convert:
|
||||
body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
|
||||
|
||||
@@ -454,35 +532,45 @@ if __name__ == '__main__':
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
uv_static, uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
|
||||
uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
|
||||
uv_static = uv_static.astype(np.float32)
|
||||
|
||||
if 'mesh_posed' in data_to_convert:
|
||||
uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
if 'mesh_unposed' in data_to_convert:
|
||||
uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
raise NotImplementedError()
|
||||
uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
if 'mesh_curvature' in data_to_convert:
|
||||
uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
|
||||
if 'body_posed' in data_to_convert:
|
||||
uv_body_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
|
||||
|
||||
|
||||
|
||||
print(time.perf_counter() - start)
|
||||
|
||||
# mesh_static
|
||||
save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
|
||||
save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
|
||||
# save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
|
||||
# save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
|
||||
np.savez_compressed(os.path.join(extras_folder, 'uv_static', f'{id}_{garment}.npz'), uv_static=uv_static)
|
||||
np.savez_compressed(os.path.join(extras_folder, 'uv_mask', f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
|
||||
|
||||
if 'mesh_posed' in data_to_convert:
|
||||
for i in range(uv_posed.shape[0]):
|
||||
save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
|
||||
for i in range(len(uv_posed)):
|
||||
np.savez_compressed(os.path.join(extras_folder, 'uv_posed' + orient_mod, f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
|
||||
# save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
|
||||
if 'mesh_unposed' in data_to_convert:
|
||||
raise NotImplementedError()
|
||||
for i in range(uv_unposed.shape[0]):
|
||||
save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
|
||||
if 'mesh_curvature' in data_to_convert:
|
||||
for i in range(uv_curvature.shape[0]):
|
||||
save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
|
||||
for i in range(len(uv_curvature)):
|
||||
np.savez_compressed(os.path.join(extras_folder, 'uv_curvature', f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
|
||||
# save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
|
||||
if 'body_posed' in data_to_convert:
|
||||
for i in range(uv_body_posed.shape[0]):
|
||||
save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
|
||||
for i in range(len(uv_body_posed)):
|
||||
np.savez_compressed(os.path.join(extras_folder, 'uv_body_posed' + orient_mod, f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
|
||||
# save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
|
||||
|
||||
|
||||
# start = time.perf_counter()
|
||||
@@ -500,9 +588,9 @@ if __name__ == '__main__':
|
||||
# for i in range(uv_gt.shape[0]):
|
||||
# np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :])
|
||||
|
||||
os.makedirs(os.path.join(folder, f'uv_curvature_{neighbourhood}_2'), exist_ok=True)
|
||||
for i in range(uv_curvature.shape[0]):
|
||||
np.savez_compressed(os.path.join(folder, f'uv_curvature_{neighbourhood}_2', id + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
|
||||
# os.makedirs(os.path.join(folder, f'uv_curvature_{neighbourhood}_2'), exist_ok=True)
|
||||
# for i in range(uv_curvature.shape[0]):
|
||||
# np.savez_compressed(os.path.join(folder, f'uv_curvature_{neighbourhood}_2', id + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
|
||||
|
||||
# os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
|
||||
# for i in range(uv_unposed.shape[0]):
|
||||
@@ -529,3 +617,5 @@ if __name__ == '__main__':
|
||||
|
||||
pass
|
||||
|
||||
|
||||
h5file_curv.close()
|
||||
Reference in New Issue
Block a user