## @package download # Module caffe2.python.models.download from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals import argparse import os import sys import signal import re import json from caffe2.proto import caffe2_pb2 # Import urllib try: import urllib.error as urlliberror import urllib.request as urllib HTTPError = urlliberror.HTTPError URLError = urlliberror.URLError except ImportError: import urllib2 as urllib HTTPError = urllib.HTTPError URLError = urllib.URLError # urllib requires more work to deal with a redirect, so not using vanity url DOWNLOAD_BASE_URL = "https://s3.amazonaws.com/download.caffe2.ai/models/" DOWNLOAD_COLUMNS = 70 # Don't let urllib hang up on big downloads def signalHandler(signal, frame): print("Killing download...") exit(0) signal.signal(signal.SIGINT, signalHandler) def deleteDirectory(top_dir): for root, dirs, files in os.walk(top_dir, topdown=False): for name in files: os.remove(os.path.join(root, name)) for name in dirs: os.rmdir(os.path.join(root, name)) os.rmdir(top_dir) def progressBar(percentage): full = int(DOWNLOAD_COLUMNS * percentage / 100) bar = full * "#" + (DOWNLOAD_COLUMNS - full) * " " sys.stdout.write(u"\u001b[1000D[" + bar + "] " + str(percentage) + "%") sys.stdout.flush() def downloadFromURLToFile(url, filename, show_progress=True): try: print("Downloading from {url}".format(url=url)) response = urllib.urlopen(url) size = int(response.info().get('Content-Length').strip()) chunk = min(size, 8192) print("Writing to {filename}".format(filename=filename)) if show_progress: downloaded_size = 0 progressBar(0) with open(filename, "wb") as local_file: while True: data_chunk = response.read(chunk) if not data_chunk: break local_file.write(data_chunk) if show_progress: downloaded_size += len(data_chunk) progressBar(int(100 * downloaded_size / size)) print("") # New line to fix for progress bar except HTTPError as e: raise Exception("Could not download model. [HTTP Error] {code}: {reason}." .format(code=e.code, reason=e.reason)) except URLError as e: raise Exception("Could not download model. [URL Error] {reason}." .format(reason=e.reason)) except Exception as e: raise e def getURLFromName(name, filename): return "{base_url}{name}/{filename}".format(base_url=DOWNLOAD_BASE_URL, name=name, filename=filename) def downloadModel(model, args): # Figure out where to store the model model_folder = '{folder}'.format(folder=model) dir_path = os.path.dirname(os.path.realpath(__file__)) if args.install: model_folder = '{dir_path}/{folder}'.format(dir_path=dir_path, folder=model) # Check if that folder is already there if os.path.exists(model_folder) and not os.path.isdir(model_folder): if not args.force: raise Exception("Cannot create folder for storing the model,\ there exists a file of the same name.") else: print("Overwriting existing file! ({filename})" .format(filename=model_folder)) os.remove(model_folder) if os.path.isdir(model_folder): if not args.force: response = "" query = "Model already exists, continue? [y/N] " try: response = raw_input(query) except NameError: response = input(query) if response.upper() == 'N' or not response: print("Cancelling download...") exit(0) print("Overwriting existing folder! ({filename})".format(filename=model_folder)) deleteDirectory(model_folder) # Now we can safely create the folder and download the model os.makedirs(model_folder) for f in ['predict_net.pb', 'init_net.pb']: try: downloadFromURLToFile(getURLFromName(model, f), '{folder}/{f}'.format(folder=model_folder, f=f)) except Exception as e: print("Abort: {reason}".format(reason=str(e))) print("Cleaning up...") deleteDirectory(model_folder) exit(0) if args.install: os.symlink("{folder}/__sym_init__.py".format(folder=dir_path), "{folder}/__init__.py".format(folder=model_folder)) def validModelName(name): invalid_names = ['__init__'] if name in invalid_names: return False if not re.match("^[/0-9a-zA-Z_-]+$", name): return False return True class ModelDownloader: def _model_dir(self, model): caffe2_home = os.path.expanduser(os.getenv('CAFFE2_HOME', '~/.caffe2')) models_dir = os.getenv('CAFFE2_MODELS', os.path.join(caffe2_home, 'models')) return os.path.join(models_dir, model) def _download(self, model): model_dir = self._model_dir(model) assert not os.path.exists(model_dir) os.makedirs(model_dir) for f in ['predict_net.pb', 'init_net.pb', 'value_info.json']: url = getURLFromName(model, f) dest = os.path.join(model_dir, f) try: try: downloadFromURLToFile(url, dest, show_progress=False) except TypeError: # show_progress not supported prior to # Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1 # (Sep 17, 2017) downloadFromURLToFile(url, dest) except Exception as e: print("Abort: {reason}".format(reason=e)) print("Cleaning up...") deleteDirectory(model_dir) exit(1) def get_c2_model(self, model_name): model_dir = self._model_dir(model_name) if not os.path.exists(model_dir): self._download(model_name) c2_predict_pb = os.path.join(model_dir, 'predict_net.pb') c2_predict_net = caffe2_pb2.NetDef() with open(c2_predict_pb, 'rb') as f: c2_predict_net.ParseFromString(f.read()) c2_predict_net.name = model_name c2_init_pb = os.path.join(model_dir, 'init_net.pb') c2_init_net = caffe2_pb2.NetDef() with open(c2_init_pb, 'rb') as f: c2_init_net.ParseFromString(f.read()) c2_init_net.name = model_name + '_init' value_info = json.load(open(os.path.join(model_dir, 'value_info.json'))) return c2_init_net, c2_predict_net, value_info if __name__ == "__main__": parser = argparse.ArgumentParser( description='Download or install pretrained models.') parser.add_argument('model', nargs='+', help='Model to download/install.') parser.add_argument('-i', '--install', action='store_true', help='Install the model.') parser.add_argument('-f', '--force', action='store_true', help='Force a download/installation.') args = parser.parse_args() for model in args.model: if validModelName(model): downloadModel(model, args) else: print("'{}' is not a valid model name.".format(model))