KAIR/utils/utils_videoio.py

import os
import cv2
import numpy as np
import torch
import random
from os import path as osp
from torchvision.utils import make_grid
import sys
from pathlib import Path
import six
from collections import OrderedDict
import math
import glob
import av
import io
from cv2 import (CAP_PROP_FOURCC, CAP_PROP_FPS, CAP_PROP_FRAME_COUNT,
                 CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH,
                 CAP_PROP_POS_FRAMES, VideoWriter_fourcc)

if sys.version_info <= (3, 3):
    FileNotFoundError = IOError
else:
    FileNotFoundError = FileNotFoundError


def is_str(x):
    """Whether the input is an string instance."""
    return isinstance(x, six.string_types)


def is_filepath(x):
    return is_str(x) or isinstance(x, Path)


def fopen(filepath, *args, **kwargs):
    if is_str(filepath):
        return open(filepath, *args, **kwargs)
    elif isinstance(filepath, Path):
        return filepath.open(*args, **kwargs)
    raise ValueError('`filepath` should be a string or a Path')


def check_file_exist(filename, msg_tmpl='file "{}" does not exist'):
    if not osp.isfile(filename):
        raise FileNotFoundError(msg_tmpl.format(filename))


def mkdir_or_exist(dir_name, mode=0o777):
    if dir_name == '':
        return
    dir_name = osp.expanduser(dir_name)
    os.makedirs(dir_name, mode=mode, exist_ok=True)


def symlink(src, dst, overwrite=True, **kwargs):
    if os.path.lexists(dst) and overwrite:
        os.remove(dst)
    os.symlink(src, dst, **kwargs)


def scandir(dir_path, suffix=None, recursive=False, case_sensitive=True):
    """Scan a directory to find the interested files.
    Args:
        dir_path (str | :obj:`Path`): Path of the directory.
        suffix (str | tuple(str), optional): File suffix that we are
            interested in. Default: None.
        recursive (bool, optional): If set to True, recursively scan the
            directory. Default: False.
        case_sensitive (bool, optional) : If set to False, ignore the case of
            suffix. Default: True.
    Returns:
        A generator for all the interested files with relative paths.
    """
    if isinstance(dir_path, (str, Path)):
        dir_path = str(dir_path)
    else:
        raise TypeError('"dir_path" must be a string or Path object')

    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
        raise TypeError('"suffix" must be a string or tuple of strings')

    if suffix is not None and not case_sensitive:
        suffix = suffix.lower() if isinstance(suffix, str) else tuple(
            item.lower() for item in suffix)

    root = dir_path

    def _scandir(dir_path, suffix, recursive, case_sensitive):
        for entry in os.scandir(dir_path):
            if not entry.name.startswith('.') and entry.is_file():
                rel_path = osp.relpath(entry.path, root)
                _rel_path = rel_path if case_sensitive else rel_path.lower()
                if suffix is None or _rel_path.endswith(suffix):
                    yield rel_path
            elif recursive and os.path.isdir(entry.path):
                # scan recursively if entry.path is a directory
                yield from _scandir(entry.path, suffix, recursive,
                                    case_sensitive)

    return _scandir(dir_path, suffix, recursive, case_sensitive)


class Cache:

    def __init__(self, capacity):
        self._cache = OrderedDict()
        self._capacity = int(capacity)
        if capacity <= 0:
            raise ValueError('capacity must be a positive integer')

    @property
    def capacity(self):
        return self._capacity

    @property
    def size(self):
        return len(self._cache)

    def put(self, key, val):
        if key in self._cache:
            return
        if len(self._cache) >= self.capacity:
            self._cache.popitem(last=False)
        self._cache[key] = val

    def get(self, key, default=None):
        val = self._cache[key] if key in self._cache else default
        return val


