luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

create wrapper for dataset

Browse Source
This commit is contained in:
luni
2023-06-30 18:28:42 +02:00
parent dbbcd5f14e
commit 8639992cd8
2 changed files with 317 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+11 -3
View File
@@ -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
data/dataset_cape.py
+306
View File
@@ -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()