diff --git a/main.py b/main.py index 61902a6..d6838f3 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,9 @@ from logging import getLogger from pathlib import Path from time import time +from typing import Any import ffmpeg -from ffmpeg.nodes import FilterableStream from rich.prompt import Prompt from utils.log import configure_logging @@ -12,71 +12,72 @@ PATH = Path(__file__).parent INPUT_DIR = PATH / 'input' OUTPUT_DIR = PATH / 'output' -OUTPUT_ARGS = { - 'vcodec': 'libvpx-vp9', - 'pix_fmt': 'yuva420p', +OUTPUT_ARGS: dict[str, Any] = { 'row-mt': 1, - 'format': 'webm', 'loglevel': 'error', - 'an': None, - 'y': None, } logger = getLogger(__name__) -def compress(input_file: Path, output_file: Path, sticker_width: int, bitrate: int | None = None): +def compress(input_file: Path, output_file: Path, sticker_width: int, max_size: int): file_input = ffmpeg.input(input_file) file_output = OUTPUT_DIR / input_file.name.replace(input_file.suffix, '.webm') - if not isinstance(file_input, FilterableStream): - logger.error(f'Failed to process file: {input_file}') - raise Exception('Failed to process file') - - stream = ffmpeg.filter(file_input, 'fps', fps=25, round='up') - stream = ffmpeg.filter( - stream, 'scale', sticker_width, -1, force_original_aspect_ratio='decrease' + stream = file_input.fps(fps=25, round='up') + stream = stream.scale(w=sticker_width, h=-1, force_original_aspect_ratio='decrease') + stream = stream.loop(loop=0) + stream = stream.crop( + out_w=f'min(min(iw,ih),{sticker_width})', out_h=f'min(min(iw,ih),{sticker_width})' ) - stream = ffmpeg.filter( - stream, 'pad', sticker_width, sticker_width, '(ow-iw)/2', '(oh-ih)/2', color='0x00000000' - ) - stream = ffmpeg.filter(stream, 'loop', 0, 1) + stream = stream.pad(color='0x00000000') - if input_file.suffix == '.gif': + if input_file.suffix in ('.gif', '.mp4'): duration = float(ffmpeg.probe(input_file)['streams'][0]['duration']) speed_factor = 1 if duration > 3: speed_factor = duration / 3 - stream = ffmpeg.filter(stream, 'setpts', f'PTS/{speed_factor}') + stream = stream.setpts(expr=f'PTS*{speed_factor}') - args = OUTPUT_ARGS - - if bitrate: - args['video_bitrate'] = f'{bitrate}k' - - output = ffmpeg.output(stream, filename=file_output, **args) + output = ffmpeg.output( + stream, + filename=file_output, + vcodec='libvpx-vp9', + pix_fmt='yuva420p', + extra_options=OUTPUT_ARGS, + ) try: - ffmpeg.run(output) - except Exception: + output.run(overwrite_output=True) + except Exception as e: + logger.error(f"Error during first encoding: {e}") exit(1) file_size = file_output.stat().st_size - if file_size / 1024 > 256: - if not bitrate: - bitrate = int(int(ffmpeg.probe(output_file)['format']['bit_rate']) / 1000) + if file_size / 1024 > max_size: + initial_duration = float(ffmpeg.probe(input_file)['format']['duration']) + bitrate = int(max_size * 8 / initial_duration * .65) - bitrate = int((256 * 1024) / (file_size / bitrate)) + temp_output = file_output.with_suffix('.tmp.webm') + stream = ffmpeg.input(output_file, c='libvpx-vp9') - logger.info( - f'File size is too high [{int(file_size / 1024)} KB], reducing to {bitrate}k...' + output = ffmpeg.output( + stream, + filename=temp_output, + vcodec='libvpx-vp9', + pix_fmt='yuva420p', + b=f'{bitrate}k', + extra_options=OUTPUT_ARGS, ) - return compress(input_file, output_file, sticker_width, bitrate) + output.run(overwrite_output=True) + + file_output.unlink() + temp_output.rename(output_file) def main(): @@ -89,29 +90,37 @@ def main(): for file in OUTPUT_DIR.iterdir(): file.unlink() - sticker_type = Prompt.ask('Do you want to create stickers or emoji?', choices=['s', 'e']) + sticker_type = Prompt.ask( + 'Do you want to create stickers or emoji or icon?', choices=['s', 'e', 'i'] + ) match sticker_type: case 's': - STICKER_WIDTH = 512 + width = 512 + size = 256 case 'e': - STICKER_WIDTH = 100 + width = 100 + size = 256 + case 'i': + width = 100 + size = 32 case _: - STICKER_WIDTH = 512 + width = 512 + size = 256 started_at = time() for file in INPUT_DIR.iterdir(): logger.info(f'Processing {file.name}...') - if file.suffix not in ('.png', '.jpg', '.webp', '.gif'): + if file.suffix not in ('.png', '.jpg', '.webp', '.gif', '.mp4'): logger.warning(f'Unsupported file type: {file.suffix}') continue file_path = INPUT_DIR / file.name file_output = OUTPUT_DIR / file.name.replace(file.suffix, '.webm') - compress(file_path, file_output, STICKER_WIDTH) + compress(file_path, file_output, width, size) logger.info(f'Finished in {round(time() - started_at, 2)} seconds') diff --git a/pyproject.toml b/pyproject.toml index 9b06dd9..c045194 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,8 +5,8 @@ description = "A simple tool to convert any images and/or gifs to webm extension readme = "README.md" requires-python = ">=3.13" dependencies = [ - "ffmpeg-python==0.2.0", "rich==13.9.4", + "typed-ffmpeg==2.6.0", ] [dependency-groups] diff --git a/utils/log.py b/utils/log.py index 906ead1..a82512c 100644 --- a/utils/log.py +++ b/utils/log.py @@ -13,3 +13,5 @@ def configure_logging(): datefmt=DATE_FORMAT, handlers=[RichHandler(show_time=False, rich_tracebacks=True)], ) + + logging.getLogger('ffmpeg').setLevel(logging.ERROR)