class VideoReader:
    """Video class with similar usage to a list object.

    This video warpper class provides convenient apis to access frames.
    There exists an issue of OpenCV's VideoCapture class that jumping to a
    certain frame may be inaccurate. It is fixed in this class by checking
    the position after jumping each time.
    Cache is used when decoding videos. So if the same frame is visited for
    the second time, there is no need to decode again if it is stored in the
    cache.

    """

    def __init__(self, filename, cache_capacity=10):
        # Check whether the video path is a url
        if not filename.startswith(('https://', 'http://')):
            check_file_exist(filename, 'Video file not found: ' + filename)
        self._vcap = cv2.VideoCapture(filename)
        assert cache_capacity > 0
        self._cache = Cache(cache_capacity)
        self._position = 0
        # get basic info
        self._width = int(self._vcap.get(CAP_PROP_FRAME_WIDTH))
        self._height = int(self._vcap.get(CAP_PROP_FRAME_HEIGHT))
        self._fps = self._vcap.get(CAP_PROP_FPS)
        self._frame_cnt = int(self._vcap.get(CAP_PROP_FRAME_COUNT))
        self._fourcc = self._vcap.get(CAP_PROP_FOURCC)

    @property
    def vcap(self):
        """:obj:`cv2.VideoCapture`: The raw VideoCapture object."""
        return self._vcap

    @property
    def opened(self):
        """bool: Indicate whether the video is opened."""
        return self._vcap.isOpened()

    @property
    def width(self):
        """int: Width of video frames."""
        return self._width

    @property
    def height(self):
        """int: Height of video frames."""
        return self._height

    @property
    def resolution(self):
        """tuple: Video resolution (width, height)."""
        return (self._width, self._height)

    @property
    def fps(self):
        """float: FPS of the video."""
        return self._fps

    @property
    def frame_cnt(self):
        """int: Total frames of the video."""
        return self._frame_cnt

    @property
    def fourcc(self):
        """str: "Four character code" of the video."""
        return self._fourcc

    @property
    def position(self):
        """int: Current cursor position, indicating frame decoded."""
        return self._position

    def _get_real_position(self):
        return int(round(self._vcap.get(CAP_PROP_POS_FRAMES)))

    def _set_real_position(self, frame_id):
        self._vcap.set(CAP_PROP_POS_FRAMES, frame_id)
        pos = self._get_real_position()
        for _ in range(frame_id - pos):
            self._vcap.read()
        self._position = frame_id

    def read(self):
        """Read the next frame.

        If the next frame have been decoded before and in the cache, then
        return it directly, otherwise decode, cache and return it.

        Returns:
            ndarray or None: Return the frame if successful, otherwise None.
        """
        # pos = self._position
        if self._cache:
            img = self._cache.get(self._position)
            if img is not None:
                ret = True
            else:
                if self._position != self._get_real_position():
                    self._set_real_position(self._position)
                ret, img = self._vcap.read()
                if ret:
                    self._cache.put(self._position, img)
        else:
            ret, img = self._vcap.read()
        if ret:
            self._position += 1
        return img

    def get_frame(self, frame_id):
        """Get frame by index.

        Args:
            frame_id (int): Index of the expected frame, 0-based.

        Returns:
            ndarray or None: Return the frame if successful, otherwise None.
        """
        if frame_id < 0 or frame_id >= self._frame_cnt:
            raise IndexError(
                f'"frame_id" must be between 0 and {self._frame_cnt - 1}')
        if frame_id == self._position:
            return self.read()
        if self._cache:
            img = self._cache.get(frame_id)
            if img is not None:
                self._position = frame_id + 1
                return img
        self._set_real_position(frame_id)
        ret, img = self._vcap.read()
        if ret:
            if self._cache:
                self._cache.put(self._position, img)
            self._position += 1
        return img

    def current_frame(self):
        """Get the current frame (frame that is just visited).

        Returns:
            ndarray or None: If the video is fresh, return None, otherwise
            return the frame.
        """
        if self._position == 0:
            return None
        return self._cache.get(self._position - 1)

    def cvt2frames(self,
                   frame_dir,
                   file_start=0,
                   filename_tmpl='{:06d}.jpg',
                   start=0,
                   max_num=0,
                   show_progress=False):
        """Convert a video to frame images.

        Args:
            frame_dir (str): Output directory to store all the frame images.
            file_start (int): Filenames will start from the specified number.
            filename_tmpl (str): Filename template with the index as the
                placeholder.
            start (int): The starting frame index.
            max_num (int): Maximum number of frames to be written.
            show_progress (bool): Whether to show a progress bar.
        """
        mkdir_or_exist(frame_dir)
        if max_num == 0:
            task_num = self.frame_cnt - start
        else:
            task_num = min(self.frame_cnt - start, max_num)
        if task_num <= 0:
            raise ValueError('start must be less than total frame number')
        if start > 0:
            self._set_real_position(start)

        def write_frame(file_idx):
            img = self.read()
            if img is None:
                return
            filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
            cv2.imwrite(filename, img)

        if show_progress:
            pass
            #track_progress(write_frame, range(file_start,file_start + task_num))
        else:
            for i in range(task_num):
                write_frame(file_start + i)

    def __len__(self):
        return self.frame_cnt

    def __getitem__(self, index):
        if isinstance(index, slice):
            return [
                self.get_frame(i)
                for i in range(*index.indices(self.frame_cnt))
            ]
        # support negative indexing
        if index < 0:
            index += self.frame_cnt
            if index < 0:
                raise IndexError('index out of range')
        return self.get_frame(index)

    def __iter__(self):
        self._set_real_position(0)
        return self

    def __next__(self):
        img = self.read()
        if img is not None:
            return img
        else:
            raise StopIteration

    next = __next__

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._vcap.release()


