Update TEED/main.py

TEED/main.py

@@ -206,7 +206,7 @@ def testPich(checkpoint_path, dataloader, model, device, output_dir, args, resiz
     print("Average time per image: %f.4" % total_duration.mean(), "seconds")
     print("Time spend in the Dataset: %f.4" % total_duration.sum(), "seconds")
 
-def parse_args(is_testing=True, pl_opt_dir = 'output/teed'):
+def parse_args(is_testing=True, pl_opt_dir='output/teed'):
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(description='TEED model')
     parser.add_argument('--choose_test_data',
@@ -214,16 +214,19 @@ def parse_args(is_testing=True, pl_opt_dir = 'output/teed'):
                         default=-1,     # UDED=15
                         help='Choose a dataset for testing: 0 - 15')
 
+    # 新增的 epoch 参数
+    parser.add_argument('--epoch', type=int, required=True, help='Epoch number')
+
     # ----------- test -------0--
     TEST_DATA = DATASET_NAMES[parser.parse_args().choose_test_data] # max 8
     test_inf = dataset_info(TEST_DATA, is_linux=IS_LINUX)
 
     # Training settings
-    # BIPED-B2=1, BIPDE-B3=2, just for evaluation, using LDC trained with 2 or 3 bloacks
     TRAIN_DATA = DATASET_NAMES[0] # BIPED=0, BRIND=6, MDBD=10, BIPBRI=13
     train_inf = dataset_info(TRAIN_DATA, is_linux=IS_LINUX)
     train_dir = train_inf['data_dir']
 
+
     # Data parameters
     parser.add_argument('--input_dir',
                         type=str,
@@ -366,7 +369,7 @@ def main(args, train_inf):
     training_dir = os.path.join(args.output_dir,args.train_data)
     os.makedirs(training_dir,exist_ok=True)
     checkpoint_path = './TEED/checkpoints/BIPED/5/5_model.pth'
-    checkpoint_path = os.path.join('./TEED/checkpoints', 'BIPED', str(args.epoch), '5_model.pth')
+    checkpoint_path = os.path.join('./TEED/checkpoints', 'BIPED', str(args.epochs), '5_model.pth')
     if args.tensorboard and not args.is_testing:
         # from tensorboardX import SummaryWriter  # previous torch version
         from torch.utils.tensorboard import SummaryWriter # for torch 1.4 or greather