diff --git a/data/dataset.py b/data/dataset.py index 99aa5e8..c9891f2 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -21,6 +21,7 @@ from torchvision.transforms.functional import InterpolationMode sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.utils import Frame +from data.cloth3d_io import readPC2, readPC2Frame def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): @@ -60,7 +61,6 @@ class StaticCache(): self.folder = folder assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) self.resolutions = resolutions - self.filled = '' self.is_cape = is_cape self.device = device @@ -74,50 +74,65 @@ class StaticCache(): self.uv_masks[res] = down.cpu().numpy() - def read(self, id): + def read(self, id, garment): data = dict() + mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool) + data['mesh_mask'] = mesh_mask + for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_static' + res_string in self.signature: - uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static + uv_static = np.load(os.path.join(self.folder, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static if 'uv_mask' + res_string in self.signature: if self.is_cape: data['uv_mask' + res_string] = self.uv_masks[res] else: - uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask + uv_mask = np.load(os.path.join(self.folder, f'uv_mask{res_string}', f'{id}_{garment}.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask if 'mesh_static' in self.signature: - mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32) - if 'mesh_mask' in self.signature: - if self.is_cape: - data['mesh_mask'] = np.ones((3638), np.float32) - else: - mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask + mesh_static = np.load(os.path.join(self.folder, 'mesh_static', f'{id}_{garment}.npy')) + data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32) + data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32) + # if 'uv_shaped_body' in self.signature: # uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body return data - def __getitem__(self, id) -> dict: - if id not in self.store: - item = self.read(id) + def __getitem__(self, frame) -> dict: + if frame not in self.store: + item = self.read(frame.id, frame.garment) stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} for key, value in stored_data.items(): if key.startswith('uv_') and value.shape[-1] == 3: stored_data[key] = torch.permute(value, [2,0,1]) - self.store[id] = stored_data + self.store[frame] = stored_data - data = { key: value.to(self.device) for key, value in self.store[id].items()} + data = { key: value.to(self.device) for key, value in self.store[frame].items()} return data + +def read_frames_file(frames_file, shuffle=False): + frames = [] + with open(frames_file, 'r') as f: + lines = f.readlines() + lines = [line.strip() for line in lines] + frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines] + + if shuffle: + random.shuffle(frames) + + return frames + + class Dataset(Dataset): - def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False): + def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): self.folder = folder self.device = device self.is_cape = is_cape @@ -125,13 +140,20 @@ class Dataset(Dataset): self.resolutions = resolutions self.transforms = tv.transforms.ToTensor() - self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True) + self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True) - self.body_faces = np.load('smpl/smpl_faces.npy') self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device) self.cache = dict() if use_cache else None + ids = set([frame.id for frame in self.frames]) + self.pose_cache = dict() + for id in ids: + self.pose_cache[id] = np.load(os.path.join(self.folder, 'pose', f'{id}.npy')) + self.shape_cache = dict() + for id in ids: + self.shape_cache[id] = np.load(os.path.join(self.folder, 'shape', f'{id}.npy')) + def __len__(self): return len(self.frames) @@ -145,73 +167,56 @@ class Dataset(Dataset): data = { **data, **self.static_cache[frame.id] } return data - separator = '_' - if self.is_cape: - separator = '.' - - uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2' - if self.is_cape: - uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' - item = dict() + mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool) + + item['id'] = frame.id + item['frame_nr'] = frame.frame_nr + item['garment'] = frame.garment - if 'id' in self.signature: - item['id'] = frame.id - if 'frame_nr' in self.signature: - item['frame_nr'] = frame.frame_nr if 'pose_shape' in self.signature: - pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32) - pose_shape[:3] = 0 - item['pose_shape'] = pose_shape + pose = self.pose_cache[frame.id][frame.frame_nr] + pose[:3] = 0 + item['pose_shape'] = np.concatenate((pose, self.shape_cache[frame.id]), axis=0) if 'mesh_body_posed' in self.signature: - mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] - item['mesh_body_posed'] = mesh_posed_body_no_orient + mesh_body_posed_no_orient = readPC2Frame(os.path.join(self.folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) + item['mesh_body_posed'] = mesh_body_posed_no_orient if 'mesh_posed' in self.signature: - mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32) - item['mesh_posed'] = mesh_posed + mesh_posed = readPC2Frame(os.path.join(self.folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + item['mesh_posed'] = np.zeros((14475, 3), dtype=np.float32) + item['mesh_posed'][mesh_mask] = mesh_posed if 'mesh_unposed' in self.signature: - mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] - item['mesh_unposed'] = mesh_unposed_no_shape - if 'mesh_curvature' in self.signature: - mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] - item['mesh_curvature'] = mesh_curvature.astype(np.int32) + mesh_unposed_no_shape = readPC2Frame(os.path.join(self.folder, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + item['mesh_unposed'] = np.zeros((14475, 3), dtype=np.float32) + item['mesh_unposed'][mesh_mask] = mesh_unposed_no_shape if 'mesh_normals' in self.signature: - mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals'] - item['mesh_normals'] = mesh_normals + mesh_normals = readPC2Frame(os.path.join(self.folder, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + item['mesh_normals'] = np.zeros((14475, 3), dtype=np.float32) + item['mesh_normals'][mesh_mask] = mesh_normals for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) - if 'uv_curvature' + res_string in self.signature: - uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); - if res_string != '': - uv_curvature = uv_curvature['uv_curvature'] - else: - uv_curvature = uv_curvature['curvature_uv'] - uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device) - item['uv_curvature' + res_string] = uv_curvature + if 'uv_normals' + res_string in self.signature: + uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_normals = uv_normals['uv_normals'] + item['uv_normals' + res_string] = uv_normals if 'uv_unposed' + res_string in self.signature: - uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] + uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed'] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in self.signature: - uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz')) - uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] - item['uv_body_posed' + res_string] = uv_posed_body_no_orient + uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed'] + item['uv_body_posed' + res_string] = uv_body_posed_no_orient if 'uv_posed' + res_string in self.signature: - uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] + uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed'] item['uv_posed' + res_string] = uv_posed - if 'uv_normals' + res_string in self.signature: - uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals'] - item['uv_normals' + res_string] = uv_normals - - data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()} + data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor or type(value) == int) else value) for key, value in item.items()} for key, value in data.items(): if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7): @@ -221,38 +226,93 @@ class Dataset(Dataset): self.cache[frame] = data data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()} - data = { **data, **self.static_cache[frame.id] } + data = { **data, **self.static_cache[frame] } return data if __name__ == '__main__': from torch.utils.data import DataLoader - + import plotly.graph_objects as go + import plotly.express as ex + from plotly.subplots import make_subplots + from custom_utils.custom_utils import _maskFaces + + F = np.load('smpl/smpl_faces.npy') + F_skirt = np.load('smpl/smpl_faces_skirt.npy') + device = "cuda" if torch.cuda.is_available() else "cpu" - signature = { 'id', - 'frame_nr', - 'pose_shape', + signature = { 'pose_shape', 'uv_posed', 'uv_body_posed', 'mesh_body_posed', 'mesh_posed', 'mesh_unposed', 'uv_unposed', - 'uv_curvature', - 'mesh_curvature', + 'uv_normals', + 'mesh_normals', 'uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', - 'uv_shaped_body', } - dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') + dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) for batch in train_dataloader: batch + b = 0 + mesh_mask = batch['mesh_mask'][b].cpu().numpy() + faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F + _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) + + mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask] + mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask] + mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask] + mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy() + + fig = make_subplots(rows=2, cols=2, + specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}], + [{'type': 'mesh3d'}, {'type': 'mesh3d'}]]) + + fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1) + fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2) + fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1) + fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2) + + fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube'), + scene2=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube'), + scene3=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube'), + scene4=dict(xaxis=dict(range=[-1, 1], autorange=False), + yaxis=dict(range=[-1, 1], autorange=False), + zaxis=dict(range=[-1, 1], autorange=False), + aspectmode='cube')) + + fig.show() + + + uv_mask = batch['uv_mask'][b].cpu().numpy() + uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0) + uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2 + uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2 + uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2 + uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2 + uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3 + + img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2) + img = np.transpose(img, (1, 2, 0)) + ex.imshow(img).show() + pass pass @@ -281,7 +341,7 @@ if __name__ == '__main__': # for i, batch in enumerate(dataset): # print(i) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) -# ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show() +# ex.imshow(batch['uv_normals'][0].numpy().argmax(-1) / 6.0).show() # sum_uv += batch['uv_unposed'].numpy().sum(axis=0) # sum_mask += batch['uv_mask'].numpy().sum(axis=0) diff --git a/data/preprocessing.py b/data/preprocessing.py index 098285c..1bf70a3 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -510,6 +510,62 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1) # fig.show() pass + +def create_frames_file(folder, out_folder, split=0.9): + ids = os.listdir(folder) + + frames_non_skirt = [] + frames_skirt = [] + ids_non_skirt = set() + ids_skirt = set() + + for id in tqdm(ids): + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + frames = info['poses'].shape[1] + + assert info['poses'].shape[0] == 72 + + for garment in garments: + for i in range(frames): + if garment == 'Skirt' or garment == 'Dress': + frames_skirt.append(Frame(id, i, garment)) + ids_skirt.add(id) + else: + frames_non_skirt.append(Frame(id, i, garment)) + ids_non_skirt.add(id) + + nr_train_non_skirt = int(len(ids_non_skirt) * split) + nr_train_skirt = int(len(ids_skirt) * split) + + ids_skirt = list(ids_skirt) + ids_non_skirt = list(ids_non_skirt) + + with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt[:nr_train_non_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt[:nr_train_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt[nr_train_non_skirt:]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt[nr_train_skirt:]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + if __name__ == '__main__': @@ -530,4 +586,6 @@ if __name__ == '__main__': # create_mesh_static(folder, out_folder) # create_mesh_body_posed_no_orient(folder, out_folder) # create_mesh_normals(folder, out_folder) - create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) \ No newline at end of file + # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) + + create_frames_file(folder, out_folder) \ No newline at end of file diff --git a/data/utils.py b/data/utils.py index 6f25887..f6b796c 100644 --- a/data/utils.py +++ b/data/utils.py @@ -1,7 +1,7 @@ import tables as tb class Frame: - def __init__(self, id, frame_nr, garment=''): + def __init__(self, id: str, frame_nr: int, garment: str=''): self.id = id self.frame_nr = frame_nr self.garment = garment