之前我有写过一篇处理视频数据集的文章,那时候做的工作是把视频帧提取处理出来再排序。现在回过头来看真的很蠢。按照之前的做法,首先,电脑内存可能没有那么大的空间,占用资源;其次,有了视频帧之后,标签不好处理,怎么输给神经网络也是个问题。
于是乎,在我了解到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*' )
第二步:
从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( )
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
完成。