YOLOv5解析之downloads.py 代碼示例

前言

代碼倉庫地址：https://github.com/Oneflow-Inc/one-yolov5歡迎star one-yolov5項目 獲取最新的動態(tài)。如果您有問題，歡迎在倉庫給我們提出寶貴的意見。如果對您有幫助，歡迎來給我Star呀~??

源碼解讀：utils/augmentations.py

這個文件主要是負責從github/googleleaps/google drive 等網(wǎng)站或者云服務器上下載所需的一些文件。由于微信會吃掉一些超鏈接影響閱讀，歡迎大家查看原始文檔網(wǎng)站解讀文章：https://start.oneflow.org/oneflow-yolo-doc/source_code_interpretation/utils/downloads_py.html

是一個工具類，代碼比較簡單，函數(shù)也比較少，主要難點還是在于一些包可能大家不是很熟悉，下面一起來學習下。

這個文件比較重要的是兩個函數(shù)：safe_download和attempt_download。在train.py或者yolo.py等文件中都會用到。

1. 導入需要的包

"""
Download?utils
"""
import?os???#?與操作系統(tǒng)進行交互的模塊
import?platform??#?提供獲取操作系統(tǒng)相關信息的模塊
import?shutil?#?Python的高階文件操作模塊
import?subprocess??#?子進程定義及操作的模塊
import?time?#?時間模塊
import?urllib?#?用于操作網(wǎng)頁?url?并對網(wǎng)頁的內(nèi)容進行抓取處理??如urllib.parse:?解析url

from?pathlib?import?Path?#?Path將str轉換為Path對象?使字符串路徑易于操作的模塊
from?zipfile?import?ZipFile?#?導入文件解壓模塊

import?oneflow?as?flow?#?導入深度學習框架oneflow包
import?requests???#?通過urllib3實現(xiàn)自動發(fā)送HTTP/1.1請求的第三方模塊

2. gsutil_getsize

這個函數(shù)是用來返回網(wǎng)站鏈接 url 對應文件的大小。

def?gsutil_getsize(url=""):
????"""用在downloads.py的print_mutation函數(shù)當中?計算某個url對應的文件大小
????用于返回網(wǎng)站鏈接url對應文件的大小，注意單位是bytes
????gs://bucket/file?size?https://cloud.google.com/storage/docs/gsutil/commands/du
????"""
????#?創(chuàng)建一個子進程在命令行執(zhí)行?gsutil?du?url?命令(訪問?Cloud?Storage)?返回執(zhí)行結果(文件)
????#?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

這兩個函數(shù)主要是用來從 github 或者 googleleaps 云服務器中下載文件的，主要是下載權重文件。

one-yolov5 倉庫中 attempt_download 函數(shù)調(diào)用 safe_download 函數(shù)。

3.1 safe_download

這個函數(shù)是用來下載 url（github） 或者 url2（谷歌云服務器） 網(wǎng)頁路徑對應的文件，

通常是下載權重文件，經(jīng)常用在 attempt_download 函數(shù)中，代碼如下：

def?safe_download(file,?url,?url2=None,?min_bytes=1e0,?error_msg=""):

????"""經(jīng)常用在?attempt_download?函數(shù)中，也可以單獨使用
????下載?url/url2?網(wǎng)頁路徑對應的文件
????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

	這個函數(shù)是實現(xiàn)從幾個云平臺 (github/googleleaps云服務器/xxx) 下載文件(在one-yolov5中一般是預訓練模型)，

	會調(diào)用上面的 safe_download 函數(shù)。會用在 experimental.py 中的 attempt_load 函數(shù)和 train.py 中，都是用來下載預訓練權重。代碼詳解如下：

	def?attempt_download(file,?repo="Oneflow-Inc/one-yolov5"):??#?from?utils.downloads?import?*;?attempt_download()
????"""用在attempt_download函數(shù)中
????下載?url/url2?網(wǎng)頁路徑對應的文件
????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...
????????????#?如果文件已經(jīng)在本地存在了就不用下載了
????????????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:?#?創(chuàng)建一個子進程在命令行執(zhí)行?git?tag?命令(返回版本號?版本號信息一般在字典最后一個?-1)?返回執(zhí)行結果(版本號?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（沒使用）

	這兩個函數(shù)是實現(xiàn)從 google drive 上下載壓縮文件并將其解壓, 再刪除掉壓縮文件。但是這好像并沒有在代碼中使用，所以這兩個函數(shù)可以隨便了解下就好，主要還是要掌握上面的兩個下載函數(shù)用的比較多。

	4.1 get_token

	這個函數(shù)實現(xiàn)從 cookie中 獲取令牌 token 。會在 gdrive_download 中被調(diào)用。

	get_token函數(shù)代碼：

	def?get_token(cookie="./cookie"):
????"""在gdrive_download中使用
????實現(xiàn)從cookie中獲取令牌token?
????"""
????with?open(cookie)?as?f:
????????for?line?in?f:
????????????if?"download"?in?line:
????????????????return?line.split()[-1]
????return?""

	4.2 gdrive_download

	這個函數(shù)實現(xiàn)從 google drive 上下載壓縮文件并將其解壓, 再刪除掉壓縮文件。這個函數(shù)貌似沒用到，隨便看下就好。

	gdrive_download函數(shù)代碼：
def?gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m',?file='tmp.zip'):
????"""
????實現(xiàn)從?google?drive?上下載壓縮文件并將其解壓,?再刪除掉壓縮文件
????:params?id:?url的?后面的?id?參數(shù)的參數(shù)值
????: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)??#?移除已經(jīng)存在的文件(可能是下載失敗/下載不完整的文件)
????cookie.unlink(missing_ok=True)??#?移除已經(jīng)存在的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()?函數(shù)在上面定義了，用于獲取當前?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}"'
????#?執(zhí)行下載指令?s?并獲得返回值?如果?cmd?命令執(zhí)行成功?則?os.system()命令會返回?0
????r?=?os.system(s)
????cookie.unlink(missing_ok=True)??#?再次移除已經(jīng)存在的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命令執(zhí)行解壓命令
????????file.unlink()??#?移除?.zip?壓縮文件

????print(f'Done?({time.time()?-?t:.1f}s)')??#?打印下載?+?解壓過程所需要的時間
????return?r


	總結

	這個文件的代碼比較少，真正有用的函數(shù)也比較少。

	也就是safe_download和attempt_download兩個函數(shù)比較重要，大家重點掌握這兩個函數(shù)即可。

	Reference

	【YOLOV5-5.x 源碼解讀】google_utils.py

	編輯：黃飛

	?