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

train encoder on raw curvature values

Browse Source
This commit is contained in:
luni
2023-01-27 11:13:10 +00:00
parent a3c5b526b8
commit 6e6b0a6604
8 changed files with 60 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aaa_helpers.txt
+4
View File
@@ -1,3 +1,7 @@
python merge_data.py -fi '/home/workspace/data/dataset/train' -fo '/home/workspace/local-data/train' -s 0 -e 1500
nohup python auto_encoder_curv.py --name="lhunor_curv_raw_r2" > nohup/review/curv_raw_r2.out 2>&1 &
one final convolutional layer for smoothing, on the final coordinates one final convolutional layer for smoothing, on the final coordinates
skirts!!!! skirts!!!!
auto_encoder_curv.py
+22 -29
View File
@@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
testing = False testing = False
def init_wandb(): def init_wandb():
if not testing: if not testing:
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/') run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/', tags=config['tags'])
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py') artifact.add_file('auto_encoder_curv.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
@@ -47,7 +47,9 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py') artifact.add_file('convnext/models/convnext_tf.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
for config_file_name in config_file_names: for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py') file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact) run.log_artifact(artifact)
uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
@@ -100,22 +102,18 @@ def train_step(batch):
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
pred = model([batch['uv_curvature'], batch['uv_mask']], training=True) pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) ce_uv_loss = tf.reduce_mean(tf.abs(batch['uv_curvature'] - pred))
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred, batch['uv_curvature'], max_val=1))
mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask'] ce_mesh_loss = tf.reduce_mean(tf.abs(batch['mesh_curvature'] - mesh_pred[:,:,0]))
ce_mesh_loss = scce(mesh_orig, mesh_pred)
loss = ce_uv_loss + ce_mesh_loss loss = ce_uv_loss + ce_mesh_loss
mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, batch['uv_curvature'])
gradients = tape.gradient(loss, model.trainable_variables) gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables)) optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, return { 'train/loss': loss, 'train/uv_l1': ce_uv_loss, 'train/ssim': ms_ssim,
'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss } 'train/mesh_l1': ce_mesh_loss }
@@ -123,35 +121,30 @@ def train_step(batch):
def valid_step(batch): def valid_step(batch):
pred = model([batch['uv_curvature'], batch['uv_mask']], training=False) pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) ce_uv_loss = tf.reduce_mean(tf.abs(batch['uv_curvature'] - pred))
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred, batch['uv_curvature'], max_val=1))
mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred = uv2mesh(pred) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask'] ce_mesh_loss = tf.reduce_mean(tf.abs(batch['mesh_curvature'] - mesh_pred[:,:,0]))
ce_mesh_loss = scce(mesh_orig, mesh_pred)
loss = ce_uv_loss + ce_mesh_loss loss = ce_uv_loss + ce_mesh_loss
mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1)) return { 'valid/loss': loss, 'valid/uv_l1': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mesh_l1': ce_mesh_loss }
dice = dice_dense(pred, batch['uv_curvature'])
return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou,
'valid/dice': dice, 'valid/ce_mesh_loss': ce_mesh_loss }
def plot_everything(batch: dict, type: str): def plot_everything(batch: dict, type: str):
pred = model([batch['uv_curvature'], batch['uv_mask']], training=False) pred = model([batch['uv_curvature'], batch['uv_mask']], training=False)
orig_img = tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.dtypes.int32) orig_img = batch['uv_curvature']
pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32) pred_img = pred
img = tf.concat((orig_img, pred_img), axis=2).numpy() img = tf.concat((orig_img, pred_img), axis=2).numpy()
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)]) # colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, 6+1)])
color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8) # color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8)
for i in range(7): # for i in range(7):
color_img[img == i, :] = colors[i] # color_img[img == i, :] = colors[i]
color_img = [wandb.Image(color_img[i]) for i in range(4)] # color_img = [wandb.Image(color_img[i]) for i in range(4)]
imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])], type + '/color_images': color_img } imgs = { type + '/images': [wandb.Image((img[i,:,:,np.newaxis]-img[i,:,:,np.newaxis].min()) / (img[i,:,:,np.newaxis].max()-img[i,:,:,np.newaxis].min())) for i in range(img.shape[0])], } # type + '/color_images': color_img
return imgs return imgs
@@ -206,7 +199,7 @@ def test():
if __name__ == '__main__': if __name__ == '__main__':
init_wandb() init_wandb()
allow_memory_growth() # allow_memory_growth()
parser = argparse.ArgumentParser(description='') parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str) parser.add_argument('--name', type=str)
args = vars(parser.parse_args()) args = vars(parser.parse_args())
configs/config_ae_curv_post.py
+5 -9
View File
@@ -31,22 +31,18 @@ config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
config['cond_ae']['activation'] = 'tanh' config['cond_ae']['activation'] = 'tanh'
config['gpu'] = '3' config['gpu'] = '1'
config['lr'] = 0.001 config['lr'] = 0.001
config['name'] = 'curv_post_1' config['name'] = 'curv_raw_r2'
config['notes'] = 'convnext + convnext, | \ config['notes'] = ''
\ config['tags'] = ['review_2']
| \
deeper model | \
| \
'
config['cond_ae']['upscale_output'] = False config['cond_ae']['upscale_output'] = False
config['pretrained_model'] = '' config['pretrained_model'] = ''
config['signature_dynamic'] = { 'id': (), 'frame_nr': (), config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
'uv_curvature': (512, 256, 7), 'uv_curvature': (512, 256, 1),
'mesh_curvature': (config['smpl_res']) 'mesh_curvature': (config['smpl_res'])
} }
config['signature_static'] = { 'mesh_mask': (config['smpl_res'],), config['signature_static'] = { 'mesh_mask': (config['smpl_res'],),
convnext/models/convnext_tf.py
+3 -3
View File
@@ -338,7 +338,7 @@ class ConvNeXtReverseLabels(tf.keras.Model):
""" """
def __init__(self, def __init__(self,
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
nr_of_classes=7): nr_of_classes=7):
super().__init__() super().__init__()
@@ -360,14 +360,14 @@ class ConvNeXtReverseLabels(tf.keras.Model):
cur = 0 cur = 0
for i in range(len(depths)): for i in range(len(depths)):
stage = tf.keras.Sequential( stage = tf.keras.Sequential(
[Block(dim=dims[i], drop_path=dp_rates[cur + j], [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
) )
self.stages.append(stage) self.stages.append(stage)
cur += depths[i] cur += depths[i]
# self.last = Dense(nr_of_classes, activation='softmax') # self.last = Dense(nr_of_classes, activation='softmax')
self.last = Dense(nr_of_classes - 1) self.last = Dense(1, activation='sigmoid')
def call(self, x, training): def call(self, x, training):
for i in reversed(range(len(self.downsample_layers))): for i in reversed(range(len(self.downsample_layers))):
curvature_labels.py
+6 -5
View File
@@ -275,6 +275,7 @@ def read_raw_curvatures(folder):
continue continue
curv = np.load(garment_dir + '/curvature_25_raw.npy') curv = np.load(garment_dir + '/curvature_25_raw.npy')
print(curv.max())
curvatures_all.append(curv) curvatures_all.append(curv)
return curvatures_all return curvatures_all
@@ -282,24 +283,24 @@ def read_raw_curvatures(folder):
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/test') parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/train')
parser.add_argument("-s", "--start", help = "start of range", default=0) parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=1500) parser.add_argument("-e", "--end", help = "end of range", default=1500)
parser.add_argument("-n", "--name", help = "end of range", default="lhunor") parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
args = parser.parse_args() args = parser.parse_args()
F = np.load("FeatureLines/superSMPL/smpl_faces.npy") # F = np.load("FeatureLines/superSMPL/smpl_faces.npy")
# k_neighbours(k=25) # k_neighbours(k=25)
folder = args.folder folder = args.folder
nr_of_labels = 6 nr_of_labels = 6
neighbourhood = 25 neighbourhood = 25
temp_file = 'temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply' temp_file = 'temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
dist = np.load('MeshUtils/dijsktra.npy') # dist = np.load('MeshUtils/dijsktra.npy')
# dist = np.load('MeshUtils/dijkstra_skirt.npy') # dist = np.load('MeshUtils/dijkstra_skirt.npy')
generate_curvatures(folder, neighbourhood) # generate_curvatures(folder, neighbourhood)
# curvatures = read_raw_curvatures(folder) curvatures = read_raw_curvatures(folder)
# bins = get_bins(curvatures, nr_of_labels) # bins = get_bins(curvatures, nr_of_labels)
# np.save('train_bins_skirt.npy', bins) # np.save('train_bins_skirt.npy', bins)
# bins = np.load('train_bins.npy') # bins = np.load('train_bins.npy')
dataset.py
+13 -14
View File
@@ -176,7 +176,7 @@ class DatasetFactory():
elif key == 'frame_nr': elif key == 'frame_nr':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'mesh_curvature': elif key == 'mesh_curvature':
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
else: else:
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
return tf_signature return tf_signature
@@ -260,8 +260,8 @@ class DatasetFactory():
if 'mesh_curvature' in signature.keys(): 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'] mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature_raw', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature_raw']
item['mesh_curvature'] = mesh_curvature.astype(np.int32) item['mesh_curvature'] = mesh_curvature
for res in resolutions: for res in resolutions:
@@ -270,13 +270,13 @@ class DatasetFactory():
res_string = '_' + str(res) res_string = '_' + str(res)
if 'uv_curvature' + res_string in signature.keys(): 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')); uv_curvature = np.load(os.path.join(folder, 'curvature_uv_raw' + res_string, id + separator + frame_nr + '.npz'));
if res_string != '': if res_string != '':
uv_curvature = uv_curvature['uv_curvature'] uv_curvature = uv_curvature['uv_curvature']
else: else:
uv_curvature = uv_curvature['curvature_uv'] uv_curvature = uv_curvature['curvature_uv_raw']
uv_curvature = tf.one_hot(uv_curvature, 7) # uv_curvature = tf.one_hot(uv_curvature, 7)
item['uv_curvature' + res_string] = uv_curvature item['uv_curvature' + res_string] = uv_curvature[:,:,tf.newaxis]
if 'uv_unposed' + res_string in signature.keys(): if 'uv_unposed' + res_string in signature.keys():
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
@@ -370,8 +370,8 @@ if __name__ == '__main__':
# 'mesh_unposed': (14475, 3,), # 'mesh_unposed': (14475, 3,),
'uv_unposed': (512, 256, 3), 'uv_unposed': (512, 256, 3),
# 'uv_unposed_32': (3, 32, 16,), # 'uv_unposed_32': (3, 32, 16,),
'uv_curvature': (512, 256, 7), 'uv_curvature': (512, 256, 1),
# 'mesh_curvature': (14475) 'mesh_curvature': (14475)
} }
signature_static = { #'uv_static': (3, 512, 256,), signature_static = { #'uv_static': (3, 512, 256,),
# 'uv_static_32': (3, 32, 16,), # 'uv_static_32': (3, 32, 16,),
@@ -381,15 +381,14 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,), # 'uv_shaped_body': (3, 512, 256,),
} }
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt', dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
batch_size=1, batch_size=1,
signature=signature_dynamic, signature=signature_dynamic,
resolutions=[512], resolutions=[512],
do_repeat=False, do_repeat=False,
use_one_hot=False, use_one_hot=False,
is_cape=False, is_cape=False,)
frames_file='frames_dress_30k.txt') static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
signature=signature_static, signature=signature_static,
resolutions=[512], resolutions=[512],
is_cape=False) is_cape=False)
@@ -404,7 +403,7 @@ if __name__ == '__main__':
for i, batch in enumerate(dataset): for i, batch in enumerate(dataset):
print(i) print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) 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_curvature'][0,:,:,0].numpy()).show()
sum_uv += batch['uv_unposed'].numpy().sum(axis=0) sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
sum_mask += batch['uv_mask'].numpy().sum(axis=0) sum_mask += batch['uv_mask'].numpy().sum(axis=0)
merge_data.py
+6 -4
View File
@@ -179,7 +179,7 @@ if __name__ == '__main__':
start = time.perf_counter() start = time.perf_counter()
# offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] # offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3]
# body_static = sio.loadmat(os.path.join(garment_dir, 'info_body.mat'))['B'] # body_static = sio.loadmat(os.path.join(garment_dir, 'info_body.mat'))['B']
# mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool)) mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))
# verts_static = (body_static + offsets_static) * mask[:, np.newaxis] # verts_static = (body_static + offsets_static) * mask[:, np.newaxis]
# # dfassert verts_static.min() > -1 and verts_static.max() < 1 # # dfassert verts_static.min() > -1 and verts_static.max() < 1
@@ -257,10 +257,12 @@ if __name__ == '__main__':
# for i in range(body_posed.shape[0]): # for i in range(body_posed.shape[0]):
# shutil.copyfile(os.path.join(garment_dir, 'curvature_uv', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_curvature', id + '_' + str(i) + '.npz')) # shutil.copyfile(os.path.join(garment_dir, 'curvature_uv', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_curvature', id + '_' + str(i) + '.npz'))
os.makedirs(os.path.join(output_folder, 'mesh_curvature'), exist_ok=True) os.makedirs(os.path.join(output_folder, 'mesh_curvature_raw'), exist_ok=True)
mesh_curvature = np.load(os.path.join(garment_dir, 'curvature_label_6_25.npy')) data = np.load(os.path.join(garment_dir, 'curvature_25_raw.npy'))
mesh_curvature_raw = np.zeros((data.shape[0], 14475), dtype=np.float32)
mesh_curvature_raw[:, mask] = data
for i in range(pose.shape[1]): for i in range(pose.shape[1]):
np.savez_compressed(os.path.join(output_folder, 'mesh_curvature', id + '_' + str(i) + '.npz'), mesh_curvature=mesh_curvature[i]) np.savez_compressed(os.path.join(output_folder, 'mesh_curvature_raw', id + '_' + str(i) + '.npz'), mesh_curvature_raw=mesh_curvature_raw[i])
# os.makedirs(os.path.join(output_folder, 'uv_posed_body_no_orient'), exist_ok=True) # os.makedirs(os.path.join(output_folder, 'uv_posed_body_no_orient'), exist_ok=True)
models/base.py
+1 -7
View File
@@ -190,13 +190,7 @@ class DecoderNew(tf.keras.Model):
x = self.decoder(x, training=training) x = self.decoder(x, training=training)
if self.is_label: # if self.is_label:
x = x * mask[:,:,:,tf.newaxis]
background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
x = tf.concat((background, x), axis=-1)
x = tf.math.softmax(x)
if len(inputs) == 2: if len(inputs) == 2:
if len(mask.shape) == 2: if len(mask.shape) == 2: