电子发烧友App
前言
代码仓库地址：https://github.com/Oneflow-Inc/one-yolov5欢迎star one-yolov5项目 获取最新的动态。如果您有问题，欢迎在仓库给我们提出宝贵的意见。如果对您有帮助，欢迎来给我Star呀~??
源码解读：utils/augmentations.py
这个文件主要是负责从github/googleleaps/google drive 等网站或者云服务器上下载所需的一些文件。由于微信会吃掉一些超链接影响阅读，欢迎大家查看原始文档网站解读文章：https://start.oneflow.org/oneflow-yolo-doc/source_code_interpretation/utils/downloads_py.html
是一个工具类，代码比较简单，函数也比较少，主要难点还是在于一些包可能大家不是很熟悉，下面一起来学习下。
这个文件比较重要的是两个函数：safe_download和attempt_download。在train.py或者yolo.py等文件中都会用到。
1. 导入需要的包
""" Download?utils """ import?os???#?与操作系统进行交互的模块 import?platform??#?提供获取操作系统相关信息的模块 import?shutil?#?Python的高阶文件操作模块 import?subprocess??#?子进程定义及操作的模块 import?time?#?时间模块 import?urllib?#?用于操作网页?url?并对网页的内容进行抓取处理??如urllib.parse:?解析url from?pathlib?import?Path?#?Path将str转换为Path对象?使字符串路径易于操作的模块 from?zipfile?import?ZipFile?#?导入文件解压模块 import?oneflow?as?flow?#?导入深度学习框架oneflow包 import?requests???#?通过urllib3实现自动发送HTTP/1.1请求的第三方模块
2. gsutil_getsize
这个函数是用来返回网站链接 url 对应文件的大小。
def?gsutil_getsize(url=""):
????"""用在downloads.py的print_mutation函数当中?计算某个url对应的文件大小
????用于返回网站链接url对应文件的大小，注意单位是bytes
????gs://bucket/file?size?https://cloud.google.com/storage/docs/gsutil/commands/du
????"""
????#?创建一个子进程在命令行执行?gsutil?du?url?命令(访问?Cloud?Storage)?返回执行结果(文件)
????#?gs://bucket/file?size?https://cloud.google.com/storage/docs/gsutil/commands/du
????s?=?subprocess.check_output(f"gsutil?du?{url}",?shell=True).decode("utf-8")
????return?eval(s.split("?")[0])?if?len(s)?else?0??#?bytes
3. safe_download、attempt_download
这两个函数主要是用来从 github 或者 googleleaps 云服务器中下载文件的，主要是下载权重文件。
one-yolov5 仓库中 attempt_download 函数调用 safe_download 函数。
3.1 safe_download
这个函数是用来下载 url（github） 或者 url2（谷歌云服务器） 网页路径对应的文件，
通常是下载权重文件，经常用在 attempt_download 函数中，代码如下：
def?safe_download(file,?url,?url2=None,?min_bytes=1e0,?error_msg=""):
????"""经常用在?attempt_download?函数中，也可以单独使用
????下载?url/url2?网页路径对应的文件
????Attempts?to?download?file?from?url?or?url2,?checks?and?removes?incomplete?downloads??min_bytes,?assert_msg??#?check
????except?Exception?as?e:??#?url2??url1?不行就尝试从?url2?中下载文件??一般是googleleaps(云服务器)
????????#?移除之前下载失败的不完整文件
????????file.unlink(missing_ok=True)??#?remove?partial?downloads
????????print(f"ERROR:?{e}
Re-attempting?{url2?or?url}?to?{file}...")
????????os.system(f"curl?-L?'{url2?or?url}'?-o?'{file}'?--retry?3?-C?-")??#?curl?download,?retry?and?resume?on?fail
????finally:
????????#?检查文件是否下载下来了?或?文件大小是否小于min_bytes
????????if?not?file.exists()?or?file.stat().st_size?
url?=?"https://github.com/Oneflow-Inc/one-yolov5/releases/download/v1.0.0/model_comparison.png"
safe_download("op.png",?url)
Downloading https://github.com/Oneflow-Inc/one-yolov5/releases/download/v1.0.0/model_comparison.png to op.png...
0%| | 0.00/118k [00:00
from?PIL?import?Image
display(Image.open("op.png"))?#?显示下载的图片
image
3.2 attempt_download
这个函数是实现从几个云平台 (github/googleleaps云服务器/xxx) 下载文件(在one-yolov5中一般是预训练模型)，
会调用上面的 safe_download 函数。会用在 experimental.py 中的 attempt_load 函数和 train.py 中，都是用来下载预训练权重。代码详解如下：
def?attempt_download(file,?repo="Oneflow-Inc/one-yolov5"):??#?from?utils.downloads?import?*;?attempt_download()
????"""用在attempt_download函数中
????下载?url/url2?网页路径对应的文件
????Attempts?to?download?file?from?url?or?url2,?checks?and?removes?incomplete?downloads??:/
????????????#?name:?要下载的文件名
????????????file?=?name.split("?")[0]??#?parse?authentication?https://url.com/file.txt?auth...
????????????#?如果文件已经在本地存在了就不用下载了
????????????if?Path(file).is_file():
????????????????print(f"Found?{url}?locally?at?{file}")??#?file?already?exists
????????????else:
????????????????safe_download(file=file,?url=url,?min_bytes=1e5)??#?下载文件
????????????return?file
????????#?GitHub?assets
????????file.parent.mkdir(parents=True,?exist_ok=True)??#?make?parent?dir?(if?required)
????????try:
????????????#?利用?github?api?获取最新的版本相关信息??这里的?response?是一个字典
????????????response?=?requests.get(f"https://api.github.com/repos/{repo}/releases/latest").json()??#?github?api
????????????assets?=?[x["name"]?for?x?in?response["assets"]]??#?release?assets,?i.e.?['yolov5s',?'yolov5m',?...]
????????????tag?=?response["tag_name"]??#?i.e.?'v1.0'
????????except:??#?fallback?plan?获取失败?就退而求其次?直接利用?git?命令强行补齐版本信息
????????????assets?=?[
????????????????"yolov5n.zip",
????????????????"yolov5s.zip",
????????????????"yolov5m.zip",
????????????????"yolov5l.zip",
????????????????"yolov5x.zip",
????????????????"yolov5n6.zip",
????????????????"yolov5s6.zip",
????????????????"yolov5m6.zip",
????????????????"yolov5l6.zip",
????????????????"yolov5x6.zip",
????????????]
????????????try:?#?创建一个子进程在命令行执行?git?tag?命令(返回版本号?版本号信息一般在字典最后一个?-1)?返回执行结果(版本号?tag)
????????????????tag?=?subprocess.check_output("git?tag",?shell=True,?stderr=subprocess.STDOUT).decode().split()[-1]
????????????except:
????????????????#?如果还是失败?就强行自己补一个版本号?tag='v1.1'?，比如这里在?one-yolov5?中直接补当前的最新版本?v1.1.
????????????????tag?=?"v1.1"??#?current?release
????????if?".zip"?not?in?name:
????????????name?=?name?+?".zip"
????????file?=?Path(name)
????????if?name?in?assets:
????????????safe_download(
????????????????file,
????????????????url=f"https://github.com/{repo}/releases/download/{tag}/{name}",
????????????????#?url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}',??#?backup?url?(optional)
????????????????min_bytes=1e5,
????????????????error_msg=f"{file}?missing,?try?downloading?from?https://github.com/{repo}/releases/",
????????????)
????????if?".zip"?in?name:
????????????new_dir?=?Path(name[:-4])
????????else:
????????????new_dir?=?Path(name)
????????if?not?os.path.exists(new_dir):??#?判断文件夹是否存在
????????????os.mkdir(new_dir)??#?新建文件夹
????????if?".zip"?in?name:
????????????print("unzipping...?",?end="")
????????????#?ZipFile(new_file).extractall(path=file.parent)??#?unzip
????????????f?=?ZipFile(file)
????????????f.extractall(new_dir)
????????????os.remove(file)??#?remove?zip
????????????tmp_dir?=?"/tmp/oneyolov5"
????????????if?os.path.isdir(tmp_dir):
????????????????shutil.rmtree(tmp_dir)
????????????if?".zip"?in?name:
????????????????path1?=?os.path.join(name[:-4],?name[:-4])
????????????else:
????????????????path1?=?os.path.join(name,?name)
????????????shutil.copytree(path1,?tmp_dir)
????????????shutil.rmtree(new_dir)
????????????shutil.copytree(tmp_dir,?new_dir)
????????????shutil.rmtree(tmp_dir)
????return?str(file)
attempt_download("yolov5n")
Downloading https://github.com/Oneflow-Inc/one-yolov5/releases/download/v1.0.0/yolov5n.zip to yolov5n.zip...
0%| | 0.00/3.53M [00:00
4. get_token & gdrive_download（没使用）
这两个函数是实现从 google drive 上下载压缩文件并将其解压, 再删除掉压缩文件。但是这好像并没有在代码中使用，所以这两个函数可以随便了解下就好，主要还是要掌握上面的两个下载函数用的比较多。
4.1 get_token
这个函数实现从 cookie中 获取令牌 token 。会在 gdrive_download 中被调用。
get_token函数代码：
def?get_token(cookie="./cookie"):
????"""在gdrive_download中使用
????实现从cookie中获取令牌token?
????"""
????with?open(cookie)?as?f:
????????for?line?in?f:
????????????if?"download"?in?line:
????????????????return?line.split()[-1]
????return?""
4.2 gdrive_download
这个函数实现从 google drive 上下载压缩文件并将其解压, 再删除掉压缩文件。这个函数貌似没用到，随便看下就好。
gdrive_download函数代码：
def?gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m',?file='tmp.zip'):
????"""
????实现从?google?drive?上下载压缩文件并将其解压,?再删除掉压缩文件
????:params?id:?url的?后面的?id?参数的参数值
????:params?file:?需要下载的压缩文件名
????"""
????t?=?time.time()??#?获取当前时间
????file?=?Path(file)???#?Path将str转换为Path对象
????cookie?=?Path('cookie')??#?gdrive?cookie
????print(f'Downloading?https://drive.google.com/uc?export=download&id={id}?as?{file}...?',?end='')
????file.unlink(missing_ok=True)??#?移除已经存在的文件(可能是下载失败/下载不完整的文件)
????cookie.unlink(missing_ok=True)??#?移除已经存在的cookie
????#?尝试下载压缩文件
????out?=?"NUL"?if?platform.system()?==?"Windows"?else?"/dev/null"
????#?使用?cmd?命令从?google?drive?上下载文件
????os.system(f'curl?-c?./cookie?-s?-L?"drive.google.com/uc?export=download&id={id}"?>?{out}')
????if?os.path.exists('cookie'):
????????#?如果文件较大?就需要有令牌?get_token?(存在?cookie?才有令牌)的指令?s?才能下载
????????#?get_token()?函数在上面定义了，用于获取当前?cookie?的令牌?token
????????s?=?f'curl?-Lb?./cookie?"drive.google.com/uc?export=download&confirm={get_token()}&id={id}"?-o?{file}'
????else:
????????#?小文件就不需要带令牌的指令?s?直接下载就行
????????s?=?f'curl?-s?-L?-o?{file}?"drive.google.com/uc?export=download&id={id}"'
????#?执行下载指令?s?并获得返回值?如果?cmd?命令执行成功?则?os.system()命令会返回?0
????r?=?os.system(s)
????cookie.unlink(missing_ok=True)??#?再次移除已经存在的cookie
????#?下载错误检测??如果?r?!=?0?则下载错误
????if?r?!=?0:
????????file.unlink(missing_ok=True)??#?下载错误?移除下载的文件(可能不完整或者下载失败)
????????print('Download?error?')??#?raise?Exception('Download?error')
????????return?r
????#?如果是压缩文件?就解压??file.suffix?方法可以获取?file?文件的后缀
????if?file.suffix?==?'.zip':
????????print('unzipping...?',?end='')
????????os.system(f'unzip?-q?{file}')??#?cmd命令执行解压命令
????????file.unlink()??#?移除?.zip?压缩文件
????print(f'Done?({time.time()?-?t:.1f}s)')??#?打印下载?+?解压过程所需要的时间
????return?r
总结
这个文件的代码比较少，真正有用的函数也比较少。
也就是safe_download和attempt_download两个函数比较重要，大家重点掌握这两个函数即可。
Reference
【YOLOV5-5.x 源码解读】google_utils.py
编辑：黄飞
?
工商网监
评论