未验证 提交 949ac1cc 编写于 作者: K Kaipeng Deng 提交者: GitHub

check md5 after weights download (#1311)

* check md5 after weights download
上级 7d873006
...@@ -22,6 +22,8 @@ import shutil ...@@ -22,6 +22,8 @@ import shutil
import requests import requests
import tqdm import tqdm
import hashlib import hashlib
import binascii
import base64
import tarfile import tarfile
import zipfile import zipfile
...@@ -304,9 +306,20 @@ def _download(url, path, md5sum=None): ...@@ -304,9 +306,20 @@ def _download(url, path, md5sum=None):
for chunk in req.iter_content(chunk_size=1024): for chunk in req.iter_content(chunk_size=1024):
if chunk: if chunk:
f.write(chunk) f.write(chunk)
shutil.move(tmp_fullname, fullname)
# check md5 after download in Content-MD5 in req.headers
content_md5 = req.headers.get('content-md5')
if not content_md5 or _md5check(
tmp_fullname,
binascii.hexlify(base64.b64decode(content_md5.strip(
'"'))).decode()):
shutil.move(tmp_fullname, fullname)
return fullname return fullname
else:
logger.warn(
"Download from url imcomplete, try downloading again...")
os.remove(tmp_fullname)
continue
def _md5check(fullname, md5sum=None): def _md5check(fullname, md5sum=None):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册