Merge remote-tracking branch 'origin/post_orig' into review_1_orig

configs/config_abl_overfit_jumpsuit_2.py

@@ -19,7 +19,7 @@ config['only_curvatures'] =           True
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '3'
+config['gpu'] = '5'
 config['name'] = 'abl_overfit_jumpsuit_2'
 config['notes'] = 'remove other conditions from input as well'
 

configs/post/config_abl_pyr_32.py

@@ -1,9 +1,9 @@
 config = dict()
 
-config['gpu'] = '2'
+config['gpu'] = '6'
 config['name'] = 'abl_pyr_32'
-config['notes'] = 'fixed noise handling'
+config['notes'] = ''
-config['tags'] = ['pyr_post']
+config['tags'] = ['review_1']
 
 # DON'T FORGET TO RESET SSIM LOSS USAGE
 # DON'T FORGET TO CHECK MODEL LOADING

curvature_analysis.py

@@ -1,283 +0,0 @@
-import os
-import sys
-os.environ["CUDA_VISIBLE_DEVICES"] = '4'
-
-
-import tqdm
-import argparse
-import random
-import string
-import time
-import numpy as np
-import trimesh as tm
-import cmapy as cm
-import scipy as sp
-import scipy.io as sio
-import plotly.graph_objects as go
-from pyntcloud import PyntCloud
-from scipy.sparse import csr_matrix, coo_matrix
-from scipy.sparse.csgraph import floyd_warshall, dijkstra
-from plotly.subplots import make_subplots
-from custom_utils.custom_utils import _maskFaces
-from postprocessing import fix_collision
-
-from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
-from custom_utils.custom_utils import _maskFaces
-from auto_encoder_curv import dice_dense
-
-def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
-    if body is not None:
-        mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :]
-    else:
-        mesh_fixed = mesh
-
-    if do_smoothing:
-        m = tm.Trimesh(mesh_fixed, faces, process=False, validate=False)
-        outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
-        outliers = np.nonzero(outliers)[0]
-        outliers_n_neighbours = get_closest_n_points(mesh_fixed, outliers, 2)
-        outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
-        faces_fixed = np.concatenate((faces, outliers_faces), axis=0)
-        m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False)
-        tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4)
-
-        mesh = m.vertices
-    else:
-        mesh = mesh_fixed
-
-    return mesh
-
-def export_to_ply(mesh, file):
-    ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
-    output_file = open(file, "wb+")
-    output_file.write(ply)
-    output_file.close()
-
-
-def get_closest_n_points(verts, idx, n):
-    dist = sp.spatial.distance.cdist(verts[idx, :], verts)
-    dist[dist == 0] = 1e9
-    idx_sort = np.argsort(dist, axis=-1)
-    closest = idx_sort[:, :n]
-    return closest
-
-
-def calculate_dist(folder, neighbourhood):
-    out = dict()
-    frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
-
-    for frame in tqdm.tqdm(frames):
-        id = frame.split('_')[0]
-        if id in out.keys():
-            continue
-
-        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
-        k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
-        out[id] = k_neighbors
-
-    return out
-
-
-def histedges_equalN(x, nbin):
-    npt = len(x)
-    return np.interp(np.linspace(0, npt, nbin + 1),
-                     np.arange(npt),
-                     np.sort(x))
-
-
-def get_bins(values, nr_of_bins):
-    flattened = []
-    for c in values:
-        flattened.extend(c.flatten())
-    flattened = np.array(flattened)
-    bins = histedges_equalN(flattened, nr_of_bins)
-    return bins
-
-
-def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
-    # id = frame.split('_')[0]
-    # frame_nr = int(frame.split('_')[1].split('.')[0])
-
-    # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
-    # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
-    # v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
-    
-    m = tm.Trimesh(v, faces, process=False, validate=False)
-
-    outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
-    outliers = np.nonzero(outliers)[0]
-    outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
-    outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
-    faces = np.concatenate((faces, outliers_faces), axis=0)
-
-    verts = v
-
-    m = tm.Trimesh(verts, faces, process=False, validate=False)
-    export_to_ply(m, temp_file)
-
-    cloud = PyntCloud.from_file(temp_file)
-
-    k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
-    #k_neighbors = k_neighbors_cache[id] # np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
-    ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
-    cloud.add_scalar_field("curvature", ev=ev)
-    curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
-    
-    bins = np.load(f'custom_data/train_bins.npy') # get_bins(curvature, nr_of_labels)
-    labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
-
-    frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
-    for j in range(bins.shape[0] - 2):
-        frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
-    frame_labels[curvature >= bins[-2]] = labels[-1]
-       
-    return frame_labels
-
-import tensorflow as tf
-
-def visualize_curvature_labels(folder, results_folder, frame):
-    id = frame.split('_')[0]
-    frame_nr = frame.split('_')[1].split('.')[0]
-
-    mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
-    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
-    verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
-    pred = np.load(os.path.join(results_folder, frame + '.npy'))
-    pred = postprocess(pred, faces)
-
-    curvature_labels_gt = generate_curvatures(verts, faces, mask, 25)
-    curvature_labels_pred = generate_curvatures(pred, faces, mask, 25)
-    # dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
-    dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:]).numpy()
-    colors_binned_all = []
-
-    for curvature_labels in [curvature_labels_gt, curvature_labels_pred]:
-        colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(curvature_labels.max()+1))])
-        colors_binned = np.zeros((curvature_labels.shape[0], 3))
-        for i in range(int(curvature_labels.max())+1):
-            # print(np.count_nonzero(curvature_labels == i))
-            colors_binned[curvature_labels == i, :] = colors[i]
-        colors_binned_all.append(colors_binned)
-
-    fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
-
-    fig.add_trace(go.Mesh3d(
-                # 8 vertices of a cube
-                x=-pred[:,0],
-                y=pred[:,2],
-                z=pred[:,1],
-                # i, j and k give the vertices of triangles
-                i = faces[:,0],
-                j = faces[:,1],
-                k = faces[:,2],
-                vertexcolor=np.ones((verts.shape[0],3))*np.array([255,200,200])[None,:]
-            ),
-            row=1, col=2)
-
-    fig.add_trace(go.Mesh3d(
-                # 8 vertices of a cube
-                x=verts[:,0],
-                y=verts[:,1],
-                z=verts[:,2],
-                # i, j and k give the vertices of triangles
-                i = faces[:,0],
-                j = faces[:,1],
-                k = faces[:,2],
-                vertexcolor=colors_binned_all[0]
-            ),
-            row=1, col=1)
-    
-    fig.add_trace(go.Mesh3d(
-                # 8 vertices of a cube
-                x=-pred[:,0],
-                y=pred[:,2],
-                z=pred[:,1],
-                # i, j and k give the vertices of triangles
-                i = faces[:,0],
-                j = faces[:,1],
-                k = faces[:,2],
-                vertexcolor=colors_binned_all[1],
-                text=f'dice: {dice}'
-            ),
-            row=1, col=3)
-
-    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
-                    scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
-                    scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
-                    # scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
-                    # scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
-                    )
-    helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
-                yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
-                zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
-        
-    fig.update_layout(scene = helper,
-                        scene2 = helper,
-                        scene3 = helper,
-                        # scene4 = helper,
-                        # scene5 = helper,
-                        )
-
-
-    # meshes = [go.Mesh3d(
-    #             x=x[:,0],
-    #             y=x[:,1],
-    #             z=x[:,2],
-    #             i = faces[:,0],
-    #             j = faces[:,1],
-    #             k = faces[:,2],
-    #             vertexcolor=colors_binned
-    #             )]
-
-    # fig = go.Figure(data=meshes)
-    # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
-    # fig.update_layout(
-    #     scene = dict(
-    #         xaxis = dict(nticks=4, range=[-1,1],),
-    #                     yaxis = dict(nticks=4, range=[-1,1],),
-    #                     zaxis = dict(nticks=4, range=[-1,1],),)
-    #                     )
-
-    return fig, dice
-
-import wandb
-
-def init_wandb():
-    run = wandb.init(project='plots', name='review_4', dir='/home/workspace/local-data/temp/')
-    artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('curvature_analysis.py')
-    run.log_artifact(artifact)
-
-if __name__ == '__main__':
-    # init_wandb()
-
-    is_skirt = False
-    skirt_modifier = '_skirt' if is_skirt else ''
-    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
-    # k_neighbours(k=25)
-
-    folder = '/home/workspace/local-data/train'
-    results_folder = '/home/workspace/data/results/post/abl_overfit_jumpsuit_2-'
-
-    frames_file = 'frames_jumpsuit_30k.txt'
-    frames = []
-    with open(os.path.join(folder, frames_file), 'r') as f:
-        lines = f.readlines()
-        lines = [line.strip() for line in lines]
-    frames = lines
-
-    temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
-    if not is_skirt:
-        dist = np.load('custom_data/dijkstra.npy')
-    else:
-        dist = np.load('custom_data/dijkstra_skirt.npy')
-
-    i = 0
-    while True:
-        frame = np.random.choice(frames)
-        # frame = '03265_193'
-        frame = '03102_49'
-        fig, dice = visualize_curvature_labels(folder, results_folder, frame)
-        wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
-        i += 1
-        print(frame)