def frames2video(frame_dir,
                 video_file,
                 fps=30,
                 fourcc='XVID',
                 filename_tmpl='{:06d}.jpg',
                 start=0,
                 end=0,
                 show_progress=False):
    """Read the frame images from a directory and join them as a video.

    Args:
        frame_dir (str): The directory containing video frames.
        video_file (str): Output filename.
        fps (float): FPS of the output video.
        fourcc (str): Fourcc of the output video, this should be compatible
            with the output file type.
        filename_tmpl (str): Filename template with the index as the variable.
        start (int): Starting frame index.
        end (int): Ending frame index.
        show_progress (bool): Whether to show a progress bar.
    """
    if end == 0:
        ext = filename_tmpl.split('.')[-1]
        end = len([name for name in scandir(frame_dir, ext)])
    first_file = osp.join(frame_dir, filename_tmpl.format(start))
    check_file_exist(first_file, 'The start frame not found: ' + first_file)
    img = cv2.imread(first_file)
    height, width = img.shape[:2]
    resolution = (width, height)
    vwriter = cv2.VideoWriter(video_file, VideoWriter_fourcc(*fourcc), fps,
                              resolution)

    def write_frame(file_idx):
        filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
        img = cv2.imread(filename)
        vwriter.write(img)

    if show_progress:
        pass
        # track_progress(write_frame, range(start, end))
    else:
        for i in range(start, end):
            write_frame(i)
    vwriter.release()


def video2images(video_path, output_dir):
    vidcap = cv2.VideoCapture(video_path)
    in_fps = vidcap.get(cv2.CAP_PROP_FPS)
    print('video fps:', in_fps)
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)
    loaded, frame = vidcap.read()
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f'number of total frames is: {total_frames:06}')
    for i_frame in range(total_frames):
        if i_frame % 100 == 0:
            print(f'{i_frame:06} / {total_frames:06}')
        frame_name = os.path.join(output_dir, f'{i_frame:06}' + '.png')
        cv2.imwrite(frame_name, frame)
        loaded, frame = vidcap.read()


def images2video(image_dir, video_path, fps=24, image_ext='png'):
    '''
    #codec = cv2.VideoWriter_fourcc(*'XVID')
    #codec = cv2.VideoWriter_fourcc('A','V','C','1')
    #codec = cv2.VideoWriter_fourcc('Y','U','V','1')
    #codec = cv2.VideoWriter_fourcc('P','I','M','1')
    #codec = cv2.VideoWriter_fourcc('M','J','P','G')
    codec = cv2.VideoWriter_fourcc('M','P','4','2')
    #codec = cv2.VideoWriter_fourcc('D','I','V','3')
    #codec =  cv2.VideoWriter_fourcc('D','I','V','X')
    #codec = cv2.VideoWriter_fourcc('U','2','6','3')
    #codec = cv2.VideoWriter_fourcc('I','2','6','3')
    #codec = cv2.VideoWriter_fourcc('F','L','V','1')
    #codec = cv2.VideoWriter_fourcc('H','2','6','4')
    #codec = cv2.VideoWriter_fourcc('A','Y','U','V')
    #codec = cv2.VideoWriter_fourcc('I','U','Y','V')
    编码器常用的几种：
    cv2.VideoWriter_fourcc("I", "4", "2", "0") 
        压缩的yuv颜色编码器，4:2:0色彩度子采样 兼容性好，产生很大的视频 avi
    cv2.VideoWriter_fourcc("P", I", "M", "1")
        采用mpeg-1编码，文件为avi
    cv2.VideoWriter_fourcc("X", "V", "T", "D")
        采用mpeg-4编码，得到视频大小平均 拓展名avi
    cv2.VideoWriter_fourcc("T", "H", "E", "O")
        Ogg Vorbis， 拓展名为ogv
    cv2.VideoWriter_fourcc("F", "L", "V", "1")
        FLASH视频，拓展名为.flv
    '''
    image_files = sorted(glob.glob(os.path.join(image_dir, '*.{}'.format(image_ext))))
    print(len(image_files))
    height, width, _ = cv2.imread(image_files[0]).shape
    out_fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')  # cv2.VideoWriter_fourcc(*'MP4V')
    out_video = cv2.VideoWriter(video_path, out_fourcc, fps, (width, height))

    for image_file in image_files:
        img = cv2.imread(image_file)
        img = cv2.resize(img, (width, height), interpolation=3)
        out_video.write(img)
    out_video.release()


