<commit_before>import os
import shutil

from taichi._lib.utils import get_os_name
from taichi.tools.image import imwrite

FRAME_FN_TEMPLATE = '%06d.png'
FRAME_DIR = 'frames'

# Write the frames to the disk and then make videos (mp4 or gif) if necessary


def scale_video(input_fn, output_fn, ratiow, ratioh):
    os.system(
        f'ffmpeg -i {input_fn}  -vf "scale=iw*{ratiow:.4f}:ih*{ratioh:.4f}" {output_fn}'
    )


def crop_video(input_fn, output_fn, x_begin, x_end, y_begin, y_end):
    os.system(
        f'ffmpeg -i {input_fn} -filter:v "crop=iw*{x_end - x_begin:.4f}:ih*{y_end - y_begin:.4f}:iw*{x_begin:0.4f}:ih*{1 - y_end:0.4f}" {output_fn}'
    )


def accelerate_video(input_fn, output_fn, speed):
    os.system(
        f'ffmpeg -i {input_fn} -filter:v "setpts={1 / speed:.4f}*PTS" {output_fn}'
    )


def get_ffmpeg_path():
    return 'ffmpeg'


def mp4_to_gif(input_fn, output_fn, framerate):
    # Generate the palette
    palette_name = 'palette.png'
    if get_os_name() == 'win':
        command = get_ffmpeg_path(
        ) + f" -loglevel panic -i {input_fn} -vf 'palettegen' -y {palette_name}"
    else:
        command = get_ffmpeg_path(
        ) + f" -loglevel panic -i {input_fn} -vf 'fps={framerate}," \
            f"scale=320:640:flags=lanczos,palettegen' -y {palette_name}"
    # print command
    os.system(command)

    # Generate the GIF
    command = get_ffmpeg_path(
    ) + f" -loglevel panic -i {input_fn} -i {palette_name} -lavfi paletteuse -y {output_fn}"
    # print command
    os.system(command)
    os.remove(palette_name)


class VideoManager:
    """Utility class for exporting results to `mp4` and `gif` formats.
    This class relies on `ffmpeg`.

    Args:
        output_dir (str): directory to save the frames.
        width, height (int): resolution of the video.
        post_processor (object): any object that implements the `process(img)`
            method, which accepts an image as a `numpy.ndarray` and returns
            the process image.
        framerate (int): frame rate of the video.
        automatic_build (bool): automatically generate the resulting video or not.

    Example::

        >>> video_manager = ti.tools.VideoManager(output_dir="./output", framerate=24, automatic_build=False)
        >>> for i in range(50):
        >>>     render()
        >>>     img = pixels.to_numpy()
        >>>     video_manager.write_frame(img)
        >>>
        >>> video_manager.make_video(gif=True, mp4=True)

    Returns:
        An instance of :class:`taichi.tools.VideoManager` class.
    """
    def __init__(self,
                 output_dir,
                 width=None,
                 height=None,
                 post_processor=None,
                 framerate=24,
                 automatic_build=True):
        assert (width is None) == (height is None)
        self.width = width
        self.height = height
        self.directory = output_dir
        self.frame_directory = os.path.join(self.directory, FRAME_DIR)
        try:
            os.makedirs(self.frame_directory)
        except:
            pass
        self.next_video_checkpoint = 4
        self.framerate = framerate
        self.post_processor = post_processor
        self.frame_counter = 0
        self.frame_fns = []
        self.automatic_build = automatic_build

    def get_output_filename(self, suffix):
        return os.path.join(self.directory, 'video' + suffix)

    def write_frame(self, img):
        """Write an `numpy.ndarray` `img` to an image file.

        The filename will be automatically determined by this manager
        and the frame counter.
        """
        if img.shape[0] % 2 != 0:
            print('Warning: height is not divisible by 2! Dropping last row')
            img = img[:-1]
        if img.shape[1] % 2 != 0:
            print('Warning: width is not divisible by 2! Dropping last column')
            img = img[:, :-1]
        if self.post_processor:
            img = self.post_processor.process(img)
        if self.width is None:
            self.width = img.shape[0]
            self.height = img.shape[1]
        assert os.path.exists(self.directory)
        fn = FRAME_FN_TEMPLATE % self.frame_counter
        self.frame_fns.append(fn)
        imwrite(img, os.path.join(self.frame_directory, fn))
        self.frame_counter += 1
        if self.frame_counter % self.next_video_checkpoint == 0:
            if self.automatic_build:
                self.make_video()
                self.next_video_checkpoint *= 2

    def get_frame_directory(self):<commit_msg>[misc] Add filename option to ti.tools.VideoManager. (#5219)<commit_after>import os
import shutil

from taichi._lib.utils import get_os_name
from taichi.tools.image import imwrite

FRAME_FN_TEMPLATE = '%06d.png'
FRAME_DIR = 'frames'

# Write the frames to the disk and then make videos (mp4 or gif) if necessary


def scale_video(input_fn, output_fn, ratiow, ratioh):
    os.system(
        f'ffmpeg -i {input_fn}  -vf "scale=iw*{ratiow:.4f}:ih*{ratioh:.4f}" {output_fn}'
    )


def crop_video(input_fn, output_fn, x_begin, x_end, y_begin, y_end):
    os.system(
        f'ffmpeg -i {input_fn} -filter:v "crop=iw*{x_end - x_begin:.4f}:ih*{y_end - y_begin:.4f}:iw*{x_begin:0.4f}:ih*{1 - y_end:0.4f}" {output_fn}'
    )


def accelerate_video(input_fn, output_fn, speed):
    os.system(
        f'ffmpeg -i {input_fn} -filter:v "setpts={1 / speed:.4f}*PTS" {output_fn}'
    )


def get_ffmpeg_path():
    return 'ffmpeg'


def mp4_to_gif(input_fn, output_fn, framerate):
    # Generate the palette
    palette_name = 'palette.png'
    if get_os_name() == 'win':
        command = get_ffmpeg_path(
        ) + f" -loglevel panic -i {input_fn} -vf 'palettegen' -y {palette_name}"
    else:
        command = get_ffmpeg_path(
        ) + f" -loglevel panic -i {input_fn} -vf 'fps={framerate}," \
            f"scale=320:640:flags=lanczos,palettegen' -y {palette_name}"
    # print command
    os.system(command)

    # Generate the GIF
    command = get_ffmpeg_path(
    ) + f" -loglevel panic -i {input_fn} -i {palette_name} -lavfi paletteuse -y {output_fn}"
    # print command
    os.system(command)
    os.remove(palette_name)


class VideoManager:
    """Utility class for exporting results to `mp4` and `gif` formats.
    This class relies on `ffmpeg`.

    Args:
        output_dir (str): directory to save the frames.
        video_filename (str): filename for the video. default filename is video.mp4.
        width, height (int): resolution of the video.
        post_processor (object): any object that implements the `process(img)`
            method, which accepts an image as a `numpy.ndarray` and returns
            the process image.
        framerate (int): frame rate of the video.
        automatic_build (bool): automatically generate the resulting video or not.

    Example::

        >>> video_manager = ti.tools.VideoManager(output_dir="./output", framerate=24, automatic_build=False)
        >>> for i in range(50):
        >>>     render()
        >>>     img = pixels.to_numpy()
        >>>     video_manager.write_frame(img)
        >>>
        >>> video_manager.make_video(gif=True, mp4=True)

    Returns:
        An instance of :class:`taichi.tools.VideoManager` class.
    """
    def __init__(self,
                 output_dir,
                 video_filename=None,
                 width=None,
                 height=None,
                 post_processor=None,
                 framerate=24,
                 automatic_build=True):
        assert (width is None) == (height is None)
        self.width = width
        self.height = height
        self.directory = output_dir
        self.frame_directory = os.path.join(self.directory, FRAME_DIR)
        try:
            os.makedirs(self.frame_directory)
        except:
            pass
        self.video_filename = video_filename
        self.next_video_checkpoint = 4
        self.framerate = framerate
        self.post_processor = post_processor
        self.frame_counter = 0
        self.frame_fns = []
        self.automatic_build = automatic_build

    def get_output_filename(self, suffix):
        if not self.video_filename:
            return os.path.join(self.directory, 'video' + suffix)
        filename, extension = os.path.splitext(self.video_filename)
        if extension is not None:
            print(f'Warning: file extension {extension} will be disregarded!')
        return os.path.join(self.directory, filename + suffix)

    def write_frame(self, img):
        """Write an `numpy.ndarray` `img` to an image file.

        The filename will be automatically determined by this manager
        and the frame counter.
        """
        if img.shape[0] % 2 != 0:
            print('Warning: height is not divisible by 2! Dropping last row')
            img = img[:-1]
        if img.shape[1] % 2 != 0:
            print('Warning: width is not divisible by 2! Dropping last column')
            img = img[:, :-1]
        if self.post_processor:
            img = self.post_processor.process(img)
        if self.width is None:
            self.width = img.shape[0]
            self.height = img.shape[1]
        assert os.path.exists(self.directory)
        fn = FRAME_FN_TEMPLATE % self.frame_counter
        self.frame_fns.append(fn)
        imwrite(img, os.path.join(self.frame_directory, fn))
        self.frame_counter += 1
        if self.frame_counter % self.next_video_checkpoint == 0:
            if self.automatic_build:
                self.make_video()
                self.next_video_checkpoint *= 2

    def get_frame_directory(self):
{'fraction_non_alphanumeric': 0.11060243471022786, 'fraction_numerical': 0.010092602226615337, 'mean_word_length': 2.356145251396648, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4983996', 'n_tokens_mistral': 3143, 'n_tokens_neox': 3027, 'n_words': 764}