video_resize_resolution_mapper.py

import math
import os
import sys

from pydantic import PositiveInt

from data_juicer.utils.constant import Fields
from data_juicer.utils.file_utils import transfer_filename
from data_juicer.utils.lazy_loader import LazyLoader
from data_juicer.utils.logger_utils import HiddenPrints
from data_juicer.utils.mm_utils import close_video, load_video

from ..base_op import OPERATORS, Mapper
from ..op_fusion import LOADED_VIDEOS

with HiddenPrints():
    ffmpeg = LazyLoader('ffmpeg', 'ffmpeg')

OP_NAME = 'video_resize_resolution_mapper'


@OPERATORS.register_module(OP_NAME)
@LOADED_VIDEOS.register_module(OP_NAME)
class VideoResizeResolutionMapper(Mapper):
    """
        Mapper to resize videos resolution. We leave the super resolution
        with deep learning for future works.
    """

    def __init__(self,
                 min_width: int = 1,
                 max_width: int = sys.maxsize,
                 min_height: int = 1,
                 max_height: int = sys.maxsize,
                 force_original_aspect_ratio: str = 'disable',
                 force_divisible_by: PositiveInt = 2,
                 *args,
                 **kwargs):
        """
        Initialization method.

        :param min_width: Videos with width less than 'min_width' will be
            mapped to videos with equal or bigger width.
        :param max_width: Videos with width more than 'max_width' will be
            mapped to videos with equal of smaller width.
        :param min_height: Videos with height less than 'min_height' will be
            mapped to videos with equal or bigger height.
        :param max_height: Videos with height more than 'max_height' will be
            mapped to videos with equal or smaller height.
        :param force_original_aspect_ratio: Enable decreasing or \
            increasing output video width or height if necessary \
            to keep the original aspect ratio, including ['disable', \
            'decrease', 'increase'].
        :param force_divisible_by: Ensures that both the output dimensions, \
            width and height, are divisible by the given integer when used \
            together with force_original_aspect_ratio, must be a positive \
            even number.
        :param args: extra args
        :param kwargs: extra args
        """
        super().__init__(*args, **kwargs)
        self._init_parameters = self.remove_extra_parameters(locals())

        force_original_aspect_ratio = force_original_aspect_ratio.lower()

        if force_original_aspect_ratio not in [
                'disable', 'decrease', 'increase'
        ]:
            raise ValueError(
                f'force_original_aspect_ratio [{force_original_aspect_ratio}]'
                f' is not supported. '
                f"Can only be one of ['disable', 'decrease', 'increase']. ")
        if (force_divisible_by <= 1 or force_divisible_by % 2
                == 1) and force_original_aspect_ratio != 'disable':
            raise ValueError(
                f'force_divisible_by [{force_divisible_by}] must be a positive'
                f' even number. ')

        self.min_width = min_width
        self.max_width = max_width
        self.min_height = min_height
        self.max_height = max_height
        self.scale_method = 'scale'
        self.force_original_aspect_ratio = force_original_aspect_ratio
        self.force_divisible_by = force_divisible_by

    def process_single(self, sample, context=False):
        # there is no video in this sample
        if self.video_key not in sample or not sample[self.video_key]:
            sample[Fields.source_file] = []
            return sample

        if Fields.source_file not in sample or not sample[Fields.source_file]:
            sample[Fields.source_file] = sample[self.video_key]

        loaded_video_keys = sample[self.video_key]

        for index, video_key in enumerate(loaded_video_keys):

            container = load_video(video_key)
            video = container.streams.video[0]
            width = video.codec_context.width
            height = video.codec_context.height
            origin_ratio = width / height
            close_video(container)

            if width >= self.min_width and width <= self.max_width and \
               height >= self.min_height and height <= self.max_height:
                continue

            # keep the original aspect ratio as possible
            if width < self.min_width:
                height = self.min_width / origin_ratio
                width = self.min_width
            if width > self.max_width:
                height = self.max_width / origin_ratio
                width = self.max_width
            if height < self.min_height:
                width = self.min_height * origin_ratio
                height = self.min_height
            if height > self.max_height:
                width = self.max_height * origin_ratio
                height = self.max_height

            # the width and height of a video must be divisible by 2.
            if self.force_original_aspect_ratio == 'disable':
                force_divisible_by = 2
            else:
                force_divisible_by = self.force_divisible_by

            # make sure in the range if possible
            width = int(max(width, self.min_width))
            width = math.ceil(width / force_divisible_by) * force_divisible_by
            width = int(min(width, self.max_width))
            width = int(width / force_divisible_by) * force_divisible_by
            height = int(max(height, self.min_height))
            height = math.ceil(
                height / force_divisible_by) * force_divisible_by
            height = int(min(height, self.max_height))
            height = int(height / force_divisible_by) * force_divisible_by

            # keep the origin aspect ratio
            if self.force_original_aspect_ratio == 'increase':
                if width / height < origin_ratio:
                    width = height * origin_ratio
                elif width / height > origin_ratio:
                    height = width / origin_ratio
            elif self.force_original_aspect_ratio == 'decrease':
                if width / height < origin_ratio:
                    height = width / origin_ratio
                elif width / height > origin_ratio:
                    width = height * origin_ratio
            width = int(round(width / force_divisible_by)) * force_divisible_by
            height = int(round(
                height / force_divisible_by)) * force_divisible_by

            # resize
            resized_video_key = transfer_filename(video_key, OP_NAME,
                                                  **self._init_parameters)
            if (not os.path.exists(resized_video_key)
                    or resized_video_key not in loaded_video_keys):
                args = ['-nostdin', '-v', 'quiet',
                        '-y']  # close the ffmpeg log
                stream = ffmpeg.input(video_key)
                stream = stream.filter('scale', width=width, height=height)
                stream = stream.output(resized_video_key).global_args(*args)
                stream.run()

            loaded_video_keys[index] = resized_video_key

        # when the file is modified, its source file needs to be updated.
        for i, value in enumerate(sample[self.video_key]):
            if sample[Fields.source_file][i] != value:
                if loaded_video_keys[i] != value:
                    sample[Fields.source_file][i] = value

        sample[self.video_key] = loaded_video_keys
        return sample