Апдейт

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')