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
skirts!!!!
auto_encoder_curv.py
+22 -29
View File
@@ -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())
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['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'],),
convnext/models/convnext_tf.py
+3 -3
View File
@@ -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))):
curvature_labels.py
+6 -5
View File
@@ -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')
dataset.py
+13 -14
View File
@@ -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)
merge_data.py
+6 -4
View File
@@ -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)
models/base.py
+1 -7
View File
@@ -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: