From fca4d9eec289f22c081daa2c61a1110e3f268f92 Mon Sep 17 00:00:00 2001 From: Tom Hochstein Date: Tue, 1 Sep 2020 14:57:07 -0500 Subject: [PATCH] Add smaller version of download_models.py Signed-off-by: Tom Hochstein --- testdata/dnn/download_models_basic.py | 159 ++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 testdata/dnn/download_models_basic.py diff --git a/testdata/dnn/download_models_basic.py b/testdata/dnn/download_models_basic.py new file mode 100644 index 0000000..5c8a616 --- /dev/null +++ b/testdata/dnn/download_models_basic.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python + +from __future__ import print_function +import hashlib +import sys +import tarfile +if sys.version_info[0] < 3: + from urllib2 import urlopen +else: + from urllib.request import urlopen + + +class Model: + MB = 1024*1024 + BUFSIZE = 10*MB + + def __init__(self, **kwargs): + self.name = kwargs.pop('name') + self.url = kwargs.pop('url', None) + self.filename = kwargs.pop('filename') + self.sha = kwargs.pop('sha', None) + self.archive = kwargs.pop('archive', None) + self.member = kwargs.pop('member', None) + + def __str__(self): + return 'Model <{}>'.format(self.name) + + def printRequest(self, r): + def getMB(r): + d = dict(r.info()) + for c in ['content-length', 'Content-Length']: + if c in d: + return int(d[c]) / self.MB + return '' + print(' {} {} [{} Mb]'.format(r.getcode(), r.msg, getMB(r))) + + def verify(self): + if not self.sha: + return False + print(' expect {}'.format(self.sha)) + sha = hashlib.sha1() + try: + with open(self.filename, 'rb') as f: + while True: + buf = f.read(self.BUFSIZE) + if not buf: + break + sha.update(buf) + print(' actual {}'.format(sha.hexdigest())) + return self.sha == sha.hexdigest() + except Exception as e: + print(' catch {}'.format(e)) + + def get(self): + if self.verify(): + print(' hash match - skipping') + return True + + if self.archive or self.member: + assert(self.archive and self.member) + print(' hash check failed - extracting') + print(' get {}'.format(self.member)) + self.extract() + else: + assert(self.url) + print(' hash check failed - downloading') + print(' get {}'.format(self.url)) + self.download() + + print(' done') + print(' file {}'.format(self.filename)) + return self.verify() + + def download(self): + try: + r = urlopen(self.url, timeout=60) + self.printRequest(r) + self.save(r) + except Exception as e: + print(' catch {}'.format(e)) + + def extract(self): + try: + with tarfile.open(self.archive) as f: + assert self.member in f.getnames() + self.save(f.extractfile(self.member)) + except Exception as e: + print(' catch {}'.format(e)) + + def save(self, r): + with open(self.filename, 'wb') as f: + print(' progress ', end='') + sys.stdout.flush() + while True: + buf = r.read(self.BUFSIZE) + if not buf: + break + f.write(buf) + print('>', end='') + sys.stdout.flush() + +models = [ + Model( + name='Fcn', + url='http://dl.caffe.berkeleyvision.org/fcn8s-heavy-pascal.caffemodel', + sha='c449ea74dd7d83751d1357d6a8c323fcf4038962', + filename='fcn8s-heavy-pascal.caffemodel'), + Model( + name='SqueezeNet_v1.1', + url='https://raw.githubusercontent.com/DeepScale/SqueezeNet/b5c3f1a23713c8b3fd7b801d229f6b04c64374a5/SqueezeNet_v1.1/squeezenet_v1.1.caffemodel', + sha='3397f026368a45ae236403ccc81cfcbe8ebe1bd0', + filename='squeezenet_v1.1.caffemodel'), + Model( + name='Colorization', + url='https://raw.githubusercontent.com/richzhang/colorization/master/models/colorization_deploy_v2.prototxt', + sha='f528334e386a69cbaaf237a7611d833bef8e5219', + filename='colorization_deploy_v2.prototxt'), + Model( + name='Colorization', + url='http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel', + sha='21e61293a3fa6747308171c11b6dd18a68a26e7f', + filename='colorization_release_v2.caffemodel'), + Model( + name='OpenPose/pose/coco', # https://github.com/CMU-Perceptual-Computing-Lab/openpose + url='http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/coco/pose_iter_440000.caffemodel', + sha='ac7e97da66f3ab8169af2e601384c144e23a95c1', + filename='openpose_pose_coco.caffemodel'), + Model( + name='YOLOv3', # https://pjreddie.com/darknet/yolo/ + url='https://pjreddie.com/media/files/yolov3.weights', + sha='520878f12e97cf820529daea502acca380f1cb8e', + filename='yolov3.weights'), + Model( + name='EAST', # https://github.com/argman/EAST (a TensorFlow model), https://arxiv.org/abs/1704.03155v2 (a paper) + url='https://www.dropbox.com/s/r2ingd0l3zt8hxs/frozen_east_text_detection.tar.gz?dl=1', + sha='3ca8233d6edd748f7ed23246c8ca24cbf696bb94', + filename='frozen_east_text_detection.tar.gz'), + Model( + name='EAST', + archive='frozen_east_text_detection.tar.gz', + member='frozen_east_text_detection.pb', + sha='fffabf5ac36f37bddf68e34e84b45f5c4247ed06', + filename='frozen_east_text_detection.pb'), +] + +# Note: models will be downloaded to current working directory +# expected working directory is opencv_extra/testdata/dnn +if __name__ == '__main__': + failedModels = [] + for m in models: + print(m) + if not m.get(): + failedModels.append(m.filename) + + if failedModels: + print("Following models have not been downloaded:") + for f in failedModels: + print("* {}".format(f)) + exit(15) -- 2.17.1