Added choices to arguments

PrimozGodec · Aug 2, 2017 · 895e046 · 895e046
1 parent 4c0bf8d
commit 895e046
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/image_colorization/main.py b/src/image_colorization/main.py
@@ -3,7 +3,12 @@
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Test model')
-    parser.add_argument('--model', type=str, default="reg_full_model", help="Choose name of the model you want to test")
+    parser.add_argument('--model',
+                        type=str,
+                        default="reg_full_model",
+                        help="Choose name of the model you want to test",
+                        choices={"reg_full_model", "reg_full_vgg_model", "reg_part_model",
+                                 "class_weights_model", "class_wo_weights_model"})
 
     args = parser.parse_args()