Files
cvpr_2023/data/cloth3d_io.py
T
2023-02-06 22:46:11 +00:00

391 lines
13 KiB
Python

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)