SumMeVideo.py

from scipy.io import loadmat
import os
from moviepy.editor import VideoFileClip
import numpy as np


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)
        video_gt = video_gt['user_score']  # (n_frames,n_annotator)
        return self.bin_classify_user_score(video_gt)

    """
    :param user scores
    :returns binary nparray
    """

    def bin_classify_user_score(self, video_gt):
        user_scores = video_gt
        result = []
        for user in range(user_scores.shape[1]):
            print(user_scores[:, user].shape)
            result.append([int(item > 0.5) for item in user_scores[:, user]])
        return np.asarray(result).T  # (n_frames,n_annotator)

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