diff --git a/aaa_helpers.txt b/aaa_helpers.txt index f8ea6ae..b276923 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -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 skirts!!!! diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index 2756698..692b0ed 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder testing = False def init_wandb(): 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.add_file('auto_encoder_curv.py') artifact.add_file('losses.py') @@ -47,7 +47,9 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] 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) 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: pred = model([batch['uv_curvature'], batch['uv_mask']], training=True) - ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) - ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) + ce_uv_loss = tf.reduce_mean(tf.abs(batch['uv_curvature'] - pred)) + 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_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask'] - ce_mesh_loss = scce(mesh_orig, mesh_pred) + ce_mesh_loss = tf.reduce_mean(tf.abs(batch['mesh_curvature'] - mesh_pred[:,:,0])) 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) 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, - 'train/dice': dice, 'train/ce_mesh_loss': ce_mesh_loss } + return { 'train/loss': loss, 'train/uv_l1': ce_uv_loss, 'train/ssim': ms_ssim, + 'train/mesh_l1': ce_mesh_loss } @@ -123,35 +121,30 @@ def train_step(batch): def valid_step(batch): pred = model([batch['uv_curvature'], batch['uv_mask']], training=False) - ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) - ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) + ce_uv_loss = tf.reduce_mean(tf.abs(batch['uv_curvature'] - pred)) + 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_orig = tf.cast(uv2mesh(tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.int32)), tf.float32) * batch['mesh_mask'] - ce_mesh_loss = scce(mesh_orig, mesh_pred) + ce_mesh_loss = tf.reduce_mean(tf.abs(batch['mesh_curvature'] - mesh_pred[:,:,0])) 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']) - - 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 } + return { 'valid/loss': loss, 'valid/uv_l1': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mesh_l1': ce_mesh_loss } def plot_everything(batch: dict, type: str): 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) - pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32) + orig_img = batch['uv_curvature'] + pred_img = pred 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)]) - color_img = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), np.uint8) - for i in range(7): - color_img[img == i, :] = colors[i] - color_img = [wandb.Image(color_img[i]) for i in range(4)] + # 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) + # for i in range(7): + # color_img[img == i, :] = colors[i] + # 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 @@ -206,7 +199,7 @@ def test(): if __name__ == '__main__': init_wandb() - allow_memory_growth() + # allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) diff --git a/configs/config_ae_curv_post.py b/configs/config_ae_curv_post.py index 9b6c971..8ddf722 100644 --- a/configs/config_ae_curv_post.py +++ b/configs/config_ae_curv_post.py @@ -31,22 +31,18 @@ config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] config['cond_ae']['activation'] = 'tanh' -config['gpu'] = '3' +config['gpu'] = '1' config['lr'] = 0.001 -config['name'] = 'curv_post_1' -config['notes'] = 'convnext + convnext, | \ - \ -| \ - deeper model | \ -| \ - ' +config['name'] = 'curv_raw_r2' +config['notes'] = '' +config['tags'] = ['review_2'] config['cond_ae']['upscale_output'] = False config['pretrained_model'] = '' config['signature_dynamic'] = { 'id': (), 'frame_nr': (), - 'uv_curvature': (512, 256, 7), + 'uv_curvature': (512, 256, 1), 'mesh_curvature': (config['smpl_res']) } config['signature_static'] = { 'mesh_mask': (config['smpl_res'],), diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py index d570c30..71212ce 100644 --- a/convnext/models/convnext_tf.py +++ b/convnext/models/convnext_tf.py @@ -338,7 +338,7 @@ class ConvNeXtReverseLabels(tf.keras.Model): """ 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., nr_of_classes=7): super().__init__() @@ -360,14 +360,14 @@ class ConvNeXtReverseLabels(tf.keras.Model): cur = 0 for i in range(len(depths)): 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])] ) self.stages.append(stage) cur += depths[i] # 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): for i in reversed(range(len(self.downsample_layers))): diff --git a/curvature_labels.py b/curvature_labels.py index f99b262..f67d90e 100644 --- a/curvature_labels.py +++ b/curvature_labels.py @@ -275,6 +275,7 @@ def read_raw_curvatures(folder): continue curv = np.load(garment_dir + '/curvature_25_raw.npy') + print(curv.max()) curvatures_all.append(curv) return curvatures_all @@ -282,24 +283,24 @@ def read_raw_curvatures(folder): if __name__ == '__main__': 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("-e", "--end", help = "end of range", default=1500) parser.add_argument("-n", "--name", help = "end of range", default="lhunor") args = parser.parse_args() - F = np.load("FeatureLines/superSMPL/smpl_faces.npy") + # F = np.load("FeatureLines/superSMPL/smpl_faces.npy") # k_neighbours(k=25) folder = args.folder nr_of_labels = 6 neighbourhood = 25 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') - generate_curvatures(folder, neighbourhood) - # curvatures = read_raw_curvatures(folder) + # generate_curvatures(folder, neighbourhood) + curvatures = read_raw_curvatures(folder) # bins = get_bins(curvatures, nr_of_labels) # np.save('train_bins_skirt.npy', bins) # bins = np.load('train_bins.npy') diff --git a/dataset.py b/dataset.py index d2ee42d..469da74 100644 --- a/dataset.py +++ b/dataset.py @@ -176,7 +176,7 @@ class DatasetFactory(): elif key == 'frame_nr': tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) 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: tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) return tf_signature @@ -260,8 +260,8 @@ class DatasetFactory(): 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) + 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 for res in resolutions: @@ -270,13 +270,13 @@ class DatasetFactory(): 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')); + uv_curvature = np.load(os.path.join(folder, 'curvature_uv_raw' + 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 + uv_curvature = uv_curvature['curvature_uv_raw'] + # uv_curvature = tf.one_hot(uv_curvature, 7) + item['uv_curvature' + res_string] = uv_curvature[:,:,tf.newaxis] 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'] @@ -370,8 +370,8 @@ if __name__ == '__main__': # 'mesh_unposed': (14475, 3,), 'uv_unposed': (512, 256, 3), # 'uv_unposed_32': (3, 32, 16,), - 'uv_curvature': (512, 256, 7), - # 'mesh_curvature': (14475) + 'uv_curvature': (512, 256, 1), + 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), # 'uv_static_32': (3, 32, 16,), @@ -381,15 +381,14 @@ if __name__ == '__main__': # '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, 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', + is_cape=False,) + static_dataset = StaticDataset(folder='/home/workspace/local-data/train', signature=signature_static, resolutions=[512], is_cape=False) @@ -404,7 +403,7 @@ if __name__ == '__main__': 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() + ex.imshow(batch['uv_curvature'][0,:,:,0].numpy()).show() sum_uv += batch['uv_unposed'].numpy().sum(axis=0) sum_mask += batch['uv_mask'].numpy().sum(axis=0) diff --git a/merge_data.py b/merge_data.py index 28122a2..f04a52c 100644 --- a/merge_data.py +++ b/merge_data.py @@ -179,7 +179,7 @@ if __name__ == '__main__': start = time.perf_counter() # 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'] - # 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] # # 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]): # 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) - mesh_curvature = np.load(os.path.join(garment_dir, 'curvature_label_6_25.npy')) + os.makedirs(os.path.join(output_folder, 'mesh_curvature_raw'), exist_ok=True) + 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]): - 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) diff --git a/models/base.py b/models/base.py index e9cfd67..ba2fd7d 100644 --- a/models/base.py +++ b/models/base.py @@ -190,13 +190,7 @@ class DecoderNew(tf.keras.Model): x = self.decoder(x, training=training) - 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 self.is_label: if len(inputs) == 2: if len(mask.shape) == 2: