299 lines
12 KiB
Python
299 lines
12 KiB
Python
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 torch.utils.data import Dataset
|
|
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]
|
|
else:
|
|
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
|
|
|
|
if testing:
|
|
frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
|
|
frames = [frame.split('.')[0] + '_4' for frame in frames]
|
|
|
|
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) 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 = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
|
|
self.uv_masks[512] = uv_mask.astype(np.float32)
|
|
for res in [32, 64, 128, 256]:
|
|
down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)
|
|
down = np.squeeze(down.numpy().astype(np.float32))
|
|
self.uv_masks[res] = down
|
|
|
|
|
|
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 Dataset(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)
|
|
|
|
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'))
|
|
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)
|
|
|
|
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).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
|
|
|
|
|
|
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:
|
|
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() |