Files
cvpr/preprocessing.py
T

217 lines
7.2 KiB
Python

import os
from custom_utils.supersmpl_io import readPC2
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
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
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)
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()
# print(gender[0,0])
# meshes = []
# x = T[0, :, :]
# meshes.append(go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,2],
# z=x[:,1],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2]
# ))
# 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,1],),
# yaxis = dict(nticks=4, range=[-1,1],),
# zaxis = dict(nticks=4, range=[-1,1],),)
# )
# fig.show()
# frame_nr = str(0)
# folder_2 = '/home/workspace/local-data/train_100/mesh_posed_body_no_orient'
# mesh_posed_body_no_orient = np.load(os.path.join(folder_2, id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
# meshes = []
# x = mesh_posed_body_no_orient
# meshes.append(go.Mesh3d(
# # 8 vertices of a cube
# x=x[:,0],
# y=x[:,2],
# z=x[:,1],
# # i, j and k give the vertices of triangles
# i = F[:,0],
# j = F[:,1],
# k = F[:,2]
# ))
# 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,1],),
# yaxis = dict(nticks=4, range=[-1,1],),
# zaxis = dict(nticks=4, range=[-1,1],),)
# )
# fig.show()
# print()
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'))
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'), exist_ok=True)
for i in range(info['poses'].shape[1]):
mesh_unposed = np.load(os.path.join(garment_dir, 'mesh_unposed', 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', 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)
if __name__ == '__main__':
src_folder = '/home/workspace/data/dataset/test_skirt' #cloth3d_skirt
# src_folder = r'D:\Work\Dataset\validation_2'
dest_folder = '/home/workspace/data/dataset/validation'
# create_no_orient_body(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)