YOLOv3+opencv实现简单的实时目标检测
实现思路
经过了各种环境配置,最近终于初步跑通了YOLOv3,用的是以下程序。
[https://github.com/eriklindernoren/PyTorch-YOLOv3]
跑通了之后想着是可以配合摄像头做一个简单的目标检测试试,主要是用opencv调用摄像头拍照,然后YOLOv3检测最后呈现。整体实现过程还算顺利,也最终感受到了GPU加速的强大,用和不用视频流畅度天差地别。这里主要是对过程的一个记录吧,第一次写,着实有点紧张。
opencv调用摄像头
opencv调用摄像头这个代码网上能搜到的很多了,代码如下:
import cv2
cap = cv2.VideoCapture(0)## 0 电脑摄像头 1 usb摄像头
while 1:
ret, frame = cap.read()
cv2.imshow("cap", frame)
if cv2.waitKey(10) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
YOLOv3算法
我也是初步接触这个算法以及这篇程序,所以对其内在底层计算原理目前了解的不是很透彻,不过如果只是简单的用它目前还不需要了解那么深。
按照程序链接配置好之后,可以根据其detect.py函数了解程序主要流程。主要包括加载模型、加载训练数据,然后调用YOLOv3网络预测结果,结果主要包括检测目标的坐标信息、目标标签等。最后通过plt函数保存图片。
实时检测实现
了解了opencv以及配置好了YOLOv3之后,则可以开始简单的实现了,主要是在detect.py的基础上进行修改,加入调用摄像头的内容。
这里遇到的一个问题主要是opencv调用摄像头拍的照片所生成的数组没法直接给YOLOv3网络作为输入(改了维度也不行),结果会出错,试了很久这里也没解决,最后只能想了一个比较笨的办法,拍照之后先保存一张图片,然后再由DataLoader函数加载这个图片,这样做出来结果是对的。
基本流程为opencv打开摄像头,拍照,保存图片,加载YOLOv3模型,加载前面保存的图片,根据图片计算结果,opencv画出目标位置及标签并展示照片。整体代码如下:
from __future__ import division
from models import *
from utils.utils import *
from utils.datasets import *
import os
import sys
import time
import datetime
import argparse
from PIL import Image
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torch.autograd import Variable
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import NullLocator
import cv2
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--image_folder", type=str, default="data/samples/ceshi"