时间戳

首页 > 解决方案 > Pytorch 视频数据集

pytorch - Pytorch 视频数据集

问题描述

嗨，我制作了一个视频帧加载器数据集，以输入到 pytorch 模型中。我想从视频中采样帧，但是应该从每个视频中统一采样帧。这是我想出的课程。我想知道是否有更好的方法来加快采样过程。
你有什么建议，特别是在read_video方法部分？
谢谢

import torch
import torchvision as tv
import cv2
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
from pathlib import Path

class VideoLoader(torch.utils.data.Dataset):

  def __init__(self, data_path, classes, transforms=None, max_frames=None, frames_ratio=None):
    super(VideoLoader, self).__init__()

    self.data_path = data_path
    self.classes = classes
    self.frames_ratio = frames_ratio

    self.transforms = transforms
    self.max_frames = max_frames
  
  def read_video(self, path):
    frames = []

    vc = cv2.VideoCapture(path)

    total_frames = int(vc.get(cv2.CAP_PROP_FRAME_COUNT))
    
    if self.frames_ratio:
      if type(self.frames_ratio) is float:
        frames_to_pick = int(total_frames * self.frames_ratio)
      else:
        frames_to_pick = self.frames_ratio
    else:
        frames_to_pick = total_frames
    
    idxs = np.linspace(0, total_frames, frames_to_pick, endpoint=False)

    for i in idxs:
      ok, f = vc.read()
      if ok:

        f = tv.transforms.ToTensor()(f)
        f = self.transforms(f) if self.transforms else f
        frames.append(f)
        
        vc.set(cv2.CAP_PROP_POS_FRAMES, i)
        if self.max_frames and len(frames) == self.max_frames: break
      else: break
    vc.release()
    return torch.stack(frames)

  def __getitem__(self, index):
    v_path, label = self.data_path[index]
    return self.read_video(v_path), self.classes[label]

  def __len__(self): return len(self.data_path)

标签： pytorchconv-neural-networkdataloaderpytorch-dataloader

解决方案

因为您不能真正并行地搜索视频，所以实际上没有任何更快的采样过程可以在本地运行。我个人遇到了这个问题，这就是为什么我开始为此构建一个名为Sieve的简单 API 。您可以直接将数据上传到 Sieve（从云存储桶或本地存储），它会快速为您切割所有帧，甚至用动作、人物、物体等标记它们。它在云中使用无服务器功能并行化，这使得它非常快，即使是数小时或数天的镜头。

然后，您可以使用仪表板快速从 Sieve 导出，该仪表板为您提供快速curl命令，您可以运行以下载所需的确切样本。

这是一个有用的回购：https ://github.com/Sieve-Data/automatic-video-processing

推荐阅读