diff --git a/paddlehub/commands/download.py b/paddlehub/commands/download.py index 490701d6cb94f9f407ceaf882596aa5a845634b6..8a06144f237985140b211aa160fa7016d00d96ef 100644 --- a/paddlehub/commands/download.py +++ b/paddlehub/commands/download.py @@ -31,40 +31,44 @@ class DownloadCommand(BaseCommand): def __init__(self, name): super(DownloadCommand, self).__init__(name) self.show_in_help = True - self.description = "Download PaddlePaddle pretrained model files." + self.description = "Download PaddlePaddle pretrained model/module files." self.parser = self.parser = argparse.ArgumentParser( description=self.__class__.__doc__, - prog='%s %s ' % (ENTRY, name), + prog='%s %s ' % (ENTRY, name), usage='%(prog)s [options]', add_help=False) # yapf: disable - self.add_arg('--output_path', str, ".", "path to save the model" ) + self.add_arg("--type", str, "All", "choice: Module/Model/All") + self.add_arg('--output_path', str, ".", "path to save the model/module" ) self.add_arg('--uncompress', bool, False, "uncompress the download package or not" ) # yapf: enable def exec(self, argv): if not argv: - print("ERROR: Please provide the model name\n") + print("ERROR: Please provide the model/module name\n") self.help() return False - model_name = argv[0] - model_version = None if "==" not in model_name else model_name.split( - "==")[1] - model_name = model_name if "==" not in model_name else model_name.split( - "==")[0] + mod_name = argv[0] + mod_version = None if "==" not in mod_name else mod_name.split("==")[1] + mod_name = mod_name if "==" not in mod_name else mod_name.split("==")[0] self.args = self.parser.parse_args(argv[1:]) - if not self.args.output_path: - self.args.output_path = "." - utils.check_path(self.args.output_path) + self.args.type = self.check_type(self.args.type) - search_result = default_hub_server.get_model_url( - model_name, version=model_version) + if self.args.type in ["Module", "Model"]: + search_result = default_hub_server.get_resource_url( + mod_name, resource_type=self.args.type, version=mod_version) + else: + search_result = default_hub_server.get_resource_url( + mod_name, resource_type="Module", version=mod_version) + if search_result == {}: + search_result = default_hub_server.get_resource_url( + mod_name, resource_type="Model", version=mod_version) url = search_result.get('url', None) except_md5_value = search_result.get('md5', None) if not url: - tips = "Can't found model %s" % model_name + tips = "Can't found model/module %s" % mod_name if model_version: - tips += " with version %s" % model_version + tips += " with version %s" % mod_version print(tips) return True @@ -93,10 +97,20 @@ class DownloadCommand(BaseCommand): result, tips, file = default_downloader.uncompress( file=file, dirname=self.args.output_path, - delete_file=True, + delete_file=False, print_progress=True) print(tips) return True + def check_type(self, mod_type): + mod_type = mod_type.lower() + if mod_type == "module": + mod_type = "Module" + elif mod_type == "model": + mod_type = "Model" + else: + mod_type = "All" + return mod_type + command = DownloadCommand.instance() diff --git a/paddlehub/commands/run.py b/paddlehub/commands/run.py index 3075ded7360f076b25abb834f6f0836892307bca..70df85f74d8cd3e80e86f3a9feb32f2348795b83 100644 --- a/paddlehub/commands/run.py +++ b/paddlehub/commands/run.py @@ -18,6 +18,7 @@ from __future__ import print_function import argparse import os +import sys from paddlehub.commands.base_command import BaseCommand, ENTRY from paddlehub.io.parser import yaml_parser, txt_parser @@ -93,7 +94,14 @@ class RunCommand(BaseCommand): if not result: return False - module = hub.Module(module_dir=module_dir) + try: + module = hub.Module(module_dir=module_dir) + except: + print( + "ERROR! %s is a model. The command is only for the module type but not the model type." + % module_name) + sys.exit(0) + self.parse_args_with_module(module, argv[1:]) if not module.default_signature: