使用Google云平臺(tái)實(shí)戰(zhàn)基于PyTorch的yolo-v3模型
對(duì)于計(jì)算機(jī)視覺愛好者來說,YOLO (You Only Look Once)是一個(gè)非常流行的實(shí)時(shí)目標(biāo)檢測(cè)算法,因?yàn)樗浅?欤瑫r(shí)性能非常好。在本文中,我將共享一個(gè)視頻處理的代碼,以獲取視頻中每個(gè)對(duì)象目標(biāo)的邊框。我們將不討論YOLO的概念或架構(gòu),因?yàn)楹芏嗪玫奈恼乱呀?jīng)在媒體中詳細(xì)闡述了這些知識(shí)點(diǎn)。這里我們只討論函數(shù)代碼。
開始
谷歌Colab地址:https://colab.research.google.com/github/vindruid/yolov3-in-colab/blob/master/yolov3_video.ipynb。yolo的git倉庫:https://github.com/ultralytics/yolov3。盡管倉庫已經(jīng)包含如何使用YOLOv3的教程,教程只需要運(yùn)行python detect.py --source file.mp4,但是我簡化了代碼,具體在谷歌Colab / Jupyter筆記本中。準(zhǔn)備YoloV3和LoadModel首先克隆YoloV3倉庫,然后導(dǎo)入通用包和repo函數(shù)!git clone https://github.com/ultralytics/yolov3import timeimport globimport torchimport os
import argparsefrom sys import platform%cd yolov3from models import *from utils.datasets import *from utils.utils import *
from IPython.display import HTMLfrom base64 import b64encode設(shè)置參數(shù)解析器,初始化設(shè)備(CPU / CUDA),初始化YOLO模型,然后加載權(quán)重。parser = argparse.ArgumentParser()parser.a(chǎn)dd_argument('--cfg', type=str, default='cfg/yolov3-spp.cfg', help='*.cfg path')parser.a(chǎn)dd_argument('--names', type=str, default='data/coco.names', help='*.names path')parser.a(chǎn)dd_argument('--weights', type=str, default='weights/yolov3-spp-ultralytics.pt', help='weights path')parser.a(chǎn)dd_argument('--img-size', type=int, default=416, help='inference size (pixels)')parser.a(chǎn)dd_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')parser.a(chǎn)dd_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')parser.a(chǎn)dd_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')parser.a(chǎn)dd_argument('--classes', nargs='+', type=int, help='filter by class')parser.a(chǎn)dd_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')opt = parser.parse_args(args = [])
weights = opt.weightsimg_size = opt.img_size
# 初始化設(shè)備device = torch_utils.select_device(opt.device)
# 初始化模型model = Darknet(opt.cfg, img_size)
# 加載權(quán)重attempt_download(weights)if weights.endswith('.pt'): # pytorch格式 model.load_state_dict(torch.load(weights, map_location=device)['model'])else: # darknet 格式 load_darknet_weights(model, weights)
model.to(device).eval();# 獲取名字和顏色names = load_classes(opt.names)colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
%cd ..我們正在使用YOLOv3-sp-ultralytics權(quán)值,該算法稱其在平均精度上遠(yuǎn)遠(yuǎn)優(yōu)于其他YOLOv3模型functiontorch_utils.select_device()將自動(dòng)找到可用的GPU,除非輸入是“cpu”對(duì)象Darknet在PyTorch上使用預(yù)訓(xùn)練的權(quán)重來初始化了YOLOv3架構(gòu)(此時(shí)我們不希望訓(xùn)練模型)預(yù)測(cè)視頻中的目標(biāo)檢測(cè)接下來,我們將讀取視頻文件并使用矩陣框重寫視頻。def predict_one_video(path_video): cap = cv2.VideoCapture(path_video) _, img0 = cap.read()
save_path = os.path.join(output_dir, os.path.split(path_video)[-1]) fps = cap.get(cv2.CAP_PROP_FPS) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'MP4V'), fps, (w, h))我們使用MP4格式寫入新的視頻,變量為vid_writer,而寬度和高度則根據(jù)原始視頻來設(shè)置。開始對(duì)視頻中的每一幀進(jìn)行循環(huán)以獲得預(yù)測(cè)。while img0 is not None:
img = letterbox(img0, new_shape=opt.img_size)[0]
# 轉(zhuǎn)換 img = img[:, :, ::-1].transpose(2, 0, 1) # BGR 到 RGB, 到 3xHxW img = np.a(chǎn)scontiguousarray(img) img = torch.from_numpy(img).to(device) img = img.float() # uint8 到 fp16/32 img /= 255.0 # 0 - 255 到 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0)
pred = model(img)[0] # 應(yīng)用 NMS pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.a(chǎn)gnostic_nms)這個(gè)模型的圖像大小是416,letterbox函數(shù)調(diào)整了圖像的大小,并給圖像進(jìn)行填充,使得圖像可以被32整除。第二部分是將圖像轉(zhuǎn)換為RGB格式,并將通道設(shè)置到第一維,即(C,H,W),然后將圖像數(shù)據(jù)放入設(shè)備(GPU或CPU)中,將像素從0-255縮放到0-1。在我們將圖像放入模型之前,我們需要使用img.unsqeeze(0)函數(shù),因?yàn)槲覀儽仨殞D像重新格式化為4維(N,C,H,W), N是圖像的數(shù)量,在本例中為1。對(duì)圖像進(jìn)行預(yù)處理后,將其放入模型中得到預(yù)測(cè)框,但是預(yù)測(cè)有很多的框,所以我們需要非最大抑制方法來過濾和合并框。
畫邊界框和標(biāo)簽,然后寫入視頻我們?cè)贜MS之后循環(huán)所有的預(yù)測(cè)(pred)來繪制邊界框,但是圖像已經(jīng)被調(diào)整為416像素大小了,我們需要使用scale_coords函數(shù)將其縮放為原始大小,然后使用plot_one_box函數(shù)來繪制框 # 檢測(cè) for i, det in enumerate(pred): #檢測(cè)每個(gè)圖片 im0 = img0
if det is not None and len(det): # 更改框的大小 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# 寫入結(jié)果 for *xyxy, conf, cls in det: label = '%s %.2f' % (names[int(cls)], conf) plot_one_box(xyxy, im0, label=label, color=colors[int(cls)]) vid_writer.write(im0) _, img0 = cap.read()播放Colab的視頻視頻在函數(shù)predict_one_video被寫入為Mp4格式,我們壓縮成h264格式,所以視頻可以在谷歌Colab / Jupyter上直接播放。顯示原始視頻我們使用IPython.display.HTML來顯示視頻,其寬度為400像素,視頻是用二進(jìn)制讀取的。path_video = os.path.join("input_video","opera_h(yuǎn)ouse.mp4")save_path = predict_one_video(path_video)
# 顯示視頻mp4 = open(path_video,'rb').read()data_url = "data:video/mp4;base64," + b64encode(mp4).decode()HTML("""<video width=400 controls> <source src="%s" type="video/mp4"></video>""" % data_url)壓縮和顯示處理過的視頻OpenCV視頻寫入器的輸出是一個(gè)比原始視頻大3倍的Mp4視頻,它不能在谷歌Colab上顯示,解決方案之一是我們對(duì)視頻進(jìn)行壓縮。我們使用ffmpeg -i {save_path} -vcodec libx264 {compressed_path}path_video = os.path.join("input_video","opera_h(yuǎn)ouse.mp4")save_path = predict_one_video(path_video)# 壓縮視頻compressed_path = os.path.join("output_compressed", os.path.split(save_path)[-1])os.system(f"ffmpeg -i {save_path} -vcodec libx264 {compressed_path}")
#顯示視頻mp4 = open(compressed_path,'rb').read()data_url = "data:video/mp4;base64," + b64encode(mp4).decode()HTML("""<video width=400 controls> <source src="%s" type="video/mp4"></video>""" % data_url)結(jié)果
左邊是原始視頻,右邊是使用代碼處理過的視頻試試你自己的視頻轉(zhuǎn)到GitHub上的谷歌Colab文件上傳你的視頻在input_video文件夾中,只需運(yùn)行最后一個(gè)單元格即可(predict & show video)。
發(fā)表評(píng)論
請(qǐng)輸入評(píng)論內(nèi)容...
請(qǐng)輸入評(píng)論/評(píng)論長度6~500個(gè)字
最新活動(dòng)更多
-
即日-10.29立即報(bào)名>> 2024德州儀器嵌入式技術(shù)創(chuàng)新發(fā)展研討會(huì)
-
10月31日立即下載>> 【限時(shí)免費(fèi)下載】TE暖通空調(diào)系統(tǒng)高效可靠的組件解決方案
-
即日-11.13立即報(bào)名>>> 【在線會(huì)議】多物理場仿真助跑新能源汽車
-
11月14日立即報(bào)名>> 2024工程師系列—工業(yè)電子技術(shù)在線會(huì)議
-
12月19日立即報(bào)名>> 【線下會(huì)議】OFweek 2024(第九屆)物聯(lián)網(wǎng)產(chǎn)業(yè)大會(huì)
-
即日-12.26火熱報(bào)名中>> OFweek2024中國智造CIO在線峰會(huì)
推薦專題
- 1 Intel宣布40年來最重大轉(zhuǎn)型:年底前裁員15000人、拋掉2/3房產(chǎn)
- 2 因美封殺TikTok,字節(jié)股價(jià)骨折!估值僅Meta1/5
- 3 宏山激光重磅發(fā)布行業(yè)解決方案,助力智能制造產(chǎn)業(yè)新飛躍
- 4 國產(chǎn)AI芯片公司破產(chǎn)!白菜價(jià)拍賣
- 5 具身智能火了,但規(guī)模落地還需時(shí)間
- 6 國產(chǎn)英偉達(dá)們,抓緊沖刺A股
- 7 三次錯(cuò)失風(fēng)口!OpenAI前員工殺回AI編程賽道,老東家捧金相助
- 8 英特爾賦能智慧醫(yī)療,共創(chuàng)數(shù)字化未來
- 9 英偉達(dá)的麻煩在后頭?
- 10 將“網(wǎng)紅”變成“商品”,AI“爆改”實(shí)力拉滿
- 高級(jí)軟件工程師 廣東省/深圳市
- 自動(dòng)化高級(jí)工程師 廣東省/深圳市
- 光器件研發(fā)工程師 福建省/福州市
- 銷售總監(jiān)(光器件) 北京市/海淀區(qū)
- 激光器高級(jí)銷售經(jīng)理 上海市/虹口區(qū)
- 光器件物理工程師 北京市/海淀區(qū)
- 激光研發(fā)工程師 北京市/昌平區(qū)
- 技術(shù)專家 廣東省/江門市
- 封裝工程師 北京市/海淀區(qū)
- 結(jié)構(gòu)工程師 廣東省/深圳市