def add_video_compression(imgs):
    codec_type = ['libx264', 'h264', 'mpeg4']
    codec_prob = [1 / 3., 1 / 3., 1 / 3.]
    codec = random.choices(codec_type, codec_prob)[0]
    # codec = 'mpeg4'
    bitrate = [1e4, 1e5]
    bitrate = np.random.randint(bitrate[0], bitrate[1] + 1)

    buf = io.BytesIO()
    with av.open(buf, 'w', 'mp4') as container:
        stream = container.add_stream(codec, rate=1)
        stream.height = imgs[0].shape[0]
        stream.width = imgs[0].shape[1]
        stream.pix_fmt = 'yuv420p'
        stream.bit_rate = bitrate
        
        for img in imgs:
            img = np.uint8((img.clip(0, 1)*255.).round())
            frame = av.VideoFrame.from_ndarray(img, format='rgb24')
            frame.pict_type = 'NONE'
            # pdb.set_trace()
            for packet in stream.encode(frame):
                container.mux(packet)

        # Flush stream
        for packet in stream.encode():
            container.mux(packet)

    outputs = []
    with av.open(buf, 'r', 'mp4') as container:
        if container.streams.video:
            for frame in container.decode(**{'video': 0}):
                outputs.append(
                    frame.to_rgb().to_ndarray().astype(np.float32) / 255.)

    #outputs = np.stack(outputs, axis=0)
    return outputs


if __name__ == '__main__':

    # -----------------------------------
    # test VideoReader(filename, cache_capacity=10)
    # -----------------------------------
#    video_reader = VideoReader('utils/test.mp4')
#    from utils import utils_image as util
#    inputs = []
#    for frame in video_reader:
#        print(frame.dtype)
#        util.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
#        #util.imshow(np.flip(frame, axis=2))

    # -----------------------------------
    # test video2images(video_path, output_dir)
    # -----------------------------------
#    video2images('utils/test.mp4', 'frames')

    # -----------------------------------
    # test images2video(image_dir, video_path, fps=24, image_ext='png')
    # -----------------------------------
#    images2video('frames', 'video_02.mp4', fps=30, image_ext='png')


    # -----------------------------------
    # test frames2video(frame_dir, video_file, fps=30, fourcc='XVID', filename_tmpl='{:06d}.png')
    # -----------------------------------
#    frames2video('frames', 'video_01.mp4', filename_tmpl='{:06d}.png')


    # -----------------------------------
    # test add_video_compression(imgs)
    # -----------------------------------
#    imgs = []
#    image_ext = 'png'
#    frames = 'frames'
#    from utils import utils_image as util
#    image_files = sorted(glob.glob(os.path.join(frames, '*.{}'.format(image_ext))))
#    for i, image_file in enumerate(image_files):
#        if i < 7:
#            img = util.imread_uint(image_file, 3)
#            img = util.uint2single(img)
#            imgs.append(img)
#
#    results = add_video_compression(imgs)
#    for i, img in enumerate(results):
#        util.imshow(util.single2uint(img))
#        util.imsave(util.single2uint(img),f'{i:05}.png')

    # run utils/utils_video.py