From 67891530a2cf4aae7cc712ee41c6f8510da085fe Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 23 Dec 2022 16:49:30 +0000 Subject: [PATCH] use raw cloth3d input --- data/dataset.py | 28 ++++++++++++++++------------ data/preprocess.py | 15 +++++---------- utils/data_utils.py | 33 +++++++++++++++++++++++++++++++++ utils/utils.py | 10 ---------- 4 files changed, 54 insertions(+), 32 deletions(-) create mode 100644 utils/data_utils.py diff --git a/data/dataset.py b/data/dataset.py index f6aeefb..44ba593 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -1,7 +1,7 @@ import os import sys -os.environ["CUDA_VISIBLE_DEVICES"] = '1' +os.environ["CUDA_VISIBLE_DEVICES"] = '0' import math import time @@ -15,12 +15,16 @@ from torch.utils.data import Dataset from render import Render, DepthRenderer +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') +from superSMPL.IO import readPC2Frame, writePC2 +from utils.data_utils import Frame, read_frames_file, mask_faces + class Dataset(Dataset): - def __init__(self, folder): + def __init__(self, folder, frames_file='frames_all.txt'): self.folder = folder self.transforms = tv.transforms.ToTensor() - self.frames = os.listdir(os.path.join(folder, 'mesh_posed')) + self.frames = read_frames_file(os.path.join(folder, frames_file)) self.render = DepthRenderer() self.body_faces = np.load('custom_data/smpl_faces.npy') @@ -31,16 +35,15 @@ class Dataset(Dataset): def __getitem__(self, idx): - path = os.path.join(self.folder, 'mesh_posed', self.frames[idx]) - mesh_posed = np.load(path)['mesh_posed'] + frame : Frame = self.frames[idx] + path = os.path.join(self.folder, frame.id) + mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True) - id = self.frames[idx].split('_')[0] - faces = np.load(os.path.join(self.folder, 'faces', id + '.npy')) - mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')) - body = np.load(os.path.join(self.folder, 'mesh_body_posed', self.frames[idx][:-1] + 'y')) + mask = np.load(os.path.join(path, frame.garment + '_mask.npy')) + faces = mask_faces(self.body_faces, mask) + body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True) - info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat' - info = sio.loadmat(info_folder) + info = sio.loadmat(os.path.join(path, 'info.mat')) combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda') combined_faces = torch.Tensor(np.concatenate((self.body_faces, faces + 14475))).to('cuda') @@ -60,7 +63,8 @@ if __name__ == '__main__': device = "cuda" if torch.cuda.is_available() else "cpu" - dataset = Dataset('/home/workspace/local-data/train') + # dataset = Dataset('/home/workspace/local-data/train') + dataset = Dataset('dataset/CLOTH3D/train') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True) for batch in train_dataloader: diff --git a/data/preprocess.py b/data/preprocess.py index 6a856d1..0d62503 100644 --- a/data/preprocess.py +++ b/data/preprocess.py @@ -7,10 +7,11 @@ 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 utils.data_utils import mask_faces, Frame from superSMPL.IO import readPC2, writePC2 -def generate_faces(folder): + +def generate_faces(folder, frames): 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) @@ -21,12 +22,6 @@ def generate_faces(folder): 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 = [] @@ -40,11 +35,11 @@ def generate_frames_file(folder, file_name): 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)) + 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) + '\n') + f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n') def genereated_registered_and_mask(folder, is_skirt = False): diff --git a/utils/data_utils.py b/utils/data_utils.py new file mode 100644 index 0000000..f754154 --- /dev/null +++ b/utils/data_utils.py @@ -0,0 +1,33 @@ +import os +import sys +import numpy as np + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + + +class Frame: + def __init__(self, id, frame_nr, garment): + self.id = id + self.frame_nr = frame_nr + self.garment = garment + + +def mask_faces(F, mask): + F_mask = np.take(mask, F.reshape(-1), axis=0) + F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 + + V_ind = np.zeros(mask.shape, dtype=np.int32) + V_ind[mask] = np.arange(mask.sum()) + F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) + + return F_new + + +def read_frames_file(file): + frames = [] + with open(file, 'r') as f: + lines = f.readlines() + for line in lines: + id, frame_nr, garment = line.split() + frames.append(Frame(id, int(frame_nr), garment)) + return frames \ No newline at end of file diff --git a/utils/utils.py b/utils/utils.py index 6ec80c9..337bc40 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -25,13 +25,3 @@ class Accumulator(): self.dict = dict() - -def mask_faces(F, mask): - F_mask = np.take(mask, F.reshape(-1), axis=0) - F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 - - V_ind = np.zeros(mask.shape, dtype=np.int32) - V_ind[mask] = np.arange(mask.sum()) - F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) - - return F_new \ No newline at end of file