Add DRF Scoliosis1K pipeline and optional wandb logging

opengait/main.py

@@ -24,13 +24,27 @@ opt = parser.parse_args()
 def initialization(cfgs, training):
     msg_mgr = get_msg_mgr()
     engine_cfg = cfgs['trainer_cfg'] if training else cfgs['evaluator_cfg']
+    logger_cfg = cfgs.get('logger_cfg', {})
     output_path = os.path.join('output/', cfgs['data_cfg']['dataset_name'],
                                cfgs['model_cfg']['model'], engine_cfg['save_name'])
     if training:
-        msg_mgr.init_manager(output_path, opt.log_to_file, engine_cfg['log_iter'],
-                             engine_cfg['restore_hint'] if isinstance(engine_cfg['restore_hint'], (int)) else 0)
+        msg_mgr.init_manager(
+            output_path,
+            opt.log_to_file,
+            engine_cfg['log_iter'],
+            engine_cfg['restore_hint'] if isinstance(engine_cfg['restore_hint'], (int)) else 0,
+            logger_cfg=logger_cfg,
+            config=cfgs,
+            phase='train',
+        )
     else:
-        msg_mgr.init_logger(output_path, opt.log_to_file)
+        msg_mgr.init_logger(
+            output_path,
+            opt.log_to_file,
+            logger_cfg=logger_cfg,
+            config=cfgs,
+            phase='test',
+        )
 
     msg_mgr.log_info(engine_cfg)
 
@@ -55,7 +69,10 @@ def run_model(cfgs, training):
     if training:
         Model.run_train(model)
     else:
-        Model.run_test(model)
+        result_dict = Model.run_test(model)
+        if result_dict:
+            msg_mgr.write_to_tensorboard(result_dict)
+            msg_mgr.write_to_wandb(result_dict)
 
 
 if __name__ == '__main__':