started writing preprocessing
This commit is contained in:
@@ -0,0 +1,391 @@
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
from struct import pack, unpack
|
||||
|
||||
|
||||
"""
|
||||
Reads PC2 files, and proposed format PC16 files
|
||||
Inputs:
|
||||
- file: path to .pc2/.pc16 file
|
||||
- float16: False for PC2 files, True for PC16
|
||||
Output:
|
||||
- data: dictionary with .pc2/.pc16 file data
|
||||
NOTE: 16-bit floats lose precision with high values (positive or negative),
|
||||
we do not recommend using this format for data outside range [-2, 2]
|
||||
"""
|
||||
def readPC2(file, float16=False):
|
||||
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
|
||||
data = {}
|
||||
bytes = 2 if float16 else 4
|
||||
dtype = np.float16 if float16 else np.float32
|
||||
with open(file, 'rb') as f:
|
||||
# Header
|
||||
data['sign'] = f.read(12)
|
||||
# data['version'] = int.from_bytes(f.read(4), 'little')
|
||||
data['version'] = unpack('<i', f.read(4))[0]
|
||||
# Num points
|
||||
# data['nPoints'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nPoints'] = unpack('<i', f.read(4))[0]
|
||||
# Start frame
|
||||
data['startFrame'] = unpack('f', f.read(4))
|
||||
# Sample rate
|
||||
data['sampleRate'] = unpack('f', f.read(4))
|
||||
# Number of samples
|
||||
# data['nSamples'] = int.from_bytes(f.read(4), 'little')
|
||||
data['nSamples'] = unpack('<i', f.read(4))[0]
|
||||
# Animation data
|
||||
size = data['nPoints']*data['nSamples']*3*bytes
|
||||
data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
|
||||
data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
|
||||
|
||||
return data
|
||||
|
||||
"""
|
||||
Reads an specific frame of PC2/PC16 files
|
||||
Inputs:
|
||||
- file: path to .pc2/.pc16 file
|
||||
- frame: number of the frame to read
|
||||
- float16: False for PC2 files, True for PC16
|
||||
Output:
|
||||
- T: mesh vertex data at specified frame
|
||||
"""
|
||||
def readPC2Frame(file, frame, float16=False):
|
||||
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
|
||||
assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer'
|
||||
bytes = 2 if float16 else 4
|
||||
dtype = np.float16 if float16 else np.float32
|
||||
with open(file,'rb') as f:
|
||||
# Num points
|
||||
f.seek(16)
|
||||
# nPoints = int.from_bytes(f.read(4), 'little')
|
||||
nPoints = unpack('<i', f.read(4))[0]
|
||||
# Number of samples
|
||||
f.seek(28)
|
||||
# nSamples = int.from_bytes(f.read(4), 'little')
|
||||
nSamples = unpack('<i', f.read(4))[0]
|
||||
if frame > nSamples:
|
||||
print("Frame index outside size")
|
||||
print("\tN. frame: " + str(frame))
|
||||
print("\tN. samples: " + str(nSamples))
|
||||
return
|
||||
# Read frame
|
||||
size = nPoints * 3 * bytes
|
||||
f.seek(size * frame, 1) # offset from current '1'
|
||||
T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
|
||||
return T.reshape(nPoints, 3)
|
||||
|
||||
"""
|
||||
Writes PC2 and PC16 files
|
||||
Inputs:
|
||||
- file: path to file (overwrites if exists)
|
||||
- V: 3D animation data as a three dimensional array (N. Frames x N. Vertices x 3)
|
||||
- float16: False for writing as PC2 file, True for PC16
|
||||
This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1
|
||||
NOTE: 16-bit floats lose precision with high values (positive or negative),
|
||||
we do not recommend using this format for data outside range [-2, 2]
|
||||
"""
|
||||
def writePC2(file, V, float16=False):
|
||||
assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
|
||||
if float16: V = V.astype(np.float16)
|
||||
else: V = V.astype(np.float32)
|
||||
with open(file, 'wb') as f:
|
||||
# Create the header
|
||||
headerFormat='<12siiffi'
|
||||
headerStr = pack(headerFormat, b'POINTCACHE2\0',
|
||||
1, V.shape[1], 0, 1, V.shape[0])
|
||||
f.write(headerStr)
|
||||
# Write vertices
|
||||
f.write(V.tobytes())
|
||||
|
||||
"""
|
||||
Reads OBJ files
|
||||
Only handles vertices, faces and UV maps
|
||||
Input:
|
||||
- file: path to .obj file
|
||||
Outputs:
|
||||
- V: 3D vertices
|
||||
- F: 3D faces
|
||||
- Vt: UV vertices
|
||||
- Ft: UV faces
|
||||
Correspondence between mesh and UV map is implicit in F to Ft correspondences
|
||||
If no UV map data in .obj file, it shall return Vt=None and Ft=None
|
||||
"""
|
||||
def readOBJ(file):
|
||||
V, Vt, F, Ft = [], [], [], []
|
||||
with open(file, 'r') as f:
|
||||
T = f.readlines()
|
||||
for t in T:
|
||||
# 3D vertex
|
||||
if t.startswith('v '):
|
||||
v = [float(n) for n in t.replace('v ','').split(' ')]
|
||||
V += [v]
|
||||
# UV vertex
|
||||
elif t.startswith('vt '):
|
||||
v = [float(n) for n in t.replace('vt ','').split(' ')]
|
||||
Vt += [v]
|
||||
# Face
|
||||
elif t.startswith('f '):
|
||||
idx = [n.split('/') for n in t.replace('f ','').split(' ')]
|
||||
f = [int(n[0]) - 1 for n in idx]
|
||||
F += [f]
|
||||
# UV face
|
||||
if '/' in t:
|
||||
f = [int(n[1]) - 1 for n in idx]
|
||||
Ft += [f]
|
||||
V = np.array(V, np.float32)
|
||||
Vt = np.array(Vt, np.float32)
|
||||
if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces'
|
||||
else: Vt, Ft = None, None
|
||||
return V, F, Vt, Ft
|
||||
|
||||
|
||||
|
||||
"""
|
||||
Writes OBJ files
|
||||
Only handles vertices, faces and UV maps
|
||||
Inputs:
|
||||
- file: path to .obj file (overwrites if exists)
|
||||
- V: 3D vertices
|
||||
- F: 3D faces
|
||||
- Vt: UV vertices
|
||||
- Ft: UV faces
|
||||
Correspondence between mesh and UV map is implicit in F to Ft correspondences
|
||||
If no UV map data as input, it will write only 3D data in .obj file
|
||||
"""
|
||||
def writeOBJ(file, V, F, Vt=None, Ft=None):
|
||||
if not Vt is None:
|
||||
assert len(F) == len(Ft), 'Inconsistent data, mesh and UV map do not have the same number of faces'
|
||||
|
||||
with open(file, 'w') as file:
|
||||
# Vertices
|
||||
for v in V:
|
||||
line = 'v ' + ' '.join([str(_) for _ in v]) + '\n'
|
||||
file.write(line)
|
||||
# UV verts
|
||||
if not Vt is None:
|
||||
for v in Vt:
|
||||
line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n'
|
||||
file.write(line)
|
||||
# 3D Faces / UV faces
|
||||
if Ft:
|
||||
F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)]
|
||||
else:
|
||||
F = [[str(i + 1) for i in f] for f in F]
|
||||
for f in F:
|
||||
line = 'f ' + ' '.join(f) + '\n'
|
||||
file.write(line)
|
||||
|
||||
|
||||
def loadInfo(filename):
|
||||
'''
|
||||
this function should be called instead of direct sio.loadmat
|
||||
as it cures the problem of not properly recovering python dictionaries
|
||||
from mat files. It calls the function check keys to cure all entries
|
||||
which are still mat-objects
|
||||
'''
|
||||
data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True)
|
||||
del data['__globals__']
|
||||
del data['__header__']
|
||||
del data['__version__']
|
||||
return _check_keys(data)
|
||||
|
||||
def _check_keys(dict):
|
||||
'''
|
||||
checks if entries in dictionary are mat-objects. If yes
|
||||
todict is called to change them to nested dictionaries
|
||||
'''
|
||||
for key in dict:
|
||||
if isinstance(dict[key], sio.matlab.mio5_params.mat_struct):
|
||||
dict[key] = _todict(dict[key])
|
||||
return dict
|
||||
|
||||
def _todict(matobj):
|
||||
'''
|
||||
A recursive function which constructs from matobjects nested dictionaries
|
||||
'''
|
||||
dict = {}
|
||||
for strg in matobj._fieldnames:
|
||||
elem = matobj.__dict__[strg]
|
||||
if isinstance(elem, sio.matlab.mio5_params.mat_struct):
|
||||
dict[strg] = _todict(elem)
|
||||
elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in elem]):
|
||||
dict[strg] = [None] * len(elem)
|
||||
for i,item in enumerate(elem):
|
||||
if isinstance(item, sio.matlab.mio5_params.mat_struct):
|
||||
dict[strg][i] = _todict(item)
|
||||
else:
|
||||
dict[strg][i] = item
|
||||
else:
|
||||
dict[strg] = elem
|
||||
return dict
|
||||
|
||||
|
||||
# import os
|
||||
# import sys
|
||||
# import numpy as np
|
||||
# from PIL import Image
|
||||
|
||||
# sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
|
||||
# from cloth3d.smpl_np import SMPLModel
|
||||
# from cloth3d.util import loadInfo, zRotMatrix, proj, mesh2UV, uv_to_pixel
|
||||
# from superSMPL.IO import readOBJ, readPC2Frame
|
||||
|
||||
# class DataReader:
|
||||
|
||||
# def __init__(self, folder, smpl_path):
|
||||
# # Data paths
|
||||
# self.SRC = folder
|
||||
# # SMPL model
|
||||
# smpl_path = smpl_path
|
||||
# self.smpl = {
|
||||
# 'f': SMPLModel(os.path.join(smpl_path, 'super_model_f.pkl')),
|
||||
# 'm': SMPLModel(os.path.join(smpl_path, 'super_model_m.pkl'))
|
||||
# }
|
||||
|
||||
# """
|
||||
# Read sample info
|
||||
# Input:
|
||||
# - sample: name of the sample e.g.:'01_01_s0'
|
||||
# """
|
||||
# def read_info(self, sample):
|
||||
# info_path = os.path.join(self.SRC, sample, 'info')
|
||||
# return loadInfo(info_path)
|
||||
|
||||
# """ Human data """
|
||||
# """
|
||||
# Read SMPL parameters for the specified sample and frame
|
||||
# Inputs:
|
||||
# - sample: name of the sample
|
||||
# - frame: frame number
|
||||
# """
|
||||
# def read_smpl_params(self, sample, frame):
|
||||
# # Read sample data
|
||||
# info = self.read_info(sample)
|
||||
# # SMPL parameters
|
||||
# gender = 'm' if info['gender'] else 'f'
|
||||
# pose = info['poses'][:, frame].reshape(self.smpl[gender].pose_shape)
|
||||
# shape = info['shape']
|
||||
# trans = info['trans'][:, frame].reshape(self.smpl[gender].trans_shape)
|
||||
# return gender, pose, shape, trans
|
||||
|
||||
# """
|
||||
# Computes human mesh for the specified sample and frame
|
||||
# Inputs:
|
||||
# - sample: name of the sample
|
||||
# - frame: frame number
|
||||
# - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint
|
||||
# Outputs:
|
||||
# - V: human mesh vertices
|
||||
# - F: mesh faces
|
||||
# """
|
||||
# def read_human(self, sample, frame, zrot, use_trans=True):
|
||||
# # SMPL parameters
|
||||
# gender, pose, shape, trans = self.read_smpl_params(sample, frame)
|
||||
# # Compute SMPL
|
||||
# V, J = self.smpl[gender].set_params(pose=pose, beta=shape, trans=trans if use_trans else None)
|
||||
# V -= J[0:1]
|
||||
# # Apply rotation on z-axis
|
||||
# zRot = zRotMatrix(zrot)
|
||||
# return zRot.dot(V.T).T, self.smpl[gender].faces, zRot.dot(trans.T).T[None,:], zRot.dot((J - J[0:1] + trans[None,:]).T).T[None,:], zRot.dot((J - J[0:1]).T).T[None,:]
|
||||
|
||||
# """ Garment data """
|
||||
# """
|
||||
# Reads garment vertices location for the specified sample, garment and frame
|
||||
# Inputs:
|
||||
# - sample: name of the sample
|
||||
# - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
|
||||
# - frame: frame number
|
||||
# - TODO: absolute: True for absolute vertex locations, False for locations relative to SMPL root joint
|
||||
# Outputs:
|
||||
# - V: 3D vertex locations for the specified sample, garment and frame
|
||||
# """
|
||||
# def read_garment_vertices(self, sample, garment, frame, trans, zrot):
|
||||
# # Read garment vertices (relative to root joint)
|
||||
# pc16_path = os.path.join(self.SRC, sample, garment + '.pc16')
|
||||
# V = readPC2Frame(pc16_path, frame, True)
|
||||
# # Read sample data
|
||||
# if trans is not None:
|
||||
# # Transform to absolute
|
||||
# V += trans[None]
|
||||
# zRot = zRotMatrix(zrot)
|
||||
# return zRot.dot(V.T).T
|
||||
|
||||
# """
|
||||
# Reads garment faces for the specified sample and garment
|
||||
# Inputs:
|
||||
# - sample: name of the sample
|
||||
# - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
|
||||
# Outputs:
|
||||
# - F: mesh faces
|
||||
# """
|
||||
# def read_garment_topology(self, sample, garment):
|
||||
# # Read OBJ file
|
||||
# obj_path = os.path.join(self.SRC, sample, garment + '.obj')
|
||||
# return readOBJ(obj_path)[1]
|
||||
|
||||
# """
|
||||
# Reads garment UV map for the specified sample and garment
|
||||
# Inputs:
|
||||
# - sample: name of the sample
|
||||
# - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
|
||||
# Outputs:
|
||||
# - Vt: UV map vertices
|
||||
# - Ft: UV map faces
|
||||
# """
|
||||
# def read_garment_UVMap(self, sample, garment):
|
||||
# # Read OBJ file
|
||||
# obj_path = os.path.join(self.SRC, sample, garment + '.obj')
|
||||
# return readOBJ(obj_path)[2:]
|
||||
|
||||
# """
|
||||
# Reads vertex colors of the specified sample and garment
|
||||
# Inputs:
|
||||
# - sample: name of the sample
|
||||
# - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
|
||||
# - F: mesh faces
|
||||
# - Vt: UV map vertices
|
||||
# - Ft: UV map faces
|
||||
# Output
|
||||
# - C: RGB colors
|
||||
# """
|
||||
# def read_garment_vertex_colors(self, sample, garment, F, Vt, Ft):
|
||||
# # Read sample texture info
|
||||
# texture = self.read_info(sample)['outfit'][garment]['texture']
|
||||
# # Plain color
|
||||
# if texture['type'] == 'color': return (255*texture['data']).astype(np.int32)
|
||||
# # Image texture
|
||||
# img_path = os.path.join(self.SRC, sample, garment + '.png')
|
||||
# img = Image.open(img_path)
|
||||
# # Get color of each UV vertex
|
||||
# color = np.array([img.getpixel(uv_to_pixel(vt)) for vt in Vt])
|
||||
# # Compute correspondece between V (mesh) and Vt (UV map)
|
||||
# m2uv = mesh2UV(F, Ft)
|
||||
# return np.array([color[list(m2uv[idx])].mean(0) for idx in sorted(list(m2uv.keys()))], np.uint8)
|
||||
|
||||
# """ Scene data """
|
||||
# """
|
||||
# Read camera location and compute projection matrix
|
||||
# Input:
|
||||
# - sample: name of the sample
|
||||
# Output:
|
||||
# - P: camera projection matrix (3 x 4)
|
||||
# """
|
||||
# def read_camera(self, sample):
|
||||
# # Read sample data
|
||||
# info = self.read_info(sample)
|
||||
# # Camera location
|
||||
# camLoc = info['camLoc']
|
||||
# # Camera projection matrix
|
||||
# return proj(camLoc)
|
||||
|
||||
# # TESTING
|
||||
# if __name__ == '__main__':
|
||||
# sample = '135_02_s8'
|
||||
# frame = 0
|
||||
# garment = 'Tshirt'
|
||||
|
||||
# reader = DataReader()
|
||||
# F = reader.read_garment_topology(sample, garment)
|
||||
# Vt, Ft = reader.read_garment_UVMap(sample, garment)
|
||||
# C = reader.read_garment_vertex_colors(sample, garment, F, Vt, Ft)
|
||||
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import shutil
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from data.cloth3d_io import readPC2, loadInfo
|
||||
from data.utils import Frame
|
||||
|
||||
|
||||
|
||||
|
||||
def clean_dataset(folder):
|
||||
ids = os.listdir(folder)
|
||||
|
||||
for id in tqdm(ids):
|
||||
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
files = os.listdir(os.path.join(folder, id))
|
||||
|
||||
for file in files:
|
||||
if file.startswith('glmarks') or file.find('nricp') != -1:
|
||||
os.remove(os.path.join(folder, id, file))
|
||||
|
||||
|
||||
def generate_frames_file(folder, file_name):
|
||||
ids = os.listdir(folder)
|
||||
frames = []
|
||||
|
||||
for id in tqdm(ids):
|
||||
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')]
|
||||
info = loadInfo(os.path.join(folder, id, 'info.mat'))
|
||||
|
||||
for garment in garments:
|
||||
for i in range(info['poses'].shape[1]):
|
||||
frames.append(Frame(id, i, garment))
|
||||
|
||||
with open(os.path.join(folder, file_name), 'w') as f:
|
||||
for i, frame in enumerate(frames):
|
||||
f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
folder = '/home/workspace/data/dataset/cloth3d/train'
|
||||
|
||||
# clean_dataset(folder)
|
||||
generate_frames_file(folder, 'frames_all.txt')
|
||||
@@ -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()
|
||||
@@ -0,0 +1,14 @@
|
||||
|
||||
|
||||
class Frame:
|
||||
def __init__(self, id, frame_nr, garment):
|
||||
self.id = id
|
||||
self.frame_nr = frame_nr
|
||||
self.garment = garment
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.id, self.frame_nr))
|
||||
|
||||
def __eq__(self, other):
|
||||
return (self.id, self.frame_nr) == (other.id, other.frame_nr)
|
||||
|
||||
Reference in New Issue
Block a user