fairseq/examples/MMPT/scripts/video_feature_extractor/videoreader.py

# Copyright Howto100M authors.
# Copyright (c) Facebook, Inc. All Rights Reserved

import torch as th
import pandas as pd
import os
import numpy as np
import ffmpeg
import random

from torch.utils.data import Dataset


class VideoLoader(Dataset):
    """modified from how2's video_feature_extractor."""
    def __init__(
        self,
        csv=None,
        video_dict=None,
        framerate=1,
        size=112,
        centercrop=False,
        hflip=False,
        **kwargs
    ):
        if csv is None and video_dict is None:
            raise ValueError("csv and video_dict cannot be both None.")
        if csv is not None:
            self.csv = pd.read_csv(csv)
        if video_dict is not None:
            self.csv = pd.DataFrame.from_dict(video_dict)

        self.centercrop = centercrop
        self.size = size
        self.framerate = framerate
        self.hflip = hflip

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

    def _get_video_dim(self, video_path):
        probe = ffmpeg.probe(video_path)
        video_stream = next((stream for stream in probe['streams']
                             if stream['codec_type'] == 'video'), None)
        width = int(video_stream['width'])
        height = int(video_stream['height'])
        return height, width

    def _get_video_info(self, video_path):
        probe = ffmpeg.probe(video_path)
        video_stream = next((stream for stream in probe['streams']
                             if stream['codec_type'] == 'video'), None)
        return video_stream

    def _get_output_dim(self, h, w):
        if isinstance(self.size, tuple) and len(self.size) == 2:
            return self.size
        elif h >= w:
            return int(h * self.size / w), self.size
        else:
            return self.size, int(w * self.size / h)

    def __getitem__(self, idx):
        video_path = self.csv['video_path'].values[idx]
        output_file = self.csv['feature_path'].values[idx]
        return self._decode(output_file, video_path)

    def _decode(self, output_file, video_path):
        if not(os.path.isfile(output_file)) and os.path.isfile(video_path):
            try:
                h, w = self._get_video_dim(video_path)
            except Exception:
                print('ffprobe failed at: {}'.format(video_path))
                return {'video': th.zeros(1), 'input': video_path,
                        'output': output_file}
            try:
                os.makedirs(os.path.dirname(output_file), exist_ok=True)
                height, width = self._get_output_dim(h, w)

                cmd = (
                    ffmpeg
                    .input(video_path)
                    .filter('fps', fps=self.framerate)
                    .filter('scale', width, height)
                )
                if self.hflip:
                    cmd = cmd.filter('hflip')

                if self.centercrop:
                    x = int((width - self.size) / 2.0)
                    y = int((height - self.size) / 2.0)
                    cmd = cmd.crop(x, y, self.size, self.size)
                video = self._run(cmd, output_file)
            except Exception:
                video = th.zeros(1)
        else:
            video = th.zeros(1)

        return {'video': video, 'input': video_path, 'output': output_file}

    def _run(self, cmd, output_file):
        out, _ = (
            cmd.output('pipe:', format='rawvideo', pix_fmt='rgb24')
            .run(capture_stdout=True, quiet=True)
        )
        if self.centercrop and isinstance(self.size, int):
            height, width = self.size, self.size
        video = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
        video = th.from_numpy(video.astype('float32'))
        return video.permute(0, 3, 1, 2)


class VideoVerifier(VideoLoader):
    def __getitem__(self, idx):
        video_path = self.csv['video_path'].values[idx]
        try:
            return self._get_video_info(video_path)
        except Exception:
            # print('ffprobe failed at: {}'.format(video_path))
            return None


class VideoCompressor(VideoLoader):
    def __init__(
        self,
        csv=None,
        video_dict=None,
        framerate=1,
        size=112,
        centercrop=False,
        hflip=False,
        crf=32,
        **kwargs
    ):
        super().__init__(
            csv,
            video_dict,
            framerate,
            size,
            centercrop,
            hflip
        )
        self.crf = crf

    def _run(self, cmd, output_file):
        out, _ = (
            cmd.output(filename=output_file, crf=self.crf)
            .run(quiet=True)
        )
        video = None
        return video


class VideoDownloader(VideoCompressor):
    """download"""
    def __getitem__(self, idx):
        video_path = self.csv['video_path'].values[idx]
        output_file = self.csv['feature_path'].values[idx]
        if not(os.path.isfile(output_file)):
            os.makedirs(os.path.dirname(output_file), exist_ok=True)
            cmd = "wget -O" + output_file + " " + video_path
            # import subprocess
            # subprocess.check_output(
            #    cmd,
            #    stderr=subprocess.STDOUT, shell=True)
            os.system(cmd)
        return {'video': None, 'input': video_path, 'output': output_file}


class AvKeyframeVideoCompressor(VideoLoader):
    """extract keyframes from a video and save it as jpg.
    TODO: consider to merge with `CodecProcessor`.
    """
    def __init__(
        self,
        csv=None,
        video_dict=None,
        framerate=1,
        size=112,
        centercrop=False,
        max_num_frames=5,
        **kwargs
    ):
        super().__init__(csv, video_dict, framerate, size, centercrop)
        self.max_num_frames = max_num_frames

    def _get_video_dim(self, video_fn):
        """decord cannot probe the size of a video, we use pyav instead."""
        import av
        with av.open(video_fn) as container:
            height = container.streams.video[0].codec_context.height
            width = container.streams.video[0].codec_context.width
        return height, width

    def _get_output_dim(self, height, width):
        """
        keep the shorter side be `self.size`, strech the other.
        """
        if height >= width:
            return int(height * self.size / width), self.size
        else:
            return self.size, int(width * self.size / height)

    def __getitem__(self, idx):
        import av
        video_path = self.csv['video_path'].values[idx]
        output_file = self.csv['feature_path'].values[idx]
        if not(os.path.isdir(output_file)) and os.path.isfile(video_path):
            try:
                h, w = self._get_video_dim(video_path)
            except Exception:
                print('probe failed at: {}'.format(video_path))
                return {'video': th.zeros(1), 'input': video_path,
                        'output': output_file}

            try:
                height, width = self._get_output_dim(h, w)

                # new for av.
                with av.open(video_path) as container:
                    container.streams.video[0].thread_type = "AUTO"
                    container.streams.video[0].codec_context.height = height
                    container.streams.video[0].codec_context.width = width
                    if self.framerate == 0:     # keyframe.
                        container.streams.video[0].codec_context.skip_frame = 'NONKEY'
                    frames = []
                    for frame in container.decode(video=0):
                        frames.append(frame)
                    frames = random.sample(frames, self.max_num_frames)

                    os.makedirs(output_file, exist_ok=True)
                    for frame in frames:
                        frame.to_image().save(
                            os.path.join(
                                output_file,
                                "%04d.jpg" % frame.index))
            except Exception:
                print('extract failed at: {}'.format(video_path))
                return {'video': th.zeros(1), 'input': video_path,
                        'output': output_file}
        video = th.zeros(1)
        return {'video': video, 'input': video_path, 'output': output_file}