首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >Google Colab上的YOLOv3 PyTorch

Google Colab上的YOLOv3 PyTorch

作者头像
代码医生工作室
发布2020-04-14 16:29:40
2.5K0
发布2020-04-14 16:29:40
举报
文章被收录于专栏:相约机器人相约机器人

作者 | Hervind Philipe

来源 | Medium

编辑 | 代码医生团队

对于计算机视觉爱好者来说,YOLO(一次只看一次)是一种非常流行的实时对象检测概念,因为它的速度非常快并且性能出色。

在本文中,将共享用于处理视频的代码,以获取Google Colab内部每一帧的每个对象的边界框

不会讨论 YOLO的概念或体系结构,这里我们只讨论功能代码

开始吧

Wahid Khene在Unsplash上拍摄的照片

可以尝试自己在这个谷歌Colab。

https://colab.research.google.com/github/vindruid/yolov3-in-colab/blob/master/yolov3_video.ipynb

从Ultralytics 写得很好的最喜欢的github 回购开始。尽管该回购已经包含了如何仅使用YOLOv3来运行视频,但是python detect.py --source file.mp4还是想通过删除一些不必要的行来分解并简化代码,并添加如何在Google Colab / Jupyter笔记本上显示经过处理的视频。

准备YoloV3和LoadModel

首先克隆Ultralytics YoloV3存储库,然后导入通用软件包和repo的功能

!git clone https://github.com/ultralytics/yolov3
import time
import glob
import torch
import os
 
import argparse
from sys import platform
%cd yolov3
from models import *
from utils.datasets import *
from utils.utils import *
 
from IPython.display import HTML
from base64 import b64encode

设置参数解析器,初始化设备(CPU / CUDA),初始化YOLO模型,然后加载权重。

parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, default='cfg/yolov3-spp.cfg', help='*.cfg path')
parser.add_argument('--names', type=str, default='data/coco.names', help='*.names path')
parser.add_argument('--weights', type=str, default='weights/yolov3-spp-ultralytics.pt', help='weights path')
parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
opt = parser.parse_args(args = [])
 
weights = opt.weights
img_size =  opt.img_size
 
# Initialize Device
device = torch_utils.select_device(opt.device)
 
# Initialize model
model = Darknet(opt.cfg, img_size)
 
# Load weights
attempt_download(weights)
if weights.endswith('.pt'):  # pytorch format
    model.load_state_dict(torch.load(weights, map_location=device)['model'])
else:  # darknet format
    load_darknet_weights(model, weights)
 
model.to(device).eval();
# Get names and colors
names = load_classes(opt.names)
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
 
%cd ..

所使用的YOLOv3-spp-ultralytics权重在平均平均精度方面远胜于其他YOLOv3

该函数torch_utils.select_device()将自动查找可用的GPU,除非输入'cpu'

对象Darknet是在PyTorch上初始化YOLOv3架构的,并且需要使用预先训练的权重来加载权重(目前不希望训练模型)

预测视频中的目标检测

接下来,将读取视频文件,并使用对象边界框重写视频。以下3个GitHub Gist是predict_one_video将在最后使用的功能的一部分。

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。虽然fps,width和height根据原始视频使用

开始循环播放视频中的每个帧以获得预测。

    while img0 is not None:
 
        # Padded resize
        img = letterbox(img0, new_shape=opt.img_size)[0]
 
        # Convert and reformating
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3xHxW
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
 
        pred = model(img)[0]
        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)

该模型的图像大小为416。函数名称letterbox正在调整图像的大小并为图像添加填充,因此宽度或高度之一变为416,另一个小于等于416,但仍可被32整除

第二部分是将图像转换为RGB格式,并将通道置于第一维(C,H,W)。将图像数据放入设备(GPU或CPU),并将像素从缩放0-255到0-1。在将图像放入模型之前,使用函数,img.unsqeeze(0)因为必须将图像重新格式化为4维(N,C,H,W),其中N是图像数,在这种情况下为1。

在对图像进行预处理之后,将其放入模型中以获得预测框。但是预测有很多盒子,因此需要non-maximum suppression过滤和合并这些盒子。

非最大抑制(NMS)

绘制边界框和标签,然后编写视频

在NMS之后循环所有预测以绘制框,但是图像已经调整为416像素,需要使用函数scale_coords将其缩放回原始大小,然后使用函数绘制框plot_one_box

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            im0 = img0 ##### Ganti im0s menjadi img0
 
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
 
                # Write results
                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上显示视频

将视频predict_one_video保存为Mp4后,h264会将其压缩为Mp4格式,然后将其压缩,以便可以直接在Google Colab / Jupyter上播放视频。

显示原始视频

使用IPython.display.HTML400宽度像素显示视频。使用二进制文件读取视频

path_video = os.path.join("input_video","opera_house.mp4")
save_path = predict_one_video(path_video)
 
# Show 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视频编写器的输出是Mp4视频,其大小是原始视频的3倍,并且无法以相同的方式显示在Google Colab上,解决方案之一是进行压缩(源)

使用以下方式将Mp4视频压缩为h264ffmpeg -i {save_path} -vcodec libx264 {compressed_path}

path_video = os.path.join("input_video","opera_house.mp4")
save_path = predict_one_video(path_video)
# compress 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}")
 
# Show video
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)

结果

左是原始视频,右是使用这些代码处理的

尝试自己的视频

转至谷歌Colab文件GitHub上

https://colab.research.google.com/github/vindruid/yolov3-in-colab/blob/master/yolov3_video.ipynb

  • 将视频上传到input_video文件夹中
  • 只需运行最后一个单元格(预测并显示视频)

资源

  • YoloV3 Torch存储库

https://github.com/ultralytics/yolov3

  • Google Colab高级输出

https://colab.research.google.com/notebooks/snippets/advanced_outputs.ipynb#scrollTo=SucxddsPhOmj

  • 在Google Colab上显示视频

https://stackoverflow.com/questions/57377185/how-play-mp4-video-in-google-colab

  • 视频压缩

https://stackoverflow.com/questions/33134985/cv2-videowriter-will-not-write-file-using-fourcc-h-264-with-logitech-c920-pyth

谢谢您的阅读,希望对您有所帮助

本文参与 腾讯云自媒体分享计划,分享自微信公众号。
原始发表:2020-04-08,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 相约机器人 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体分享计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档