85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
import scipy as sp
|
|
import scipy.io as sio
|
|
from tqdm import tqdm
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
from utils.utils import mask_faces
|
|
from superSMPL.IO import readPC2, writePC2
|
|
|
|
def generate_faces(folder):
|
|
frames = os.listdir(os.path.join(folder, 'mesh_mask'))
|
|
F = np.load('custom_data/smpl_faces.npy')
|
|
os.makedirs(os.path.join(folder, 'faces'), exist_ok=True)
|
|
|
|
for frame in tqdm(frames):
|
|
mask = np.load(os.path.join(folder, 'mesh_mask', frame))
|
|
faces = mask_faces(F, mask)
|
|
np.save(os.path.join(folder, 'faces', frame), faces)
|
|
|
|
|
|
class Frame:
|
|
def __init__(self, id, frame_nr):
|
|
self.id = id
|
|
self.frame_nr = frame_nr
|
|
|
|
|
|
def generate_frames_file(folder, file_name):
|
|
ids = os.listdir(folder)#[:1]
|
|
frames = []
|
|
|
|
for id in tqdm(ids):
|
|
garment_dir = os.path.join(folder, id)
|
|
if not os.path.isdir(garment_dir):
|
|
continue
|
|
garments = [f.replace('_registered.pc16','') for f in os.listdir(garment_dir) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
|
|
assert len(garments) == 1
|
|
garment = garments[0]
|
|
verts = readPC2(os.path.join(garment_dir, garment + '_registered.pc16'), True)['V']
|
|
for i in range(verts.shape[0]):
|
|
frames.append(Frame(id, i))
|
|
|
|
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) + '\n')
|
|
|
|
|
|
def genereated_registered_and_mask(folder, is_skirt = False):
|
|
ids = os.listdir(folder)
|
|
|
|
for id in 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 not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')]
|
|
|
|
assert len(garments) == 1
|
|
|
|
if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')):
|
|
continue
|
|
|
|
garment = garments[0]
|
|
|
|
verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
|
|
map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1))
|
|
verts = verts[:, map, :] # ---------------> This is the garment after registration
|
|
raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1))
|
|
mask = raw_mask > 0
|
|
# boundaries = raw_mask == 2
|
|
|
|
writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True)
|
|
np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# folder = '/home/workspace/local-data/train'
|
|
# folder = '/home/workspace/local-data/valid'
|
|
folder = 'dataset/CLOTH3D/train'
|
|
|
|
# generate_faces(folder)
|
|
# genereated_registered_and_mask(folder)
|
|
generate_frames_file(folder, 'frames_all.txt')
|