SumMeVideo.py

from scipy.io import loadmat
import os
from moviepy.editor import VideoFileClip


class SumMeVideo(VideoFileClip):
    def __init__(self, video_name, video_path, gt_base_dir):
        self.video_clip = VideoFileClip(video_path)
        self.fps = int(self.video_clip.fps)
        self.duration = int(self.video_clip.duration)
        self.gt_path = os.path.join(gt_base_dir, video_name + '.mat')
        self.video_name = video_name

    def get_gt(self):
        video_gt = loadmat(self.gt_path)
        return video_gt['gt_score']

    def get_frames(self):
        return list(self.video_clip.iter_frames(with_times=False))