61 lines
1.5 KiB
Python
61 lines
1.5 KiB
Python
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') |