started writing preprocessing
This commit is contained in:
@@ -0,0 +1,112 @@
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pickle as pkl
|
||||
from scipy import sparse
|
||||
|
||||
from time import time
|
||||
|
||||
from data.supersmpl_io import writeOBJ, readOBJ
|
||||
|
||||
|
||||
def makeGraph(F):
|
||||
G = {}
|
||||
for f in F:
|
||||
N = len(f)
|
||||
for i,v in enumerate(f):
|
||||
if not v in G: G[v] = {}
|
||||
G[v][f[(i+1) % N]] = 1
|
||||
return G
|
||||
|
||||
def extrapolateMissings(M, G):
|
||||
for i in range(6890, M.shape[0]):
|
||||
if np.isnan(M[i]).all():
|
||||
M[i] = neighAvg(M, G, i)
|
||||
elif np.isnan(M[i]).any():
|
||||
print("HOW?")
|
||||
print(M[i])
|
||||
if np.isnan(M).any(): extrapolateMissings(M, G)
|
||||
|
||||
def neighAvg(M, G, i):
|
||||
neigh = set(G[i].keys())
|
||||
avg = 0
|
||||
count = 0
|
||||
for n in neigh:
|
||||
if not np.isnan(M[n]).any():
|
||||
avg += M[n]
|
||||
count += 1
|
||||
if count: return avg / count
|
||||
return np.nan
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
g = 'f'
|
||||
# Read Super topology
|
||||
V, F, Vt, Ft = readOBJ('super_f_avg.obj')
|
||||
N = V.shape[0]
|
||||
G = makeGraph(F)
|
||||
|
||||
# Read PKL
|
||||
print("Loading .pkl file")
|
||||
with open('custom_data/model_' + g + '.pkl', 'rb') as f:
|
||||
M = pkl.load(f, encoding='latin1')
|
||||
|
||||
# Fields
|
||||
# print("Checking fields and shapes")
|
||||
# for k in M:
|
||||
# print(k, M[k].shape)
|
||||
# Faces
|
||||
print("New faces...")
|
||||
M['f'] = F
|
||||
# J_regressor (24, 6890)
|
||||
# Sparse matrix, just change shape to add zeros
|
||||
print("New regressor")
|
||||
A = np.zeros((24, N), np.float32)
|
||||
A[:,:6890] = M['J_regressor'].todense()
|
||||
M['J_regressor'] = sparse.csr_matrix(A)
|
||||
|
||||
# Weights (6890, 24)
|
||||
print("New weights")
|
||||
start = time()
|
||||
A = np.empty((N, 24), np.float32)
|
||||
A[:] = np.nan
|
||||
A[:6890] = M['weights'][:6890]
|
||||
extrapolateMissings(A, G)
|
||||
M['weights'] = A
|
||||
print(time() - start)
|
||||
|
||||
# Posedirs (6890, 3, 207)
|
||||
print("New posedirs")
|
||||
start = time()
|
||||
A = np.empty((N, 3, 207), np.float32)
|
||||
A[:] = np.nan
|
||||
A[:6890] = M['posedirs'][:6890]
|
||||
extrapolateMissings(A, G)
|
||||
M['posedirs'] = A
|
||||
print(time() - start)
|
||||
|
||||
# v_tempalte (6890, 3)
|
||||
print("New template")
|
||||
start = time()
|
||||
A = np.empty((N, 3), np.float32)
|
||||
A[:] = np.nan
|
||||
A[:6890] = M['v_template'][:6890]
|
||||
extrapolateMissings(A, G)
|
||||
M['v_template'] = A
|
||||
print(time() - start)
|
||||
|
||||
# Shapedirs (6890, 3, 10)
|
||||
print("New shapedirs")
|
||||
start = time()
|
||||
A = np.empty((N, 3, 10), np.float32)
|
||||
A[:] = np.nan
|
||||
A[:6890] = M['shapedirs'][:6890]
|
||||
extrapolateMissings(A, G)
|
||||
M['shapedirs'] = A
|
||||
print(time() - start)
|
||||
|
||||
# Dump
|
||||
print("Saving new model")
|
||||
with open('super_model_' + g + '.pkl', 'wb') as f:
|
||||
pkl.dump(M, f)
|
||||
@@ -0,0 +1,310 @@
|
||||
import glob
|
||||
import os
|
||||
|
||||
from data.supersmpl_io import readPC2
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
|
||||
|
||||
import sys
|
||||
import shutil
|
||||
import tqdm
|
||||
import random
|
||||
import math
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
import scipy.io as sio
|
||||
|
||||
from custom_utils.smpl import SMPL
|
||||
import plotly.graph_objects as go
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
|
||||
|
||||
def create_no_orient_body(folder):
|
||||
ids = os.listdir(folder)
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
|
||||
smpl = SMPL()
|
||||
for id in tqdm.tqdm(ids):
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
# if os.path.exists(os.path.join(folder, id, 'mesh_posed_body_no_orient')):
|
||||
# continue
|
||||
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
|
||||
pose, shape, gender = info['poses'], info['shape'], info['gender']
|
||||
|
||||
# if gender == 1:
|
||||
# continue
|
||||
|
||||
pose = np.transpose(pose, [1, 0]).astype(np.float32)
|
||||
|
||||
if len(os.listdir(os.path.join(folder, id, 'mesh_posed_body_no_orient'))) == pose.shape[0]:
|
||||
continue
|
||||
|
||||
pose[:, :3] = 0
|
||||
beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
|
||||
R = smpl.run_rodrigues(pose)
|
||||
T, _, _ = smpl.run(R, np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), True)
|
||||
T = T.numpy()
|
||||
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'mesh_posed_body_no_orient'), exist_ok=True)
|
||||
for i in range(pose.shape[0]):
|
||||
np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :])
|
||||
|
||||
|
||||
|
||||
def create_no_shape_unposed(folder):
|
||||
import tensorflow as tf
|
||||
ids = os.listdir(folder)
|
||||
|
||||
smpl = SMPL()
|
||||
for id in tqdm.tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
# if os.path.exists(os.path.join(folder, id, 'mesh_unposed_no_shape')):
|
||||
# continue
|
||||
|
||||
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
|
||||
# if len(os.listdir(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot'))) == info['poses'].shape[1]:
|
||||
# continue
|
||||
|
||||
shape = info['shape']
|
||||
gender = int(info['gender'][0,0] > 0)
|
||||
garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
|
||||
assert len (garments) == 1
|
||||
assert gender == info['gender'][0,0]
|
||||
# if len (garments) != 1:
|
||||
# continue
|
||||
garment = garments[0]
|
||||
mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
|
||||
shape_offsets = tf.tensordot(shape.astype(np.float32), smpl.shapedirs[gender], axes=[[1], [2]]).numpy()
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot'), exist_ok=True)
|
||||
for i in range(info['poses'].shape[1]):
|
||||
mesh_unposed = np.load(os.path.join(garment_dir, 'mesh_unposed_one_hot', id + '_' + str(i) + '.npz'))
|
||||
mesh_unposed = mesh_unposed['mesh_unposed']
|
||||
mesh_unposed[mask, :] = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :]
|
||||
np.savez_compressed(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz'), mesh_unposed_no_shape = mesh_unposed)
|
||||
|
||||
|
||||
|
||||
|
||||
def check_if_skirt(folder):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
for id in tqdm.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 f != 'info.mat' and not f.endswith('.obj')]
|
||||
# assert len (garments) == 1
|
||||
# garment = garments[0]
|
||||
|
||||
# assert garment != 'Skirt'
|
||||
|
||||
assert "Skirt" not in garments
|
||||
|
||||
def gather_ids(folder, file_name):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
to_print = []
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
|
||||
if len(id) == 5:
|
||||
to_print.append(id)
|
||||
|
||||
to_print.sort()
|
||||
|
||||
with open(file_name, 'w') as f:
|
||||
f.write('\n'.join(to_print))
|
||||
|
||||
|
||||
def verify_poses(folder):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
|
||||
assert len (garments) == 1
|
||||
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
poses = info['poses']
|
||||
verts = readPC2(os.path.join(folder, id, garments[0] + '_registered.pc16'), True)['V']
|
||||
# print(id)
|
||||
# assert type(poses) == np.ndarray and poses.shape[0] == 72 and poses.shape[1] == verts.shape[0], 'mv ' + os.path.join(folder, id) + ' /home/workspace/data/dataset/invalid/'
|
||||
if not (type(poses) == np.ndarray and poses.shape[0] == 72 and poses.shape[1] == verts.shape[0]):
|
||||
shutil.move(os.path.join(folder, id), '/home/workspace/data/dataset/invalid/')
|
||||
print(id)
|
||||
|
||||
|
||||
def fix_split_garment_generated_data(folder):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
|
||||
assert len (garments) == 1
|
||||
|
||||
if len(id) == 6 and id.endswith('2'):
|
||||
frames_folder = os.path.join(folder, id, 'mesh_posed_body_no_orient') # os.path.join(folder, id, 'mesh_unposed_no_shape')
|
||||
frames = os.listdir(frames_folder)
|
||||
frame_id = frames[0].split('_')[0]
|
||||
if len(frame_id) == 6:
|
||||
assert frame_id == id, 'Wrong id found!'
|
||||
if len(frame_id) == 5 and id.startswith(frame_id):
|
||||
for frame in frames:
|
||||
frame_nr = frame.split('.')[0].split('_')[1]
|
||||
shutil.move(os.path.join(frames_folder, frame), os.path.join(frames_folder, id + '_' + frame_nr + '.npz'))
|
||||
|
||||
|
||||
def create_garment_type_table():
|
||||
folder = '/home/workspace/data/dataset'
|
||||
# sub_folders = ['train', 'validation', 'test', 'train_skirt', 'valid_skirt', 'test_skirt']
|
||||
# sub_folders = ['train', 'validation', 'train_skirt', 'valid_skirt']
|
||||
sub_folders = ['test_skirt']
|
||||
# sub_folders = ['test', 'test_skirt']
|
||||
|
||||
garment_type_table = dict()
|
||||
|
||||
for sub_f in sub_folders:
|
||||
ids = os.listdir(os.path.join(folder, sub_f))
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
if not os.path.isdir(os.path.join(folder, sub_f, id)):
|
||||
continue
|
||||
|
||||
garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, sub_f, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
|
||||
assert len (garments) == 1
|
||||
|
||||
if id in garment_type_table.keys():
|
||||
print(id + ': ' + garments[0] + ' ' + garment_type_table[id])
|
||||
|
||||
garment_type_table[id] = garments[0]
|
||||
|
||||
np.save('custom_data/garment_type_table_test.npy', garment_type_table)
|
||||
|
||||
|
||||
def generate_faces():
|
||||
folder = '/home/workspace/data/dataset/test_skirt_ready'
|
||||
output_folder = '/home/workspace/data/results/support/faces_skirt'
|
||||
F = np.load("custom_utils/smpl_faces_skirt.npy")
|
||||
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
|
||||
ids = os.listdir(os.path.join(folder, 'pose_shape'))
|
||||
ids = [id.split('_')[0] for id in ids]
|
||||
ids = set(ids)
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool)
|
||||
|
||||
_, faces = _maskFaces(np.zeros((14475, 3), np.float32), F, mask)
|
||||
np.save(os.path.join(output_folder, id + '.npy'), faces.astype(np.int32))
|
||||
|
||||
|
||||
def generate_frames_file(folder, ratio=0.1):
|
||||
all_frames = os.listdir(os.path.join(folder, 'pose_shape'))
|
||||
ids = [frame.split('_')[0] for frame in all_frames]
|
||||
all_frames = [frame.split('.')[0] for frame in all_frames]
|
||||
ids = sorted(list(set(ids)))
|
||||
|
||||
garment_type_table = np.load('custom_data/garment_type_table_train.npy', allow_pickle=True).item()
|
||||
|
||||
|
||||
random.shuffle(ids)
|
||||
|
||||
sub_ids = ids[:int(len(ids) * ratio)]
|
||||
|
||||
selected = []
|
||||
|
||||
for id in tqdm.tqdm(sub_ids):
|
||||
if garment_type_table[id] != 'Dress':
|
||||
continue
|
||||
if len(selected) > 30000:
|
||||
break
|
||||
|
||||
# frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
|
||||
frames = [frame for frame in all_frames if frame.startswith(id + '_')]
|
||||
frames = [frame.split('/')[-1].split('.')[0] for frame in frames]
|
||||
selected.extend(frames)
|
||||
|
||||
with open(os.path.join(folder, 'frames_dress_30k.txt'), 'w') as f:
|
||||
for i, frame in enumerate(selected):
|
||||
f.write(frame + '\n')
|
||||
|
||||
|
||||
def generate_frames_file_filtered_cape():
|
||||
folder = '/home/workspace/local-data/cape_processed/train/female'
|
||||
|
||||
all_frames = os.listdir(os.path.join(folder, 'pose_shape'))
|
||||
ids = [frame.split('.')[0] for frame in all_frames]
|
||||
all_frames = ['.'.join(frame.split('.')[:-1]) for frame in all_frames]
|
||||
ids = sorted(list(set(ids)))
|
||||
|
||||
seqs_train_00032_SS = ['00032_shortshort_hips', '00032_shortshort_shoulders_mill']
|
||||
seqs_test_00032_SS = [ '00032_shortshort_tilt_twist_left']
|
||||
|
||||
seqs_train_00032_SL = ['00032_shortlong_hips', '00032_shortlong_shoulders_mill']
|
||||
seqs_test_00032_SL = [ '00032_shortlong_tilt_twist_left']
|
||||
|
||||
seqs_train_00032_SL_2 = ['00032_shortlong_hips', '00032_shortlong_tilt_twist_left']
|
||||
seqs_test_00032_SL_2 = ['00032_shortlong_shoulders_mill']
|
||||
|
||||
seqs_train_03223_SL_2 = ['03223_shortlong_hips', '03223_shortlong_tilt_twist_left']
|
||||
seqs_test_03223_SL_2 = ['03223_shortlong_shoulders_mill']
|
||||
|
||||
seqs_train_03223_SS_2 = ['03223_shortshort_hips', '03223_shortshort_tilt_twist_left']
|
||||
seqs_test_03223_SS_2 = ['03223_shortshort_shoulders_mill']
|
||||
|
||||
ids = [id for id in ids if id in seqs_train_03223_SS_2]
|
||||
|
||||
random.shuffle(ids)
|
||||
|
||||
sub_ids = ids
|
||||
|
||||
selected = []
|
||||
|
||||
for id in tqdm.tqdm(sub_ids):
|
||||
# frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
|
||||
frames = [frame for frame in all_frames if frame.startswith(id + '.')]
|
||||
# frames = [frame.split('/')[-1].split('.')[0] for frame in frames]
|
||||
selected.extend(frames)
|
||||
|
||||
random.shuffle(selected)
|
||||
|
||||
with open(os.path.join(folder, 'frames_03223_SS_train_2.txt'), 'w') as f:
|
||||
for i, frame in enumerate(selected):
|
||||
f.write(frame + '\n')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
src_folder = '/home/workspace/local-data/train_skirt' #cloth3d_skirt
|
||||
# src_folder = r'D:\Work\Dataset\validation_2'
|
||||
dest_folder = '/home/workspace/data/dataset/validation'
|
||||
|
||||
|
||||
# create_no_orient_body(src_folder)
|
||||
# fix_split_garment_generated_data(src_folder)
|
||||
# create_no_shape_unposed(src_folder)
|
||||
# check_if_skirt(src_folder)
|
||||
# gather_ids(src_folder, 'ids/test_non_skirt_ids.txt')
|
||||
# verify_poses(src_folder)
|
||||
# create_garment_type_table()
|
||||
# generate_faces()
|
||||
generate_frames_file(src_folder, ratio=1)
|
||||
# generate_frames_file_filtered_cape()
|
||||
Reference in New Issue
Block a user