Update app_predict.py

app_predict.py

@@ -21,7 +21,7 @@ model = ['swin',
          'vit',
          'convnext']
 
-def main():
+def main(input_model=None):
     parser = argparse.ArgumentParser(description='Quick demo Image Classification')
     parser.add_argument('--input_dir', default='./test/', type=str, help='Input images root')
     parser.add_argument('--result_dir', default='./result/', type=str, help='Results images root')
@@ -29,7 +29,7 @@ def main():
     parser.add_argument('--model', default='convnext', type=str, help='Classifier')
 
     args = parser.parse_args()
-
+    args.model = input_model
     inp_dir = args.input_dir
     out_dir = args.result_dir
     os.makedirs(out_dir, exist_ok=True)
@@ -290,4 +290,4 @@ def build_ensemble_model(model: dict, pretrained: bool):
     return ensemble_model
     
 if __name__ == '__main__':
-    main()
+    main(model=None)