Caffe2 - Python API
A deep learning, cross platform ML framework
download.py
1 # Copyright (c) 2016-present, Facebook, Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 ##############################################################################
15 
16 ## @package download
17 # Module caffe2.python.models.download
18 from __future__ import absolute_import
19 from __future__ import division
20 from __future__ import print_function
21 from __future__ import unicode_literals
22 import argparse
23 import os
24 import sys
25 import signal
26 import re
27 
28 # Import urllib
29 try:
30  import urllib.error as urlliberror
31  import urllib.request as urllib
32  HTTPError = urlliberror.HTTPError
33  URLError = urlliberror.URLError
34 except ImportError:
35  import urllib2 as urllib
36  HTTPError = urllib.HTTPError
37  URLError = urllib.URLError
38 
39 # urllib requires more work to deal with a redirect, so not using vanity url
40 DOWNLOAD_BASE_URL = "https://s3.amazonaws.com/download.caffe2.ai/models/"
41 DOWNLOAD_COLUMNS = 70
42 
43 
44 # Don't let urllib hang up on big downloads
45 def signalHandler(signal, frame):
46  print("Killing download...")
47  exit(0)
48 
49 
50 signal.signal(signal.SIGINT, signalHandler)
51 
52 
53 def deleteDirectory(top_dir):
54  for root, dirs, files in os.walk(top_dir, topdown=False):
55  for name in files:
56  os.remove(os.path.join(root, name))
57  for name in dirs:
58  os.rmdir(os.path.join(root, name))
59  os.rmdir(top_dir)
60 
61 
62 def progressBar(percentage):
63  full = int(DOWNLOAD_COLUMNS * percentage / 100)
64  bar = full * "#" + (DOWNLOAD_COLUMNS - full) * " "
65  sys.stdout.write(u"\u001b[1000D[" + bar + "] " + str(percentage) + "%")
66  sys.stdout.flush()
67 
68 
69 def downloadFromURLToFile(url, filename, show_progress=True):
70  try:
71  print("Downloading from {url}".format(url=url))
72  response = urllib.urlopen(url)
73  size = int(response.info().get('Content-Length').strip())
74  chunk = min(size, 8192)
75  print("Writing to {filename}".format(filename=filename))
76  if show_progress:
77  downloaded_size = 0
78  progressBar(0)
79  with open(filename, "wb") as local_file:
80  while True:
81  data_chunk = response.read(chunk)
82  if not data_chunk:
83  break
84  local_file.write(data_chunk)
85  if show_progress:
86  downloaded_size += len(data_chunk)
87  progressBar(int(100 * downloaded_size / size))
88  print("") # New line to fix for progress bar
89  except HTTPError as e:
90  raise Exception("Could not download model. [HTTP Error] {code}: {reason}."
91  .format(code=e.code, reason=e.reason))
92  except URLError as e:
93  raise Exception("Could not download model. [URL Error] {reason}."
94  .format(reason=e.reason))
95  except Exception as e:
96  raise e
97 
98 
99 def getURLFromName(name, filename):
100  return "{base_url}{name}/{filename}".format(base_url=DOWNLOAD_BASE_URL,
101  name=name, filename=filename)
102 
103 
104 def downloadModel(model, args):
105  # Figure out where to store the model
106  model_folder = '{folder}'.format(folder=model)
107  dir_path = os.path.dirname(os.path.realpath(__file__))
108  if args.install:
109  model_folder = '{dir_path}/{folder}'.format(dir_path=dir_path,
110  folder=model)
111 
112  # Check if that folder is already there
113  if os.path.exists(model_folder) and not os.path.isdir(model_folder):
114  if not args.force:
115  raise Exception("Cannot create folder for storing the model,\
116  there exists a file of the same name.")
117  else:
118  print("Overwriting existing file! ({filename})"
119  .format(filename=model_folder))
120  os.remove(model_folder)
121  if os.path.isdir(model_folder):
122  if not args.force:
123  response = ""
124  query = "Model already exists, continue? [y/N] "
125  try:
126  response = raw_input(query)
127  except NameError:
128  response = input(query)
129  if response.upper() == 'N' or not response:
130  print("Cancelling download...")
131  exit(0)
132  print("Overwriting existing folder! ({filename})".format(filename=model_folder))
133  deleteDirectory(model_folder)
134 
135  # Now we can safely create the folder and download the model
136  os.makedirs(model_folder)
137  for f in ['predict_net.pb', 'init_net.pb']:
138  try:
139  downloadFromURLToFile(getURLFromName(model, f),
140  '{folder}/{f}'.format(folder=model_folder,
141  f=f))
142  except Exception as e:
143  print("Abort: {reason}".format(reason=str(e)))
144  print("Cleaning up...")
145  deleteDirectory(model_folder)
146  exit(0)
147 
148  if args.install:
149  os.symlink("{folder}/__sym_init__.py".format(folder=dir_path),
150  "{folder}/__init__.py".format(folder=model_folder))
151 
152 
153 def validModelName(name):
154  invalid_names = ['__init__']
155  if name in invalid_names:
156  return False
157  if not re.match("^[/0-9a-zA-Z_]+$", name):
158  return False
159  return True
160 
161 
162 if __name__ == "__main__":
163  parser = argparse.ArgumentParser(
164  description='Download or install pretrained models.')
165  parser.add_argument('model', nargs='+',
166  help='Model to download/install.')
167  parser.add_argument('-i', '--install', action='store_true',
168  help='Install the model.')
169  parser.add_argument('-f', '--force', action='store_true',
170  help='Force a download/installation.')
171  args = parser.parse_args()
172  for model in args.model:
173  if validModelName(model):
174  downloadModel(model, args)
175  else:
176  print("'{}' is not a valid model name.".format(model))