making ready for skirts

auto_encoder_static.py

@@ -12,13 +12,18 @@ from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
 import plotly.graph_objects as go
 import plotly.express as ex
 
+if config['is_skirt']:
+    from losses import edge_loss_skirt as edge_loss
+else:
+    from losses import edge_loss
+
 from tf_utils import check_tf_version, allow_memory_growth
-from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict, plot_3d
 from models.cond_encoder import AutoEncoder
 
 
+
 def init_wandb():
     run = wandb.init(project='vae', name=config['name'],
             notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
@@ -35,7 +40,7 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256)
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
 
 
 
@@ -93,7 +98,7 @@ def plot_everything(batch: dict, type: str):
     img = tf.concat((orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
-    F = np.load("custom_utils/smpl_faces.npy")
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
     garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
     gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
@@ -101,7 +106,7 @@ def plot_everything(batch: dict, type: str):
     body = gt
     garment_posed = mesh_pred
     collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
-    collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
     collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
 
     figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
@@ -162,6 +167,7 @@ if __name__ == '__main__':
     resolution_modifier = ''
     if config['resolution'] != 512:
         resolution_modifier = '_' + str(config['resolution'])
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     factory = DatasetFactory()