【pytorch】制作自己的数据集—基于VQA视频数据集

之前我有写过一篇处理视频数据集的文章，那时候做的工作是把视频帧提取处理出来再排序。现在对之前的方法做个改进。按照之前的做法，首先，电脑内存可能没有那么大的空间，占用资源；其次，有了视频帧之后，标签不好处理，怎么输给神经网络也是个问题。

于是，我在了解到pytorch可以自己做自己的数据集后，就迫不及待地动起手来了。

在开始之前，先来了解一下今天的处理的数据集。ODV-VQA数据集里面全部都是全景视频。他把这些全景视频分成10组，如下图：

每组里面数据共有54个受损视频和6个参考视频，每个参考视频对应9个受损视频，然后这9个受损视频呢，包含三种投影方式（全景视频投影到平面），分别是ERP、RCMP和TSP，每个投影方式里又包含3个qp。以group01中的两组视频为例：

每个视频的命名方式为：视频名称_投影方式_每帧大小_fps_qp_视频大小。关于视频的标签，采用的是主观评价打分的形式。在Group_ScoreData的压缩包里，存放的就是每个视频的分数。

这里由于我想做的是无参考视频评价，所以我把每个group的参考视频和参考视频的标签都给删了。（数据集在这篇文章里有：ODV-VQA数据集，翻不了墙的小伙伴在评论区call我哦）

下面就开始吧

首先我想的是读取视频和标签所在的文件夹，把他们放在列表里：

scorelist = glob.glob('./score/Group*')  # 视频和标签所在的文件夹
videolist = glob.glob('./data/Group*')   # 即，group1-10

第二步：从videolist和scorelist中读取视频和标签。这个过程靠定义两个函数，get_video_list(),get_score_file来完成。

def get_video_list(grouplist):
    video = []
    Video = []
    for group in grouplist:
        for videolist in glob.glob(str(group) + '/*.mp4'):
            video.append(videolist)
            video.sort()
        # Video.append(video)
    return video

def get_score_file(scorelist):
    SCORE=[]
    for score in scorelist:
        scorefile = []
        for file in glob.glob(score+"/*.xlsx"):
            scorefile.append(file)
            scorefile.sort()
        SCORE.append(scorefile)
    return SCORE
Video = get_video_list(grouplist)
scores = get_score_file(scorelist)

第三步，也是最后一步，就是写数据集啦

class VideoDataset(Data.Dataset):
    def __init__(self,videos,scores):

        super(VideoDataset,self).__init__()
        self.videos = videos

        self.scores = scores


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

    def __getitem__(self, idx):
        video = self.videos[idx]

        video_data = skvideo.io.vread(video)
        video_score = self.scores[idx]

        transform = transforms.Compose([transforms.ToTensor])

        video_length = video_data.shape[0]
        video_channel = video_data.shape[3]
        video_height = video_data.shape[1]
        video_width = video_data.shape[2]

        transformed_video = torch.zeros([video_length, video_channel, video_height, video_width])
        for frame_idx in range(video_length):
    
                frame = video_data[frame_idx]
                transformed_video[frame_idx] = frame
     
                if torch.cuda.is_available():
                transformed_video=transformed_video.cuda()
                video_score= video_score.cuda()
        sample = {'video': transformed_video,
                  'score': video_score}

        return sample