cvae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_overfit_jumpsuit_2 import config
+from configs.config_abl_all_post import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -41,7 +41,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAE
+from models.cvae_model import CVAENew as CVAE
 
 
 def init_wandb():

evaluate_ae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_overfit_jumpsuit_2 import config
+from configs.config_abl_all_post import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -42,7 +42,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAE
+from models.cvae_model import CVAENew as CVAE
 from custom_utils.supersmpl_io import writeOBJ
 # from unif_metric import chamfer_and_score
 
@@ -104,7 +104,7 @@ if __name__ == '__main__':
 
     iteration = ''
 
-    load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
+    load_model(model, 'checkpoints/post/' + checkpoint_name, iteration=iteration)
     # load_model(model, 'checkpoints/' + 'abl_all_120k')
 
     output_folder = '/home/workspace/data/results/post/' + checkpoint_name + '-' + iteration

preprocessing.py

@@ -232,17 +232,17 @@ def generate_frames_file(folder, ratio=0.1):
     selected = []
 
     for id in tqdm.tqdm(sub_ids):
-        if garment_type_table[id] != 'Dress':
+        # if garment_type_table[id] != 'Dress':
-            continue
+        #     continue
-        if len(selected) > 30000:
+        # if len(selected) > 30000:
-            break
+        #     break
 
         # frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
         frames = [frame for frame in all_frames if frame.startswith(id + '_')]
         frames = [frame.split('/')[-1].split('.')[0] for frame in frames]
         selected.extend(frames)
 
-    with open(os.path.join(folder, 'frames_dress_30k.txt'), 'w') as f:
+    with open(os.path.join(folder, 'frames_skirt_10p.txt'), 'w') as f:
         for i, frame in enumerate(selected):
             f.write(frame + '\n')
 
@@ -293,9 +293,9 @@ def generate_frames_file_filtered_cape():
 
 if __name__ == '__main__':
 
-    src_folder = '/home/workspace/local-data/train_skirt' #cloth3d_skirt
+    src_folder = '/home/workspace/local-data/valid_skirt' #cloth3d_skirt
     # src_folder = r'D:\Work\Dataset\validation_2'
-    dest_folder = '/home/workspace/data/dataset/validation' 
+    dest_folder = '/home/workspace/data/dataset/local-cache/train_skirt' 
 
 
     # create_no_orient_body(src_folder)
@@ -306,5 +306,5 @@ if __name__ == '__main__':
     # verify_poses(src_folder)
     # create_garment_type_table()
     # generate_faces()
-    generate_frames_file(src_folder, ratio=1)
+    generate_frames_file(src_folder, ratio=0.1)
     # generate_frames_file_filtered_cape()