VSUMMVideo.py

import os
from moviepy.editor import VideoFileClip
import numpy as np
from src.utils import digits_in_string, has_string


class VSUMMVideo(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)
        self.video_name = video_name
        self.nframes = self.video_clip.reader.infos['video_nframes']

    def get_gt(self):
        vid_scores = []
        for i, user in enumerate(os.listdir(self.gt_path)):
            user_dir = os.path.join(self.gt_path, user)
            frame_idx = []
            for idx, summary in enumerate(os.listdir(user_dir)):
                if has_string(summary, 'frame'):
                    frame_idx.append(digits_in_string(summary))  # Frame123--> 123
            vid_scores.append(frame_idx)
        binary_scores = self.bin_classify_user_score(vid_scores)

        return binary_scores

    def bin_classify_user_score(self, user_scores_list):
        result = []
        for frames_idx in list(user_scores_list):
            vid_frames = np.zeros(self.nframes)
            np.put(vid_frames, frames_idx, [1])  # replace all selected frames with 1
            result.append(vid_frames)
        return np.asarray(result).T  # (n_frames,n_annotator)

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