对于计算机视觉爱好者来说，yolo (you only look once)是一个非常流行的实时目标检测算法，因为它非常快，同时性能非常好。在本文中，我将共享一个视频处理的代码，以获取视频中每个对象目标的边框。我们将不讨论yolo的概念或架构，因为很多好的文章已经在媒体中详细阐述了这些知识点。这里我们只讨论函数代码。开始
谷歌colab地址：32 img /= 255.0 # 0 - 255 到 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0)
pred = model(img)[0] # 应用 nms pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)这个模型的图像大小是416，letterbox函数调整了图像的大小，并给图像进行填充，使得图像可以被32整除。第二部分是将图像转换为rgb格式，并将通道设置到第一维，即(c,h,w)，然后将图像数据放入设备(gpu或cpu)中，将像素从0-255缩放到0-1。在我们将图像放入模型之前，我们需要使用img.unsqeeze(0)函数，因为我们必须将图像重新格式化为4维(n,c,h,w)， n是图像的数量，在本例中为1。对图像进行预处理后，将其放入模型中得到预测框，但是预测有很多的框，所以我们需要非最大抑制方法来过滤和合并框。
画边界框和标签，然后写入视频我们在nms之后循环所有的预测(pred)来绘制边界框，但是图像已经被调整为416像素大小了，我们需要使用scale_coords函数将其缩放为原始大小，然后使用plot_one_box函数来绘制框 # 检测 for i, det in enumerate(pred): #检测每个图片 im0 = img0
if det is not none and len(det): # 更改框的大小 det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# 写入结果 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格式，所以视频可以在谷歌colab / jupyter上直接播放。显示原始视频我们使用ipython.display.html来显示视频，其宽度为400像素，视频是用二进制读取的。path_video = os.path.join("input_video","opera_house.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视频写入器的输出是一个比原始视频大3倍的mp4视频，它不能在谷歌colab上显示，解决方案之一是我们对视频进行压缩。我们使用ffmpeg -i {save_path} -vcodec libx264 {compressed_path}path_video = os.path.join("input_video","opera_house.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)结果
左边是原始视频，右边是使用代码处理过的视频试试你自己的视频转到github上的谷歌colab文件(/github/vindruid/yolov3-in-colab/blob/master/yolov3_video.ipynb)上传你的视频在input_video文件夹中，只需运行最后一个单元格即可(predict & show video)