new dataloader
This commit is contained in:
+204
-458
@@ -1,79 +1,101 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
||||
|
||||
import math
|
||||
import time
|
||||
import cv2
|
||||
import torch
|
||||
import random
|
||||
import pickle
|
||||
import pytorch3d
|
||||
import torchvision as tv
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
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 custom_utils.custom_utils import nest_conditions, batch_to_dict
|
||||
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame
|
||||
from data.utils import Frame
|
||||
|
||||
|
||||
class StaticDataset():
|
||||
def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False, use_no_orient=True) -> None:
|
||||
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]
|
||||
|
||||
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], in_memory=True, is_cape=False) -> None:
|
||||
self.store = dict()
|
||||
self.signature = signature
|
||||
self.folder = folder
|
||||
self.in_memory = in_memory
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
|
||||
'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256',
|
||||
'mesh_static', 'mesh_mask', 'uv_shaped_body']
|
||||
self.keys = [value for value in signature.keys() if value in self.possible_keys]
|
||||
if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
|
||||
self.keys.append('uv_static_128')
|
||||
self.resolutions = resolutions
|
||||
self.filled = ''
|
||||
self.is_cape = is_cape
|
||||
self.use_no_orient = use_no_orient
|
||||
|
||||
|
||||
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 = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
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, garment):
|
||||
def read(self, id):
|
||||
data = dict()
|
||||
|
||||
info = loadInfo(os.path.join(self.folder, id))
|
||||
|
||||
data['shape'] = info['shape']
|
||||
data['poses'] = info['poses']
|
||||
|
||||
data['map'] = sio.loadmat(os.path.join(self.folder, id, garment + '_map.mat'))['map'].reshape((-1))
|
||||
|
||||
|
||||
if 'mesh_mask' in self.signature.keys():
|
||||
if self.is_cape:
|
||||
data['mesh_mask'] = np.ones((3638), np.float32)
|
||||
else:
|
||||
data['mesh_mask'] = read_garment_mask(self.folder, id, garment).astype(np.float32)
|
||||
|
||||
for res in self.resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_static' + res_string in self.signature.keys():
|
||||
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.keys():
|
||||
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.keys():
|
||||
data['mesh_static'] = read_garment_template(self.folder, id, garment)
|
||||
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
|
||||
|
||||
@@ -90,462 +112,186 @@ class StaticDataset():
|
||||
|
||||
|
||||
def expand_batch(self, batch):
|
||||
for key in self.keys:
|
||||
batch[key] = tf.zeros(((0,) + self.signature[key]))
|
||||
for i in range(batch['id'].shape[0]):
|
||||
item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'), batch['garment'][i].numpy().decode('UTF-8'))
|
||||
for key in self.keys:
|
||||
if key == 'pose':
|
||||
frame_nr = int(batch['frame_nr'][i].numpy().decode('UTF-8'))
|
||||
pose = item[key][frame_nr, :]
|
||||
if self.use_no_orient:
|
||||
pose[:3] = 0
|
||||
batch[key] = tf.concat((batch[key], tf.expand_dims(pose, axis=0)), axis=0)
|
||||
else:
|
||||
batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
|
||||
batch_size = batch['id'].shape[0]
|
||||
items = []
|
||||
for i in range(batch_size):
|
||||
items.append(self.__getitem__(batch['id'][i].numpy().decode('UTF-8')))
|
||||
for key in items[0].keys():
|
||||
batch[key] = torch.stack([item[key] for item in items], dim=0)
|
||||
return batch
|
||||
|
||||
|
||||
class Dataset(Dataset):
|
||||
def __init__(self, folder, signature, resolutions, 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
|
||||
|
||||
class DatasetIdx(tf.data.Dataset):
|
||||
self.transforms = tv.transforms.ToTensor()
|
||||
self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
|
||||
|
||||
def __new__(cls, 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))
|
||||
self.body_faces = np.load('smpl/smpl_faces.npy')
|
||||
self.static_cache = StaticCache(folder, signature, is_cape=is_cape)
|
||||
|
||||
# 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 = []
|
||||
with open(os.path.join(folder, frames_file), 'r') as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
frames = lines
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(frames)
|
||||
|
||||
return tf.data.Dataset.from_tensor_slices(frames)
|
||||
self.cache = dict() if use_cache else None
|
||||
|
||||
|
||||
class DatasetFactory():
|
||||
def __len__(self):
|
||||
return len(self.frames)
|
||||
|
||||
def __init__(self):
|
||||
super(DatasetFactory, self).__init__()
|
||||
|
||||
def _build_tf_signature(self, signature):
|
||||
tf_signature = []
|
||||
for key in signature.keys():
|
||||
if key == 'id':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'frame_nr':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'garment':
|
||||
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
|
||||
elif key == 'mesh_curvature':
|
||||
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
|
||||
else:
|
||||
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
|
||||
return tf_signature
|
||||
|
||||
|
||||
def get_dataset(self,
|
||||
folder,
|
||||
batch_size,
|
||||
signature,
|
||||
shuffle=True,
|
||||
resolutions=[512],
|
||||
in_memory=False,
|
||||
do_repeat=True,
|
||||
is_cape=False,
|
||||
use_no_orient=True,
|
||||
testing=False,
|
||||
frames_file=None):
|
||||
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
||||
|
||||
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file)
|
||||
def __getitem__(self, idx):
|
||||
frame : Frame = self.frames[idx]
|
||||
if self.cache is not None:
|
||||
if frame in self.cache.keys():
|
||||
return self.cache[frame]
|
||||
|
||||
separator = '_'
|
||||
if is_cape:
|
||||
if self.is_cape:
|
||||
separator = '.'
|
||||
|
||||
uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
|
||||
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'
|
||||
|
||||
if use_no_orient == False:
|
||||
assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!'
|
||||
item = dict()
|
||||
|
||||
to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
|
||||
shapedirs = []
|
||||
with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
|
||||
female = pickle.load(f, encoding="latin1")
|
||||
shapedirs.append(female['shapedirs'][to_keep, :])
|
||||
with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
|
||||
male = pickle.load(f, encoding="latin1")
|
||||
shapedirs.append(male['shapedirs'][to_keep, :])
|
||||
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
|
||||
|
||||
def read_item(frame):
|
||||
frame_split = tf.strings.split(frame, sep=separator)
|
||||
id, frame_nr, garment = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8"), frame_split[2].numpy().decode("utf-8")
|
||||
item = dict()
|
||||
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
|
||||
|
||||
# redunda ntly loading info again
|
||||
pose = loadInfo(os.path.join(folder, id))['poses'][:,frame_nr]
|
||||
# !!!!!!!!!!!!!!!!!!!
|
||||
map = None
|
||||
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 'id' in signature.keys():
|
||||
item['id'] = id
|
||||
if 'frame_nr' in signature.keys():
|
||||
item['frame_nr'] = frame_nr
|
||||
if 'garment' in signature.keys():
|
||||
item['garment'] = garment
|
||||
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 'mesh_body_posed' in signature.keys():
|
||||
if use_no_orient:
|
||||
item['mesh_body_posed'] = read_posed_body_frame(folder, id, frame_nr, pose, deorient=use_no_orient)
|
||||
if 'mesh_posed' in signature.keys():
|
||||
if use_no_orient:
|
||||
item['mesh_posed'] = read_posed_garment_frame(folder, id, garment, frame, map, mask=None, shape=None, shapedirs=None, pose=None, unshape=False, deorient=False)
|
||||
# if 'mesh_unposed' in signature.keys():
|
||||
# mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
||||
# item['mesh_unposed'] = mesh_unposed_no_shape
|
||||
data = { **item, **self.static_cache[frame.id] }
|
||||
|
||||
if 'mesh_curvature' in signature.keys():
|
||||
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
|
||||
|
||||
|
||||
for res in resolutions:
|
||||
res_string = ''
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_curvature' + res_string in signature.keys():
|
||||
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz'));
|
||||
if res_string != '':
|
||||
uv_curvature = uv_curvature['uv_curvature']
|
||||
else:
|
||||
uv_curvature = uv_curvature['curvature_uv']
|
||||
uv_curvature = tf.one_hot(uv_curvature, 7)
|
||||
item['uv_curvature' + res_string] = uv_curvature
|
||||
|
||||
if 'uv_unposed' + res_string in signature.keys():
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_string, id + separator + 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 signature.keys():
|
||||
if use_no_orient:
|
||||
uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + 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
|
||||
else:
|
||||
uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
|
||||
item['uv_body_posed' + res_string] = uv_body_posed
|
||||
|
||||
if 'uv_posed' + res_string in signature.keys():
|
||||
if use_no_orient:
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
else:
|
||||
uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
|
||||
as_list = []
|
||||
for key in signature.keys():
|
||||
if key in item:
|
||||
as_list.append(item[key])
|
||||
|
||||
return as_list
|
||||
|
||||
|
||||
if in_memory:
|
||||
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
inp=[i],
|
||||
Tout=self._build_tf_signature(signature)
|
||||
), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True)
|
||||
else:
|
||||
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
|
||||
inp=[i],
|
||||
Tout=self._build_tf_signature(signature)
|
||||
), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True)
|
||||
|
||||
if do_repeat:
|
||||
dataset = dataset.repeat()
|
||||
|
||||
return dataset
|
||||
data = { key: (torch.from_numpy(value).to(self.device) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in data.items()}
|
||||
|
||||
if self.cache is not None:
|
||||
self.cache[frame] = data
|
||||
return data
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
# dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
|
||||
# for d in dataset:
|
||||
# print(d)
|
||||
|
||||
factory = DatasetFactory()
|
||||
signature_dynamic = { 'id': (), 'frame_nr': (), 'garment': (),
|
||||
# 'pose_shape': (72 + 11,),
|
||||
# 'uv_posed': (3, 512, 256,),
|
||||
# 'uv_body_posed': (512, 256, 3),
|
||||
# 'uv_body_posed_32': (32, 16, 3),
|
||||
# 'mesh_body_posed': (14475, 3,),
|
||||
# 'mesh_posed': (14475, 3,),
|
||||
# 'mesh_unposed': (14475, 3,),
|
||||
'uv_unposed': (512, 256, 3),
|
||||
# 'uv_unposed_32': (3, 32, 16,),
|
||||
'uv_curvature': (512, 256, 7),
|
||||
# 'mesh_curvature': (14475)
|
||||
}
|
||||
signature_static = { #'uv_static': (3, 512, 256,),
|
||||
# 'uv_static_32': (3, 32, 16,),
|
||||
# 'mesh_static': (14475, 3,),
|
||||
# 'mesh_mask': (14475,),
|
||||
'uv_mask': (512, 256,),
|
||||
# 'uv_shaped_body': (3, 512, 256,),
|
||||
}
|
||||
|
||||
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
|
||||
batch_size=1,
|
||||
signature=signature_dynamic,
|
||||
resolutions=[512],
|
||||
do_repeat=False,
|
||||
use_one_hot=False,
|
||||
is_cape=False,
|
||||
frames_file='frames_dress_30k.txt')
|
||||
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
|
||||
signature=signature_static,
|
||||
resolutions=[512],
|
||||
is_cape=False)
|
||||
|
||||
|
||||
# for i, batch in enumerate(dataset):
|
||||
# print(i)
|
||||
|
||||
|
||||
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()
|
||||
|
||||
# pass
|
||||
|
||||
|
||||
# uv_res = 256
|
||||
|
||||
# def m2uv_array(m2uv):
|
||||
# result = np.zeros((14475, 2), dtype=np.int32)
|
||||
# for k, v in m2uv.items():
|
||||
# if len(v) == 2:
|
||||
# result[k] = list(v)
|
||||
# elif len(v) == 1:
|
||||
# x = v.pop()
|
||||
# result[k] = [x, x]
|
||||
# else:
|
||||
# assert False
|
||||
|
||||
# return result
|
||||
|
||||
# def unique_uv_map_mask(uv_map_discreet):
|
||||
# mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
||||
# temp = np.zeros((uv_res * 2, uv_res))
|
||||
# for p in uv_map_discreet:
|
||||
# temp[p[0], p[1]] += 1
|
||||
# # temp = temp == 1
|
||||
# for i, p in enumerate(uv_map_discreet):
|
||||
# if temp[p[0], p[1]] == 1:
|
||||
# mask[i] = True
|
||||
# elif temp[p[0], p[1]] == 2:
|
||||
# mask[i] = True
|
||||
# temp[p[0], p[1]] = 0
|
||||
# return mask
|
||||
|
||||
# def uv_to_mesh(uv, m2uv):
|
||||
# res = uv.shape[0]
|
||||
# mesh = np.zeros((14475, 3))
|
||||
# uv_pixels = uv_map_discreet[m2uv]
|
||||
# # counter = np.ones((14475))
|
||||
# # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
|
||||
# mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
|
||||
# return mesh
|
||||
|
||||
|
||||
# uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
||||
# m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
||||
# m2uv = m2uv_array(m2uv)
|
||||
# uv_mask = np.load('custom_utils/mask_256x512.npy')
|
||||
# curvature_label_type = '6_25'
|
||||
|
||||
# uv_res = 256
|
||||
# uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
||||
# uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
||||
|
||||
# uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
|
||||
# uv_map_mask = unique_uv_map_mask(uv_map_discreet)
|
||||
# uv_counter = np.zeros((uv_map.shape[0]))
|
||||
# uv_counter[m2uv[:, 0]] = 1
|
||||
# uv_counter[m2uv[:, 1]] += 1
|
||||
# uv_counter = uv_counter[uv_map_mask]
|
||||
|
||||
|
||||
|
||||
# for i, batch in enumerate(dataset):
|
||||
# print(i)
|
||||
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
|
||||
# uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
|
||||
# mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
|
||||
# mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
|
||||
|
||||
|
||||
# values = mesh_unposed
|
||||
# uv_values = np.zeros((uv_map.shape[0], 3))
|
||||
# uv_values[m2uv[:, 0]] = values
|
||||
# uv_values[m2uv[:, 1]] += values
|
||||
# uv_counter = np.zeros((uv_map.shape[0]))
|
||||
# uv_counter[m2uv[:, 0]] = 1
|
||||
# uv_counter[m2uv[:, 1]] += 1
|
||||
# uv_values /= uv_counter[:, np.newaxis]
|
||||
# x = uv_values
|
||||
# # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
|
||||
# im = ((x - x.min()) / (x.max() - x.min())) * 255
|
||||
|
||||
|
||||
# x = np.transpose(uv_unposed, [1,2,0])
|
||||
# x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
|
||||
# _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
|
||||
# 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 = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# )]).show()
|
||||
|
||||
|
||||
# print()
|
||||
|
||||
|
||||
|
||||
|
||||
# # @tf.function
|
||||
# def euclid_error(mask, gt, pred):
|
||||
# # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1)
|
||||
# return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
|
||||
|
||||
# all_error = []
|
||||
# for batch in dataset:
|
||||
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
# mesh_reconstruct = uv2mesh(batch['uv_posed'])
|
||||
# err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
|
||||
|
||||
# i = 296
|
||||
# err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
|
||||
# colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
|
||||
|
||||
# _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
|
||||
# x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
|
||||
# y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
|
||||
# colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
|
||||
# fig = 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 = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2]
|
||||
# ),
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=y[:,0],
|
||||
# y=y[:,1],
|
||||
# z=y[:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# )])
|
||||
# fig.show()
|
||||
|
||||
# fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
|
||||
# fig.update_traces(marker={'color': colors, 'size': 5})
|
||||
# fig.show()
|
||||
|
||||
# all_error.append(err.numpy().sum())
|
||||
# # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
|
||||
# print(len(all_error))
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
# dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# from losses import collision_loss as collision_penalty
|
||||
# from losses import collision_penalty as collision_penalty_2
|
||||
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
|
||||
|
||||
|
||||
# for d in dataset:
|
||||
# d
|
||||
|
||||
# v = uv2mesh(d[0])
|
||||
# b = d[3]
|
||||
# m = d[4]
|
||||
|
||||
# b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
|
||||
|
||||
# c = collision_penalty(v, b, m)
|
||||
# c2 = collision_penalty(v, b2, m)
|
||||
|
||||
# c_ = collision_penalty_2(v, b, m)
|
||||
# c2_ = collision_penalty_2(v, b2, m)
|
||||
|
||||
# print()
|
||||
# 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()
|
||||
+3
-3
@@ -1,16 +1,16 @@
|
||||
import tables as tb
|
||||
|
||||
class Frame:
|
||||
def __init__(self, id, frame_nr, garment):
|
||||
def __init__(self, id, frame_nr, garment=''):
|
||||
self.id = id
|
||||
self.frame_nr = frame_nr
|
||||
self.garment = garment
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.id, self.frame_nr))
|
||||
return hash((self.id, self.frame_nr, self.garment))
|
||||
|
||||
def __eq__(self, other):
|
||||
return (self.id, self.frame_nr) == (other.id, other.frame_nr)
|
||||
return (self.id, self.frame_nr, self.garment) == (other.id, other.frame_nr, self.garment)
|
||||
|
||||
|
||||
def save_to_h5file(h5file, id, frame_nr, garment, data, name, type):
|
||||
|
||||
Reference in New Issue
Block a user