From 8639992cd835f718ef897934d9ce510176d88bf2 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Jun 2023 18:28:42 +0200 Subject: [PATCH] create wrapper for dataset --- data/dataset.py | 14 +- data/dataset_cape.py | 306 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 317 insertions(+), 3 deletions(-) create mode 100644 data/dataset_cape.py diff --git a/data/dataset.py b/data/dataset.py index c9891f2..03e24c3 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -1,7 +1,7 @@ import os import sys -# os.environ["CUDA_VISIBLE_DEVICES"] = '6' +# os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3' import math import time @@ -14,7 +14,6 @@ import scipy.io as sio import plotly.express as ex import torchvision -from torch.utils.data import Dataset from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode @@ -22,6 +21,7 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.utils import Frame from data.cloth3d_io import readPC2, readPC2Frame +from data.dataset_cape import DatasetCape def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): @@ -131,7 +131,7 @@ def read_frames_file(frames_file, shuffle=False): return frames -class Dataset(Dataset): +class DatasetCloth3D(torch.utils.data.Dataset): def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): self.folder = folder self.device = device @@ -229,6 +229,14 @@ class Dataset(Dataset): data = { **data, **self.static_cache[frame] } return data +def Dataset(folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False): + if not is_cape: + return DatasetCloth3D(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=False) + else: + return DatasetCape(folder, signature, resolutions=resolutions, frames_file=frames_file, device=device, use_cache=use_cache, is_cape=True) + + + if __name__ == '__main__': from torch.utils.data import DataLoader diff --git a/data/dataset_cape.py b/data/dataset_cape.py new file mode 100644 index 0000000..4c4c370 --- /dev/null +++ b/data/dataset_cape.py @@ -0,0 +1,306 @@ +import os +import sys + +# os.environ["CUDA_VISIBLE_DEVICES"] = '6' + +import math +import time +import torch +import random +import pytorch3d +import torchvision as tv +import numpy as np +import scipy.io as sio +import plotly.express as ex +import torchvision + +from torchvision.transforms import Resize +from torchvision.transforms.functional import InterpolationMode + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +from data.utils import Frame + + +def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None): + if frames_file is None: + uv_static_folder = os.path.join(folder, 'pose_shape') + frames = sorted(os.listdir(uv_static_folder)) + + if is_cape: + frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('.')[0], frame.split('.')[1]) for frame in frames] + else: + frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] + + # if testing: + # frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) + # frames = [frame.split('.')[0] + '_4' for frame in frames] + + + if frames_file is not None: + frames = [] + with open(os.path.join(folder, frames_file), 'r') as f: + lines = f.readlines() + lines = [line.strip() for line in lines] + frames = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines] + + if shuffle: + random.shuffle(frames) + + return frames + + +class StaticCache(): + def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None: + self.store = dict() + self.signature = signature + 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 + + if is_cape: + self.uv_masks = dict() + uv_mask = torch.tensor(np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32), device=device) + self.uv_masks[512] = uv_mask.to(torch.float32).cpu().numpy() + for res in [32, 64, 128, 256]: + down = Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT)(uv_mask[None,:,:]) + down = torch.squeeze(down).to(torch.float32) + self.uv_masks[res] = down.cpu().numpy() + + + def read(self, id): + data = dict() + + 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 + 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 + + 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 + # 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) + 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 + + data = { key: value.to(self.device) for key, value in self.store[id].items()} + return data + + +class DatasetCape(torch.utils.data.Dataset): + def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False): + self.folder = folder + self.device = device + self.is_cape = is_cape + self.signature = signature + self.resolutions = resolutions + + self.transforms = tv.transforms.ToTensor() + self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, 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 + + + def __len__(self): + return len(self.frames) + + + def __getitem__(self, idx): + frame : Frame = self.frames[idx] + if self.cache is not None: + if frame in self.cache.keys(): + data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()} + 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() + + 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 + 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 + 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 + 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) + 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 + + 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_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'] + 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 + + 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'] + 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()} + + for key, value in data.items(): + if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7): + data[key] = torch.permute(value, [2,0,1]) + + if self.cache is not None: + 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] } + return data + + +if __name__ == '__main__': + from torch.utils.data import DataLoader + + device = "cuda" if torch.cuda.is_available() else "cpu" + + signature = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_posed', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'mesh_unposed', + 'uv_unposed', + 'uv_curvature', + 'mesh_curvature', + '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') + train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) + + for batch in train_dataloader: + batch + pass + + pass + + + + + + + + + + +# if __name__ == '__main__': +# os.environ["CUDA_VISIBLE_DEVICES"] = "6" +# from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces +# import plotly.graph_objects as go +# import plotly.express as ex +# # import cmapy as cm +# uv2mesh = UV2Mesh(is_skirt=True) + +# F = np.load("custom_utils/smpl_faces.npy") + +# sum_uv = np.zeros((512, 256, 3), dtype=np.float32) +# sum_mask = np.zeros((512, 256), dtype=np.float32) +# 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() +# sum_uv += batch['uv_unposed'].numpy().sum(axis=0) +# sum_mask += batch['uv_mask'].numpy().sum(axis=0) + +# sum_mask[sum_mask == 0] = 1 +# avg_garment = sum_uv / sum_mask[:, :, np.newaxis] + +# # x = np.transpose(avg_garment, [1,2,0]) +# x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0] +# go.Figure(data=[ +# go.Mesh3d( +# # 8 vertices of a cube +# x=x[:,0], +# y=x[:,1], +# z=x[:,2], +# # i, j and k give the vertices of triangles +# i = F[:,0], +# j = F[:,1], +# k = F[:,2] +# )]).show() + +# np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment) + +# print() \ No newline at end of file