configs save

scripts/ae_body_train.py

@@ -11,10 +11,12 @@ parser.add_argument('--config', type=str)
 parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
-config_module = importlib.import_module(args['config'])
-config = config_module.config
 if args['debug']:
     from configs.iccvw.ae_body_cape_m import config
+else:
+    config_module = importlib.import_module(args['config'])
+    config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -32,10 +34,7 @@ import plotly.graph_objects as go
 import plotly.express as ex
 import torch.optim as optim
 
-if config['is_skirt']:
-    from losses.losses import edge_loss_skirt as edge_loss
-else:
-    from losses.losses import edge_loss
+from losses.losses import edge_loss
 
 from data.dataset import Dataset
 # from torch.utils.data import DataLoader
@@ -65,7 +64,7 @@ def init_wandb():
             artifact.add_file(file_name)
     run.log_artifact(artifact)
 
-uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+uv2mesh = UV2Mesh(res=256, is_skirt=False, is_smpl_orig=config['is_cape'])
 
 
 
@@ -122,7 +121,7 @@ def plot_everything(batch: dict, type: str):
     img = torch.concat((orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
 
-    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    F = np.load("smpl/smpl_faces.npy")
     garment_unposed, faces = mesh_pred.cpu().numpy(), F
     gt = [mesh_orig[i].cpu().numpy() for i in range(pred.shape[0])]
     
@@ -176,9 +175,7 @@ if __name__ == '__main__':
     if not args['debug']:
         init_wandb()
 
-    skirt_modifier = '_skirt' if config['is_skirt'] else ''
-
-    uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
+    uv_mask_body = np.load('custom_utils/mask_256x512.npy')
     uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
     uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
 

scripts/ae_static_train.py

@@ -94,7 +94,7 @@ def train_step(batch):
         mesh_orig = batch['mesh_static']
         l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
         e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
     
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }