prepare for cape
This commit is contained in:
+470
-13
@@ -1,14 +1,15 @@
|
||||
from glob import glob
|
||||
import os
|
||||
import tqdm
|
||||
import shutil
|
||||
import numpy as np
|
||||
import plotly.express as ex
|
||||
import plotly.graph_objects as go
|
||||
from simpleicp import PointCloud, SimpleICP
|
||||
from scipy.spatial.distance import cdist
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
|
||||
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '2'
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
|
||||
import tensorflow as tf
|
||||
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
@@ -24,26 +25,44 @@ if gpus:
|
||||
print(e)
|
||||
|
||||
from custom_utils.smpl import SMPL
|
||||
# from custom_utils.smpl import SMPL
|
||||
|
||||
|
||||
folder = '/home/workspace/data/dataset/cape_release'
|
||||
# smpl = SMPL(orig_smpl=True)
|
||||
|
||||
|
||||
def get_frames():
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
# sequences = np.load()
|
||||
test = np.load(os.path.join(folder, 'sequences/03383/shortlong_ATUsquat/shortlong_ATUsquat.000105.npz'))
|
||||
gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True)
|
||||
test = np.load(os.path.join(folder, 'train/male/02474/longshort_ATUsquat/longshort_ATUsquat.000100.npz'))
|
||||
faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy'))
|
||||
min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy'))
|
||||
min_body = np.load(os.path.join(folder, 'minimal_body_shape/02474/02474_minimal.npy'))
|
||||
shape = np.load(os.path.join(folder, 'minimal_body_shape/02474/betas_2.npy'))
|
||||
offsets_cano = test['v_cano'] - min_body
|
||||
pose = test['pose']
|
||||
mask = np.abs(offsets_cano.sum(axis=-1)) < 1e-2
|
||||
pose = test['pose'].astype(np.float32)
|
||||
shape = np.concatenate((shape, [0] if gender_table['02474'] == 'female' else [1])).astype(np.float32)
|
||||
posed_body = smpl.run(smpl.run_rodrigues(pose[np.newaxis,:]), shape[np.newaxis,:], remove=False)[0].numpy()[0]
|
||||
transl = test['transl']
|
||||
posed_body = posed_body + transl[np.newaxis, :]
|
||||
offset_posed = test['v_posed'] - posed_body
|
||||
mask = np.abs(offsets_cano.sum(axis=-1)) > 1e-2
|
||||
|
||||
smpl = SMPL()
|
||||
smpl.run_garment()
|
||||
zero_body = smpl.run(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape[np.newaxis,:], remove=True)[0]
|
||||
zero_body_root, J, _ = smpl.run(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape[np.newaxis,:], remove=True)
|
||||
# _, J0, _ = smpl.run(smpl.run_rodrigues(pose[np.newaxis,:]), shape[np.newaxis,:], remove=True, subtract_root=True)
|
||||
|
||||
# verts, faces = _maskFaces(test['v_posed'], faces, mask)
|
||||
posed_body = smpl.run_garment(pose[np.newaxis,:], shape[np.newaxis,:], zero_body)[0].numpy()
|
||||
posed_body_root = smpl.run_garment(pose[np.newaxis,:], shape[np.newaxis,:], zero_body_root)
|
||||
posed_body_root = posed_body_root[0].numpy()
|
||||
# verts, faces = _maskFaces(test['v_cano'], faces, mask)
|
||||
|
||||
x = test['v_cano']
|
||||
y = min_body
|
||||
# x = test['v_cano']
|
||||
# y = min_body
|
||||
x = min_body[smpl.to_keep, :] - J[0,0,:][np.newaxis, :]
|
||||
y = zero_body_root[0].numpy()
|
||||
faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
@@ -171,8 +190,446 @@ def get_frames_pop():
|
||||
print()
|
||||
|
||||
|
||||
def split_data():
|
||||
from CAPE.data.dataset_configs import dataset_config_dicts
|
||||
count = 0
|
||||
cut_first = 2
|
||||
sample_rate = 1
|
||||
|
||||
src = '/home/workspace/local-data/temp/cape_release/sequences'
|
||||
dst = '/home/workspace/local-data/cape/' + phase + '/' + gender
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
|
||||
config = dataset_config_dicts[gender]
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
if id not in config[phase + '_' + 'subjs']:
|
||||
continue
|
||||
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
if sequence_name in config['exclude_seqs']:
|
||||
continue
|
||||
if sequence_name not in config[phase + '_' + 'seqs']:
|
||||
continue
|
||||
|
||||
cloth = sequence.split('_')[0]
|
||||
if cloth not in config[phase + '_' + 'cloth']:
|
||||
continue
|
||||
|
||||
if [id, cloth] in config['exclude_cases']:
|
||||
continue
|
||||
|
||||
frames = sorted(glob(os.path.join(garment_dir, sequence, '*.npz')))
|
||||
if len(frames) == 0:
|
||||
print('{} no files here, skipping ..\n'.format(id+'_'+sequence))
|
||||
return
|
||||
frames = frames[cut_first: -cut_first : sample_rate] # take every k frame
|
||||
|
||||
if len(frames) == 0:
|
||||
continue
|
||||
|
||||
os.makedirs(os.path.join(dst, id, sequence), exist_ok=True)
|
||||
for frame in frames:
|
||||
try:
|
||||
np.load(frame)
|
||||
except: # skip and report the broken file if any
|
||||
print('{} broken, skipping..'.format(frame))
|
||||
continue
|
||||
|
||||
frame_name = frame.split('/')[-1]
|
||||
count += 1
|
||||
print(id + '_' + frame_name)
|
||||
|
||||
test_name = frame.split('/')
|
||||
test_name.index(id)
|
||||
test_name.index(sequence)
|
||||
test_name.index(frame_name)
|
||||
|
||||
shutil.move(frame, os.path.join(dst, id, sequence, frame_name))
|
||||
|
||||
print(count)
|
||||
|
||||
|
||||
def verify_split_data():
|
||||
from CAPE.data.dataset_configs import dataset_config_dicts
|
||||
count = 0
|
||||
cut_first = 2
|
||||
sample_rate = 1
|
||||
phase = 'test'
|
||||
gender = 'male'
|
||||
src = '/home/workspace/data/dataset/cape_release/sequences'
|
||||
dst = '/home/workspace/data/dataset/cape_release/' + phase + '/' + gender
|
||||
|
||||
config = dataset_config_dicts[gender]
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
if id not in config[phase + '_' + 'subjs']:
|
||||
continue
|
||||
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
if sequence_name in config['exclude_seqs']:
|
||||
continue
|
||||
if sequence_name not in config[phase + '_' + 'seqs']:
|
||||
continue
|
||||
|
||||
cloth = sequence.split('_')[0]
|
||||
if cloth not in config[phase + '_' + 'cloth']:
|
||||
continue
|
||||
|
||||
if [id, cloth] in config['exclude_cases']:
|
||||
continue
|
||||
|
||||
frames = sorted(glob(os.path.join(garment_dir, sequence, '*.npz')))
|
||||
if len(frames) == 0:
|
||||
print('{} no files here, skipping ..\n'.format(id+'_'+sequence))
|
||||
return
|
||||
frames = frames[cut_first: -cut_first : sample_rate] # take every k frame
|
||||
|
||||
if len(frames) == 0:
|
||||
continue
|
||||
|
||||
# os.makedirs(os.path.join(dst, id, sequence), exist_ok=True)
|
||||
for frame in frames:
|
||||
try:
|
||||
np.load(frame)
|
||||
except: # skip and report the broken file if any
|
||||
print('{} broken, skipping..'.format(frame))
|
||||
continue
|
||||
|
||||
frame_name = frame.split('/')[-1]
|
||||
count += 1
|
||||
print(id + '_' + frame_name)
|
||||
|
||||
# shutil.move(frame, os.path.join(dst, id, sequence, frame_name))
|
||||
|
||||
print(count)
|
||||
|
||||
|
||||
def merge_pose_shape():
|
||||
src = os.path.join('/home/workspace/local-data/cape', phase, gender)
|
||||
dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'pose_shape')
|
||||
gender_table = np.load('/home/workspace/local-data/temp/cape_release/misc/subj_genders.pkl', allow_pickle=True)
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
# sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
# cloth = sequence.split('_')[0]
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(garment_dir, sequence)))
|
||||
|
||||
for frame in frames:
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
|
||||
output_id = id + '_' + frame_name
|
||||
# continue
|
||||
|
||||
data = np.load(os.path.join(garment_dir, sequence, frame))
|
||||
pose = data['pose']
|
||||
shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release_old/minimal_body_shape', id, 'betas_2.npy'))
|
||||
seq_gender = [0] if gender_table[id] == 'female' else [1]
|
||||
shape = np.concatenate((shape, seq_gender)).astype(np.float32)
|
||||
# body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=False)[0][0]
|
||||
pose_shape = np.concatenate((pose, shape))
|
||||
|
||||
np.save(os.path.join(dst, id + '_' + frame_name + '.npy'), pose_shape)
|
||||
|
||||
print(id + '_' + frame)
|
||||
|
||||
|
||||
def create_posed_body():
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
src = os.path.join('/home/workspace/local-data/cape', phase, gender)
|
||||
dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_posed_body_no_orient')
|
||||
gender_table = np.load('/home/workspace/data/dataset/cape_release_old/misc/subj_genders.pkl', allow_pickle=True)
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
# sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
# cloth = sequence.split('_')[0]
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(garment_dir, sequence)))
|
||||
|
||||
for frame in frames:
|
||||
data = np.load(os.path.join(garment_dir, sequence, frame))
|
||||
pose = data['pose'].astype(np.float32)
|
||||
shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release_old/minimal_body_shape', id, 'betas_2.npy'))
|
||||
seq_gender = [0] if gender_table[id] == 'female' else [1]
|
||||
shape = np.concatenate((shape, seq_gender)).astype(np.float32)
|
||||
# body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=True)[0][0]
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
# np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_body_posed=body)
|
||||
|
||||
pose_2 = pose.copy()
|
||||
pose_2[:3] = 0
|
||||
body_2 = smpl.run(smpl.run_rodrigues(pose_2[np.newaxis, :]), shape[np.newaxis,:], remove=True)[0][0]
|
||||
np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_posed_body_no_orient=body_2)
|
||||
|
||||
|
||||
print(id + '_' + frame)
|
||||
|
||||
|
||||
def create_static():
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
src = os.path.join('/home/workspace/local-data/cape', phase, gender)
|
||||
dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_static')
|
||||
src_orig = os.path.join('/home/workspace/local-data/temp/cape_release/sequences')
|
||||
J0 = np.load('custom_data/smpl_root_joint.npy')
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
# sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
# cloth = sequence.split('_')[0]
|
||||
|
||||
avaialble_frames = sorted(os.listdir(os.path.join(src_orig, id, sequence)))
|
||||
first_frame = avaialble_frames[0]
|
||||
|
||||
print(first_frame)
|
||||
data = np.load(os.path.join(src_orig, id, sequence, first_frame))
|
||||
frame = data['v_cano'][smpl.to_keep, :] - J0[np.newaxis, :]
|
||||
|
||||
np.savez_compressed(os.path.join(dst, id + '_' + sequence + '.npz'), mesh_static = frame)
|
||||
|
||||
# x = data['v_posed'][smpl.to_keep, :] - J0[np.newaxis, :] - data['transl'][np.newaxis, :]
|
||||
# faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
# fig = go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# ),
|
||||
# ])
|
||||
|
||||
# fig.update_layout(scene_aspectmode='manual',
|
||||
# scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
# fig.update_layout(
|
||||
# scene = dict(
|
||||
# xaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
# yaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
# zaxis = dict(nticks=4, range=[-1.5,1.5],),)
|
||||
# )
|
||||
|
||||
# fig.show()
|
||||
|
||||
# print()
|
||||
|
||||
|
||||
|
||||
|
||||
def create_unposed():
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
src = os.path.join('/home/workspace/local-data/cape', phase, gender)
|
||||
dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_unposed_no_shape')
|
||||
gender_table = np.load('/home/workspace/data/dataset/cape_release_old/misc/subj_genders.pkl', allow_pickle=True)
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
# faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
J0 = np.load('custom_data/smpl_root_joint.npy')
|
||||
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
# sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
# cloth = sequence.split('_')[0]
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(garment_dir, sequence)))
|
||||
|
||||
for frame in frames:
|
||||
data = np.load(os.path.join(garment_dir, sequence, frame))
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
|
||||
shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release_old', 'minimal_body_shape/' + id + '/betas_2.npy'))
|
||||
seq_gender = 0 if gender_table[id] == 'female' else 1
|
||||
|
||||
shape_offsets = tf.tensordot(shape[tf.newaxis, :], smpl.shapedirs[seq_gender], axes=[[1], [2]]).numpy()
|
||||
unposed = data['v_cano'][smpl.to_keep, :] - J0[np.newaxis, :] - np.squeeze(shape_offsets)[smpl.to_keep, :]
|
||||
|
||||
np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_unposed_no_shape=unposed)
|
||||
|
||||
# x = unposed
|
||||
# faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
# fig = go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# ),
|
||||
# ])
|
||||
|
||||
# fig.update_layout(scene_aspectmode='manual',
|
||||
# scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
# fig.update_layout(
|
||||
# scene = dict(
|
||||
# xaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
# yaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
# zaxis = dict(nticks=4, range=[-1.5,1.5],),)
|
||||
# )
|
||||
|
||||
# fig.show()
|
||||
|
||||
print(id + '_' + frame)
|
||||
|
||||
|
||||
def create_posed():
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
src = os.path.join('/home/workspace/local-data/cape', phase, gender)
|
||||
dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_posed')
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
J0 = np.load('custom_data/smpl_root_joint.npy')
|
||||
|
||||
|
||||
ids = os.listdir(src)
|
||||
|
||||
for id in ids:
|
||||
garment_dir = os.path.join(src, id)
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
# sequence_name = '_'.join(sequence.split('_')[1:])
|
||||
# cloth = sequence.split('_')[0]
|
||||
|
||||
frames = sorted(os.listdir(os.path.join(garment_dir, sequence)))
|
||||
|
||||
for frame in frames:
|
||||
data = np.load(os.path.join(garment_dir, sequence, frame))
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
|
||||
posed = data['v_posed'][smpl.to_keep, :] - J0[np.newaxis, :] - data['transl'][np.newaxis, :]
|
||||
|
||||
np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_posed=posed)
|
||||
|
||||
# x = posed
|
||||
# faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
# fig = go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# ),
|
||||
# ])
|
||||
|
||||
# fig.update_layout(scene_aspectmode='manual',
|
||||
# scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
# fig.update_layout(
|
||||
# scene = dict(
|
||||
# xaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
# yaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
# zaxis = dict(nticks=4, range=[-1.5,1.5],),)
|
||||
# )
|
||||
|
||||
# fig.show()
|
||||
|
||||
print(id + '_' + frame)
|
||||
|
||||
|
||||
def create_mask():
|
||||
pass
|
||||
|
||||
|
||||
def verify():
|
||||
root = os.path.join('/home/workspace/local-data/cape_processed', phase, gender)
|
||||
|
||||
mesh_static = os.listdir(os.path.join(root, 'mesh_static'))
|
||||
|
||||
mesh_posed = sorted(os.listdir(os.path.join(root, 'mesh_posed')))
|
||||
mesh_posed_body_no_orient = sorted(os.listdir(os.path.join(root, 'mesh_posed_body_no_orient')))
|
||||
mesh_unposed_no_shape = sorted(os.listdir(os.path.join(root, 'mesh_unposed_no_shape')))
|
||||
pose_shape = sorted(os.listdir(os.path.join(root, 'pose_shape')))
|
||||
mesh_curvature = sorted(os.listdir(os.path.join(root, 'mesh_curvature')))
|
||||
|
||||
mesh_posed = ['.'.join(x.split('.')[:-1]) for x in mesh_posed]
|
||||
mesh_posed_body_no_orient = ['.'.join(x.split('.')[:-1]) for x in mesh_posed_body_no_orient]
|
||||
mesh_unposed_no_shape = ['.'.join(x.split('.')[:-1]) for x in mesh_unposed_no_shape]
|
||||
pose_shape = ['.'.join(x.split('.')[:-1]) for x in pose_shape]
|
||||
mesh_curvature = ['.'.join(x.split('.')[:-1]) for x in mesh_curvature]
|
||||
|
||||
mesh_static = ['.'.join(x.split('.')[:-1]) for x in mesh_static]
|
||||
mesh_posed = ['.'.join(x.split('.')[:-1]) for x in mesh_posed]
|
||||
mesh_posed_body_no_orient = ['.'.join(x.split('.')[:-1]) for x in mesh_posed_body_no_orient]
|
||||
mesh_unposed_no_shape = ['.'.join(x.split('.')[:-1]) for x in mesh_unposed_no_shape]
|
||||
pose_shape = ['.'.join(x.split('.')[:-1]) for x in pose_shape]
|
||||
mesh_curvature = ['.'.join(x.split('.')[:-1]) for x in mesh_curvature]
|
||||
|
||||
|
||||
assert mesh_posed == mesh_posed_body_no_orient == mesh_unposed_no_shape == pose_shape == mesh_curvature
|
||||
assert set(mesh_posed) == set(mesh_posed_body_no_orient) == set(mesh_unposed_no_shape) == set(pose_shape) == set(mesh_static) == set(mesh_curvature)
|
||||
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
train_split = get_frames()
|
||||
phases = ['test', 'test', 'train', 'train']
|
||||
genders = ['male', 'female', 'male', 'female']
|
||||
|
||||
for phase, gender in zip(phases, genders):
|
||||
# get_frames()
|
||||
# split_data()
|
||||
# verify_split_data()
|
||||
# merge_pose_shape()
|
||||
# create_posed_body()
|
||||
# create_static()
|
||||
# create_unposed()
|
||||
# create_posed()
|
||||
verify()
|
||||
|
||||
print()
|
||||
@@ -0,0 +1,487 @@
|
||||
import os
|
||||
import sys
|
||||
import tqdm
|
||||
import math
|
||||
import time
|
||||
import random
|
||||
import shutil
|
||||
import collections
|
||||
import argparse
|
||||
import threading
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import trimesh as tm
|
||||
import cmapy as cm
|
||||
import scipy.io as sio
|
||||
import scipy.interpolate as spi
|
||||
import matplotlib.pyplot as plt
|
||||
import scipy.stats as stat
|
||||
import cv2
|
||||
import plotly.express as ex
|
||||
import plotly.graph_objects as go
|
||||
|
||||
|
||||
from custom_utils.supersmpl_io import readPC2
|
||||
|
||||
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
|
||||
|
||||
|
||||
def add_triangle_centers():
|
||||
pass
|
||||
|
||||
def dilate(img):
|
||||
uv_mask_2 = img != 0
|
||||
m = img.min()
|
||||
# norm = img.max() - img.min()
|
||||
img = (img - m) * uv_mask_2 #* (np.invert(uv_mask_2) * -1 + uv_mask_2)
|
||||
size = 1
|
||||
shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT
|
||||
element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size))
|
||||
result = cv2.dilate(img, element, iterations=1)
|
||||
result = (result + m) * uv_mask
|
||||
return result
|
||||
|
||||
def erode(img):
|
||||
# m = img.min()
|
||||
# norm = img.max() - img.min()
|
||||
# img = (img - m)
|
||||
size = 1
|
||||
shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT
|
||||
element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size))
|
||||
result = cv2.erode(img, element, iterations=2)
|
||||
# result = result + m
|
||||
return result
|
||||
|
||||
def chunks(lst, n):
|
||||
"""Yield successive n-sized chunks from lst."""
|
||||
for i in range(0, len(lst), n):
|
||||
yield lst[i:i + n]
|
||||
|
||||
|
||||
class Result():
|
||||
def __init__(self) -> None:
|
||||
self.result = dict()
|
||||
|
||||
result_g = Result()
|
||||
|
||||
class RBFThread(threading.Thread):
|
||||
def __init__(self, uv_values, id, result):
|
||||
threading.Thread.__init__(self)
|
||||
self.uv_values = uv_values
|
||||
self.id = id
|
||||
self.result = result
|
||||
|
||||
def run(self):
|
||||
print('processing: ' + str(self.id))
|
||||
rbf = spi.RBFInterpolator(y=uv_map_discreet_unique, d=self.uv_values, neighbors=10)
|
||||
res = rbf(rbf_dest)
|
||||
self.result[self.id] = res
|
||||
pass
|
||||
|
||||
def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None):
|
||||
verts_uv = np.zeros((res * 2, res, 3))
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = verts
|
||||
uv_values[m2uv[:, 1]] += verts
|
||||
uv_values = uv_values[uv_map_mask, :] # add garment mask
|
||||
uv_values /= uv_counter[:, np.newaxis]
|
||||
return_garment_uv_mask = False
|
||||
|
||||
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, axis=-1) != 0
|
||||
|
||||
# test
|
||||
verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
|
||||
# verts_uv_save = verts_uv.copy()
|
||||
|
||||
# garment_mask = np.sum(uv_values, 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
|
||||
|
||||
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, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4)
|
||||
garment_uv_mask_border = garment_uv_mask.astype(bool)
|
||||
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)
|
||||
garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
|
||||
garment_uv_mask = garment_uv_mask * uv_mask[:,:,0]
|
||||
rbf_dest_current = rbf_dest * garment_uv_mask
|
||||
return_garment_uv_mask = True
|
||||
|
||||
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
|
||||
if method == 'rbf':
|
||||
interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear')
|
||||
elif method == 'nearest':
|
||||
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
final = interpolator(rbf_dest_current)
|
||||
|
||||
result = np.zeros((res * 2, res, 3))
|
||||
|
||||
# not sure, shouldn't be needed
|
||||
# rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
|
||||
|
||||
result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
|
||||
result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
|
||||
if garment_uv_mask is not None:
|
||||
result = result * garment_uv_mask[:,:,np.newaxis]
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return result, garment_uv_mask, garment_uv_mask_border
|
||||
else:
|
||||
return result
|
||||
|
||||
|
||||
|
||||
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):
|
||||
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
|
||||
# verts_uv_save = verts_uv.copy()
|
||||
|
||||
# unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0
|
||||
# unique_uv_values = uv_values[:, unique_uv_values_mask, :]
|
||||
# uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :]
|
||||
|
||||
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
|
||||
|
||||
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).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
|
||||
final.append(interpolator(rbf_dest_current))
|
||||
final = np.asarray(final)
|
||||
|
||||
result = np.zeros((batch, res * 2, res, 3))
|
||||
result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
|
||||
# result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
|
||||
result[:, uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values
|
||||
if garment_uv_mask is not None:
|
||||
result = result * garment_uv_mask[np.newaxis,:,:,np.newaxis]
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return result, garment_uv_mask
|
||||
else:
|
||||
return result
|
||||
|
||||
def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
|
||||
batch = values.shape[0]
|
||||
values = (values + 1) * mask
|
||||
values_uv = np.zeros((batch, res * 2, res))
|
||||
missing = np.zeros((res * 2, res))
|
||||
uv_values = np.zeros((batch, uv_map.shape[0]))
|
||||
uv_values[:, m2uv[:, 0]] = values
|
||||
uv_values[:, m2uv[:, 1]] = values
|
||||
uv_values = uv_values[:, uv_map_mask]
|
||||
return_garment_uv_mask = False
|
||||
# test
|
||||
# verts_uv_save = verts_uv.copy()
|
||||
|
||||
# unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0
|
||||
# unique_uv_values = uv_values[:, unique_uv_values_mask, :]
|
||||
# uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :]
|
||||
|
||||
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:
|
||||
raise Exception
|
||||
|
||||
# uv_values = uv_values * garment_mask
|
||||
|
||||
# test
|
||||
values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
|
||||
missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0)
|
||||
|
||||
unique_uv_values = uv_values[:, garment_mask]
|
||||
uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
|
||||
# rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0]
|
||||
|
||||
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
|
||||
|
||||
rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask
|
||||
rbf_dest_current = np.nonzero(rbf_dest_current)
|
||||
rbf_dest_current = np.array(rbf_dest_current).T
|
||||
|
||||
# uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
|
||||
final = []
|
||||
for i in range(batch):
|
||||
# print('processing frame: ' + str(i))
|
||||
interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:])
|
||||
final.append(interpolator(rbf_dest_current))
|
||||
final = np.asarray(final)
|
||||
|
||||
result = np.zeros((batch, res * 2, res))
|
||||
result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
|
||||
result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final
|
||||
if garment_uv_mask is not None:
|
||||
result = result * garment_uv_mask[np.newaxis,:,:]
|
||||
|
||||
if return_garment_uv_mask:
|
||||
return result, garment_uv_mask
|
||||
else:
|
||||
return result
|
||||
|
||||
|
||||
def m2uv_array(m2uv):
|
||||
result = np.zeros((3638, 2), dtype=np.int32)
|
||||
for k, v in m2uv.items():
|
||||
if len(v) == 2:
|
||||
result[k] = list(v)
|
||||
elif len(v) == 1:
|
||||
x = v.pop()
|
||||
result[k] = [x, x]
|
||||
else:
|
||||
assert False
|
||||
|
||||
return result
|
||||
|
||||
|
||||
|
||||
def verify_reconstruction(orig, reconstructed):
|
||||
error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1))
|
||||
colors = np.array([cm.color('jet', 1 - i) for i in error])
|
||||
error_sum = np.sum(error)
|
||||
|
||||
|
||||
verts_uv = np.zeros((uv_res * 2, uv_res, 3))
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = colors
|
||||
uv_values[m2uv[:, 1]] += colors
|
||||
uv_counter = np.zeros((uv_map.shape[0]))
|
||||
uv_counter[m2uv[:, 0]] = 1
|
||||
uv_counter[m2uv[:, 1]] += 1
|
||||
uv_values /= uv_counter[:, np.newaxis]
|
||||
verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(111)
|
||||
ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255)
|
||||
plt.show()
|
||||
|
||||
|
||||
|
||||
fig = plt.figure(figsize=(15, 15))
|
||||
ax = fig.add_subplot(121, projection='3d')
|
||||
x = orig
|
||||
ax.set_xlim3d(-1, 1)
|
||||
ax.set_ylim3d(-1, 1)
|
||||
ax.set_zlim3d(-1, 1)
|
||||
ax.set_xlabel("X")
|
||||
ax.set_ylabel("Y")
|
||||
ax.set_zlabel("Z")
|
||||
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
|
||||
|
||||
ax = fig.add_subplot(122, projection='3d')
|
||||
x = reconstructed#[error > 0.1, :]
|
||||
ax.set_xlim3d(-1, 1)
|
||||
ax.set_ylim3d(-1, 1)
|
||||
ax.set_zlim3d(-1, 1)
|
||||
ax.set_xlabel("X")
|
||||
ax.set_ylabel("Y")
|
||||
ax.set_zlabel("Z")
|
||||
ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :]
|
||||
|
||||
plt.show()
|
||||
|
||||
pass
|
||||
|
||||
# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs
|
||||
def unique_pairs(a):
|
||||
a.sort(axis=1)
|
||||
b = np.ascontiguousarray(a).view(
|
||||
np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
|
||||
)
|
||||
_, inv, ct = np.unique(b, return_inverse=True, return_counts=True)
|
||||
return ct[inv] == 1
|
||||
|
||||
|
||||
def unique_uv_map_mask(uv_map_discreet):
|
||||
mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
||||
temp = np.zeros((uv_res * 2, uv_res))
|
||||
for p in uv_map_discreet:
|
||||
temp[p[0], p[1]] += 1
|
||||
# temp = temp == 1
|
||||
for i, p in enumerate(uv_map_discreet):
|
||||
if temp[p[0], p[1]] == 1:
|
||||
mask[i] = True
|
||||
elif temp[p[0], p[1]] == 2:
|
||||
mask[i] = True
|
||||
temp[p[0], p[1]] = 0
|
||||
return mask
|
||||
|
||||
|
||||
def test(verts, faces, values):
|
||||
uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
uv_values[m2uv[:, 0]] = values
|
||||
uv_values[m2uv[:, 1]] += values
|
||||
uv_counter = np.zeros((uv_map.shape[0]))
|
||||
uv_counter[m2uv[:, 0]] = 1
|
||||
uv_counter[m2uv[:, 1]] += 1
|
||||
uv_values /= uv_counter[:, np.newaxis]
|
||||
x = uv_values
|
||||
verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
|
||||
im = ((x - x.min()) / (x.max() - x.min())) * 255
|
||||
tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show()
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cape_processed/train/male')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=100)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
args = parser.parse_args()
|
||||
|
||||
# smpl = SMPL('FeatureLines/superSMPL')
|
||||
|
||||
|
||||
F = np.load("custom_utils/smpl_faces_orig.npy")
|
||||
folder = args.folder
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_3_orig.npy')
|
||||
m2uv = np.load('custom_utils/m2uv_orig.npy', allow_pickle=True).item()
|
||||
m2uv = m2uv_array(m2uv)
|
||||
uv_mask = np.load('custom_utils/mask_256x512.npy')
|
||||
curvature_label_type = '6_25'
|
||||
|
||||
uv_res = 256
|
||||
uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
||||
# uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
|
||||
# uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
|
||||
uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
||||
|
||||
uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
|
||||
uv_map_mask = unique_uv_map_mask(uv_map_discreet)
|
||||
|
||||
uv_counter = np.zeros((uv_map.shape[0]))
|
||||
uv_counter[m2uv[:, 0]] = 1
|
||||
uv_counter[m2uv[:, 1]] += 1
|
||||
uv_counter = uv_counter[uv_map_mask]
|
||||
|
||||
uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :]
|
||||
|
||||
rbf_dest = np.ones((uv_res * 2, uv_res))
|
||||
rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
|
||||
rbf_dest[np.invert(uv_mask[:,:,0])] = 0
|
||||
# rbf_dest = np.nonzero(rbf_dest)
|
||||
# rbf_dest = np.array(rbf_dest).T
|
||||
# rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
|
||||
# rbf_dest = rbf_dest[unique_pairs(rbf_dest), :] # not good, unique_pairs doesn't consider coord order
|
||||
|
||||
|
||||
# os.makedirs(os.path.join(folder, 'uv_posed_body_no_orient'), exist_ok=True)
|
||||
# os.makedirs(os.path.join(folder, 'uv_posed'), exist_ok=True)
|
||||
# os.makedirs(os.path.join(folder, 'uv_curvature'), exist_ok=True)
|
||||
# os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
|
||||
# # os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
|
||||
# os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, 'pose_shape'))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
frame_name = frame.split('.')
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
static_frame_name = frame_name.split('.')[0]
|
||||
|
||||
verts_static = np.load(os.path.join(folder, 'mesh_static', static_frame_name + '.npz')); verts_static = verts_static['mesh_static']
|
||||
mask = np.ones((3638), bool)
|
||||
body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', frame_name + '.npz')); body_no_orient = body_no_orient['mesh_posed_body_no_orient']
|
||||
gt = np.load(os.path.join(folder, 'mesh_posed', frame_name + '.npz')); gt = gt['mesh_posed']
|
||||
curvature = np.load(os.path.join(folder, 'mesh_curvature', frame_name + '.npz')); curvature = curvature['mesh_curvature']
|
||||
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', frame_name + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
|
||||
body_no_orient, gt, curvature, mesh_unposed_no_shape = body_no_orient[np.newaxis,:,:], gt[np.newaxis,:,:], curvature[np.newaxis,:], mesh_unposed_no_shape[np.newaxis,:,:]
|
||||
|
||||
# start = time.perf_counter()
|
||||
|
||||
|
||||
uv_verts_static, garment_uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, verts_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
|
||||
uv_verts_static = uv_verts_static.astype(np.float32)
|
||||
# uv_body_static = convert_to_uv(uv_res, m2uv, uv_map, body_static, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
|
||||
uv_posed_body_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, body_no_orient[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_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, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
# print(time.perf_counter() - start)
|
||||
|
||||
uv_posed_body_no_orient, uv_gt, uv_curvature, uv_unposed_no_shape = uv_posed_body_no_orient[0], uv_gt[0], uv_curvature[0], uv_unposed_no_shape[0]
|
||||
|
||||
# start = time.perf_counter()
|
||||
if not os.path.exists(os.path.join(folder, 'uv_static', static_frame_name + '.npz')):
|
||||
np.savez_compressed(os.path.join(folder, 'uv_static', static_frame_name + '.npz'), uv_static=uv_verts_static)
|
||||
# np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static)
|
||||
# np.savez_compressed(os.path.join(folder, 'uv_mask', static_frame_name + '.npz'), uv_mask=garment_uv_mask)
|
||||
|
||||
np.savez_compressed(os.path.join(folder, 'uv_posed_body_no_orient', frame_name + '.npz'), uv_posed_body_no_orient=uv_posed_body_no_orient)
|
||||
np.savez_compressed(os.path.join(folder, 'uv_posed', frame_name + '.npz'), uv_gt=uv_gt)
|
||||
np.savez_compressed(os.path.join(folder, 'uv_curvature', frame_name + '.npz'), curvature_uv=uv_curvature)
|
||||
np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape', frame_name + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape)
|
||||
|
||||
|
||||
# print(time.perf_counter() - start)
|
||||
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
import os
|
||||
import sys
|
||||
import tqdm
|
||||
import argparse
|
||||
import random
|
||||
import string
|
||||
import numpy as np
|
||||
import trimesh as tm
|
||||
import cmapy as cm
|
||||
import scipy as sp
|
||||
import scipy.io as sio
|
||||
import plotly.graph_objects as go
|
||||
from pyntcloud import PyntCloud
|
||||
from scipy.sparse import csr_matrix, coo_matrix
|
||||
from scipy.sparse.csgraph import floyd_warshall, dijkstra
|
||||
|
||||
from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
# from FeatureLines.superSMPL.smpl import SMPL
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
import tensorflow as tf
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
|
||||
except RuntimeError as e:
|
||||
print(e)
|
||||
from custom_utils.smpl import SMPL
|
||||
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
|
||||
|
||||
def get_skirt_body():
|
||||
body = readOBJ('FeatureLines/superSMPL_skirt.obj')
|
||||
|
||||
return body[0]
|
||||
|
||||
|
||||
def export_to_ply(mesh, file):
|
||||
ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
|
||||
output_file = open(file, "wb+")
|
||||
output_file.write(ply)
|
||||
output_file.close()
|
||||
|
||||
|
||||
def get_closest_n_points(verts, idx, n):
|
||||
dist = sp.spatial.distance.cdist(verts[idx, :], verts)
|
||||
dist[dist == 0] = 1e9
|
||||
idx_sort = np.argsort(dist, axis=-1)
|
||||
closest = idx_sort[:, :n]
|
||||
return closest
|
||||
|
||||
|
||||
def histedges_equalN(x, nbin):
|
||||
npt = len(x)
|
||||
return np.interp(np.linspace(0, npt, nbin + 1),
|
||||
np.arange(npt),
|
||||
np.sort(x))
|
||||
|
||||
|
||||
|
||||
def k_neighbours(k):
|
||||
body = get_skirt_body()
|
||||
adj = np.ones((14475, 14475)) * 9999 # ???
|
||||
diagonal = np.linspace(0, 14474, 14475, dtype=np.int32)
|
||||
m = tm.Trimesh(body, F, process=False, validate=False)
|
||||
edges = m.edges_unique
|
||||
lengths = m.edges_unique_length
|
||||
adj[edges[:, 0], edges[:, 1]] = lengths
|
||||
adj[edges[:, 1], edges[:, 0]] = lengths
|
||||
adj[diagonal, diagonal] = 0
|
||||
|
||||
adj = csr_matrix(adj)
|
||||
dist = dijkstra(adj, directed=False, unweighted=False, limit=0.5)
|
||||
|
||||
np.save('MeshUtils/dijkstra_skirt.npy', dist)
|
||||
|
||||
# dist = floydWarshall(adj, 14475)
|
||||
|
||||
# idx_sort = np.argsort(dist, axis=-1)
|
||||
# closest = idx_sort[:, :k]
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def generate_curvatures(folder, neighbourhood, save_folder):
|
||||
ids = os.listdir(folder)[int(args.start):int(args.end)]
|
||||
# gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True)
|
||||
faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
k_neighbors = np.load('custom_data/dijkstra_orig_argsort_25.npy')
|
||||
|
||||
os.makedirs(os.path.join(save_folder, 'curvatures_raw'), exist_ok=True)
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
frames = os.listdir(os.path.join(garment_dir, sequence))
|
||||
|
||||
for frame in frames:
|
||||
data = np.load(os.path.join(folder, id, sequence, frame))
|
||||
# pose = data['poses']
|
||||
# shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release/minimal_body_shape', id, id + '_minimal.npy'))
|
||||
# gender = [0] if gender_table[id] == 'female' else [1]
|
||||
# shape = np.concatenate((shape, gender)).astype(np.float32)
|
||||
# body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=False)[0][0]
|
||||
|
||||
verts = data['v_posed'][smpl.to_keep, :]
|
||||
|
||||
m = tm.Trimesh(verts, faces, process=False, validate=False)
|
||||
export_to_ply(m, temp_file)
|
||||
|
||||
cloud = PyntCloud.from_file(temp_file)
|
||||
# cloud.plot(mesh=True)
|
||||
|
||||
ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
|
||||
cloud.add_scalar_field("curvature", ev=ev)
|
||||
curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
|
||||
np.save(os.path.join(save_folder, 'curvatures_raw', id + '_' + frame_name + '.npy'), np.array(curvature))
|
||||
|
||||
print(id + '_' + frame)
|
||||
|
||||
|
||||
def get_bins(values, nr_of_bins):
|
||||
flattened = []
|
||||
for c in values:
|
||||
flattened.extend(c.flatten())
|
||||
flattened = np.array(flattened)
|
||||
bins = histedges_equalN(flattened, nr_of_bins)
|
||||
return bins
|
||||
|
||||
|
||||
def generate_curvature_labels(folder, bins):
|
||||
os.makedirs(os.path.join(save_folder, 'mesh_curvature'), exist_ok=True)
|
||||
|
||||
frames = os.listdir(folder)#[int(args.start):int(args.end)]
|
||||
labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
|
||||
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
curvature = np.load(os.path.join(folder, frame))
|
||||
frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
|
||||
for j in range(bins.shape[0] - 2):
|
||||
frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
|
||||
frame_labels[curvature >= bins[-2]] = labels[-1]
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npy', 'It\'s not an npy file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
np.savez_compressed(os.path.join(save_folder, 'mesh_curvature', frame_name + '.npz'), mesh_curvature=frame_labels)
|
||||
|
||||
# ids = os.listdir(folder)#[:10]
|
||||
# labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
|
||||
|
||||
# i = 0
|
||||
# for id in tqdm.tqdm(ids):
|
||||
# garment_dir = os.path.join(folder, id)
|
||||
|
||||
|
||||
# curvature = curvatures[i]
|
||||
# i += 1
|
||||
|
||||
# curvature_label = np.zeros((curvature.shape[0], curvature.shape[1]))
|
||||
# frame = 0
|
||||
# for frame in range(curvature.shape[0]):
|
||||
# frame_labels = np.zeros((curvature.shape[1]), dtype=np.float32)
|
||||
# for j in range(bins.shape[0] - 2):
|
||||
# frame_labels[(curvature[frame, :] >= bins[j]) * (curvature[frame, :] < bins[j+1])] = labels[j]
|
||||
# frame_labels[curvature[frame, :] >= bins[-2]] = labels[-1]
|
||||
# curvature_label[frame, :] = frame_labels
|
||||
# frame += 1
|
||||
|
||||
# unmasked_label = np.zeros((curvature.shape[0], 14475), dtype=np.float32)
|
||||
# unmasked_label[:, mask] = curvature_label
|
||||
# np.save(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".npy"), unmasked_label)
|
||||
# writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)
|
||||
|
||||
|
||||
def visualize_curvature_labels(folder):
|
||||
frames = os.listdir(os.path.join(folder, 'mesh_curvature'))#[:10]
|
||||
faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
|
||||
for _ in range(10):
|
||||
|
||||
frame = random.choice(frames)
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', frame)); mesh_posed = mesh_posed['mesh_posed']
|
||||
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', frame)); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
|
||||
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(mesh_curvature.max()+1))])
|
||||
colors_binned = np.zeros((mesh_curvature.shape[0], 3))
|
||||
for i in range(int(mesh_curvature.max())+1):
|
||||
print(np.count_nonzero(mesh_curvature == i))
|
||||
colors_binned[mesh_curvature == i, :] = colors[i]
|
||||
|
||||
x = mesh_posed
|
||||
meshes = [go.Mesh3d(
|
||||
x=x[:,0],
|
||||
y=x[:,2],
|
||||
z=x[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
vertexcolor=colors_binned
|
||||
),]
|
||||
|
||||
fig = go.Figure(data=meshes)
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
fig.update_layout(
|
||||
scene = dict(
|
||||
xaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
yaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
zaxis = dict(nticks=4, range=[-1.5,1.5],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
print()
|
||||
|
||||
|
||||
def read_raw_curvatures(folder):
|
||||
|
||||
frames = os.listdir(folder)#[int(args.start):int(args.end)]
|
||||
|
||||
curvatures_all = []
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
curv = np.load(os.path.join(folder, frame))
|
||||
curvatures_all.append(curv)
|
||||
|
||||
return curvatures_all
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-if", "--in_folder", help = "process the given folder", default=r'/home/workspace/local-data/cape_processed/test/male/curvatures_raw')
|
||||
parser.add_argument("-of", "--out_folder", help = "process the given folder", default=r'/home/workspace/local-data/cape_processed/test/male')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=1500)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
|
||||
args = parser.parse_args()
|
||||
|
||||
F = np.load("custom_utils/smpl_faces_orig.npy")
|
||||
# k_neighbours(k=25)
|
||||
|
||||
folder = args.in_folder
|
||||
save_folder = args.out_folder #'/home/workspace/data/dataset/cape_processed'
|
||||
nr_of_labels = 6
|
||||
neighbourhood = 25
|
||||
temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
|
||||
# dist = np.load('custom_data/dijkstra.npy')
|
||||
# dist = np.load('MeshUtils/dijkstra_skirt.npy')
|
||||
|
||||
# generate_curvatures(folder, neighbourhood, save_folder)
|
||||
# curvatures = read_raw_curvatures(folder)
|
||||
# bins = get_bins(curvatures, nr_of_labels)
|
||||
# np.save('custom_data/cape_2_train_male_bins.npy', bins)
|
||||
bins = np.load('custom_data/cape_2_train_male_bins.npy')
|
||||
generate_curvature_labels(folder, bins)
|
||||
# visualize_curvature_labels(save_folder)
|
||||
|
||||
pass
|
||||
+40
-15
@@ -54,7 +54,7 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
|
||||
|
||||
|
||||
class StaticDataset():
|
||||
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True) -> None:
|
||||
def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None:
|
||||
self.store = dict()
|
||||
self.signature = signature
|
||||
self.folder = folder
|
||||
@@ -68,9 +68,20 @@ class StaticDataset():
|
||||
self.keys.append('uv_static_128')
|
||||
self.resolutions = resolutions
|
||||
self.filled = ''
|
||||
self.is_cape = is_cape
|
||||
if use_filled:
|
||||
self.filled = '_filled'
|
||||
|
||||
if is_cape:
|
||||
self.uv_masks = dict()
|
||||
uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
|
||||
self.uv_masks[512] = uv_mask.astype(np.float32)
|
||||
for res in [32, 64, 128, 256]:
|
||||
down = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(np.float32))
|
||||
self.uv_masks[res] = down
|
||||
|
||||
|
||||
|
||||
def read(self, id):
|
||||
data = dict()
|
||||
@@ -83,12 +94,18 @@ class StaticDataset():
|
||||
if 'uv_static' + res_string in self.signature.keys():
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
|
||||
if 'uv_mask' + res_string in self.signature.keys():
|
||||
if self.is_cape:
|
||||
data['uv_mask' + res_string] = self.uv_masks[res]
|
||||
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
|
||||
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
|
||||
@@ -118,9 +135,13 @@ class StaticDataset():
|
||||
|
||||
class DatasetIdx(tf.data.Dataset):
|
||||
|
||||
def __new__(cls, folder, shuffle=True):
|
||||
def __new__(cls, folder, shuffle=True, is_cape=False):
|
||||
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 shuffle:
|
||||
random.shuffle(frames)
|
||||
@@ -146,18 +167,22 @@ class DatasetFactory():
|
||||
return tf_signature
|
||||
|
||||
|
||||
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True):
|
||||
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True, is_cape=False):
|
||||
# assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
|
||||
# static_data = StaticDataset(folder, signature)
|
||||
dataset_idx = DatasetIdx(folder, shuffle)
|
||||
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape)
|
||||
filled = ''
|
||||
if use_filled:
|
||||
filled = '_filled'
|
||||
separator = '_'
|
||||
if is_cape:
|
||||
separator = '.'
|
||||
|
||||
|
||||
def read_item(frame):
|
||||
frame_split = tf.strings.split(frame, sep='_')
|
||||
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]
|
||||
@@ -166,7 +191,7 @@ class DatasetFactory():
|
||||
if 'id' in signature.keys():
|
||||
item['id'] = id
|
||||
if 'pose_shape' in signature.keys():
|
||||
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy'))
|
||||
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy'))
|
||||
|
||||
##############################
|
||||
pose_shape[:3] = 0
|
||||
@@ -174,7 +199,7 @@ class DatasetFactory():
|
||||
|
||||
item['pose_shape'] = pose_shape
|
||||
if 'uv_posed' in signature.keys():
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
item['uv_posed'] = uv_posed
|
||||
# if 'uv_body_posed' in signature.keys():
|
||||
# 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])
|
||||
@@ -187,10 +212,10 @@ class DatasetFactory():
|
||||
# item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient
|
||||
|
||||
if 'mesh_body_posed' in signature.keys():
|
||||
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
||||
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
||||
item['mesh_body_posed'] = mesh_posed_body_no_orient
|
||||
if 'mesh_posed' in signature.keys():
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed']
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = 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']
|
||||
@@ -199,12 +224,12 @@ class DatasetFactory():
|
||||
# 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():
|
||||
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + 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
|
||||
|
||||
|
||||
if 'mesh_curvature' in signature.keys():
|
||||
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
|
||||
|
||||
|
||||
@@ -214,7 +239,7 @@ class DatasetFactory():
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_curvature' + res_string in signature.keys():
|
||||
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled + res_string, id + '_' + frame_nr + '.npz'));
|
||||
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled + res_string, id + separator + frame_nr + '.npz'));
|
||||
if filled != '' or res_string != '':
|
||||
uv_curvature = uv_curvature['uv_curvature' + filled]
|
||||
else:
|
||||
@@ -223,11 +248,11 @@ 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' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in signature.keys():
|
||||
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
|
||||
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + separator + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
import argparse
|
||||
import os
|
||||
import tqdm
|
||||
import cv2
|
||||
import tensorflow as tf
|
||||
import scipy.interpolate as spi
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
# Currently, memory growth needs to be the same across GPUs
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
|
||||
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
|
||||
except RuntimeError as e:
|
||||
# Memory growth must be set before GPUs have been initialized
|
||||
print(e)
|
||||
|
||||
def get_img_idxs():
|
||||
img_idxs = dict()
|
||||
img_coords = dict()
|
||||
uv_masks = dict()
|
||||
uv_mask = np.load('custom_utils/mask_256x512.npy')
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - uv_mask))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((uv_mask))), (1,0))
|
||||
img_idxs[512] = img_idx
|
||||
img_coords[512] = coords
|
||||
uv_masks[512] = uv_mask.astype(np.int32)
|
||||
for res in [32, 64, 128, 256]:
|
||||
down = tf.image.resize(uv_mask[:,:,tf.newaxis].astype(np.int32), size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(bool))
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((down))), (1,0))
|
||||
img_idxs[res] = img_idx
|
||||
img_coords[res] = coords
|
||||
uv_masks[res] = down.astype(np.int32)
|
||||
|
||||
return img_idxs, img_coords, uv_masks
|
||||
|
||||
|
||||
def fill(img, mask):
|
||||
coords = img_coords[img.shape[0]]
|
||||
img_idx = img_idxs[img.shape[0]]
|
||||
# coords = np.transpose(np.stack(np.nonzero((1 - uv_masks[img.shape[0]]))), (1,0))
|
||||
# img_idx = np.transpose(np.stack(np.nonzero((uv_masks[img.shape[0]]))), (1,0))
|
||||
img_data = img[img_idx[:,0], img_idx[:,1], :]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256,3), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1], :] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1], :] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def fill_label(img, mask):
|
||||
coords = img_coords[img.shape[0]]
|
||||
img_idx = img_idxs[img.shape[0]]
|
||||
img_data = img[img_idx[:,0], img_idx[:,1]]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1]] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1]] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def generate_downscaled_uv_map(folder, id, size):
|
||||
img = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); img = img['uv_mask'].astype(np.float32)
|
||||
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(bool))
|
||||
np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz'), uv_mask = down)
|
||||
|
||||
|
||||
class UVMaskStore():
|
||||
def __init__(self, folder, size):
|
||||
self.size = size
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
return uv_masks[self.size]
|
||||
|
||||
|
||||
def downscale_uv_maps_posed_body(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('_')
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
|
||||
mask_full = uv_mask_store[512][id]
|
||||
img = fill(img, mask_full)
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][id]
|
||||
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
down = down * mask[:,:,np.newaxis]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_curv(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('_')
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
|
||||
mask_full = uv_mask_store[512][id]
|
||||
img = fill_label(img, mask_full)
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][id]
|
||||
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down) * mask[:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_curvature = down)
|
||||
|
||||
|
||||
|
||||
def downscale_uv_maps_unposed(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('_')
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
|
||||
mask_full = uv_mask_store[512][id]
|
||||
img = fill(img, mask_full)
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][id]
|
||||
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
down = down * mask[:,:,np.newaxis]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_static(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('.')
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
|
||||
mask_full = uv_mask_store[512][id]
|
||||
img = fill(img, mask_full)
|
||||
|
||||
for size in sizes:
|
||||
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
mask = uv_mask_store[size][id]
|
||||
down = down * mask[:,:,np.newaxis]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_mask(folder, type, size):
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id, _ = frame.split('.')
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
|
||||
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(bool))
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/cape_processed/test/male')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=10000)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
sizes = [32, 64, 128, 256]
|
||||
folder = args.folder
|
||||
|
||||
img_idxs, img_coords, uv_masks = get_img_idxs()
|
||||
|
||||
uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
|
||||
uv_mask_store[512] = UVMaskStore(folder, size=512)
|
||||
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
|
||||
downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
|
||||
|
||||
# downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
|
||||
# downscale_uv_maps_curv(folder=folder, type='uv_curvature', sizes=sizes)
|
||||
# downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
|
||||
|
||||
|
||||
print()
|
||||
Reference in New Issue
Block a user