如何解决视频的 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)
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。