added choices for domain

line · Oct 7, 2024 · e2cc675 · e2cc675
1 parent 5e4de26
commit e2cc675
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/training/train.py b/training/train.py
@@ -293,7 +293,10 @@ def check_valid_combination(dataset, feature, domain):
                         help='feature name. select from [resnet_glove, clip, clip_slowfast, clip_slowfast_pann, i3d_clip, clap].'
                              'NOTE: i3d_clip and clip_slowfast_pann are only for TVSum and QVHighlight, respectively.')
     parser.add_argument('--resume', '-r', type=str, help='specify model path for fine-tuning. If None, train the model from scratch.')
-    parser.add_argument('--domain', '-dm', type=str, help='domain for highlight detection dataset (e.g., BK for TVSum, dog for YouTube Highlight).')
+    parser.add_argument('--domain', '-dm', type=str,
+                        choices=['BK', 'BT', 'DS', 'FM', 'GA', 'MS', 'PK', 'PR', 'VT', 'VU',
+                                 'dog', 'gymnastics', 'parkour', 'skating', 'skiing', 'surfing'],
+                        help='domain for highlight detection dataset (e.g., BK for TVSum, dog for YouTube Highlight).')
     args = parser.parse_args()
 
     is_valid = check_valid_combination(args.dataset, args.feature, args.domain)