from logging import getLogger
from pathlib import Path
from time import time
from typing import Any
import ffmpeg
from rich.prompt import Prompt
from utils.log import configure_logging
PATH = Path(__file__).parent
INPUT_DIR = PATH / 'input'
OUTPUT_DIR = PATH / 'output'
OUTPUT_ARGS: dict[str, Any] = {
'row-mt': 1,
'loglevel': 'warning',
'max-intra-rate': 100,
'quality': 'default',
'auto-alt-ref': 1,
}
logger = getLogger(__name__)
def compress(input_file: Path, output_file: Path, sticker_width: int, max_size: int):
file_input = ffmpeg.input(input_file)
args = OUTPUT_ARGS.copy()
stream = file_input.fps(fps=25, round='up')
stream = stream.scale(w=sticker_width, h=-1, force_original_aspect_ratio='decrease')
stream = stream.crop(
out_w=f'min(min(iw,ih),{sticker_width})', out_h=f'min(min(iw,ih),{sticker_width})'
)
if input_file.suffix in ('.gif', '.mp4'):
duration = float(ffmpeg.probe(input_file)['streams'][0]['duration'])
max_bitrate = int(max_size * 8 / duration) * .95
speed_factor = 1
if duration > 3:
speed_factor = duration / 3
stream = stream.setpts(expr=f'PTS*{speed_factor}')
args = args | {
'minrate': f'{(max_bitrate * .2):.2f}k',
'maxrate': f'{max_bitrate:.2f}k',
'b:v': f'{(max_bitrate * .65):.2f}k',
}
else:
args = args | {'crf': 0}
stream = stream.output(
filename=output_file,
vcodec='libvpx-vp9',
pix_fmt='yuva420p',
an=True,
extra_options=args,
)
stream.run(overwrite_output=True)
file_size = output_file.stat().st_size
if file_size / 1024 > max_size:
print(stream.compile_line())
logger.info(
f'Average bitrate: {(int(ffmpeg.probe(output_file)['format']['bit_rate']) / 1024):.2f} KB/s'
)
logger.error(
f'File size ({file_size / 1024:.2f} KB) exceeds the maximum size ({max_size} KB).'
)
exit(1)
def main():
if not INPUT_DIR.exists():
INPUT_DIR.mkdir(parents=True)
if not OUTPUT_DIR.exists():
OUTPUT_DIR.mkdir(parents=True)
for file in OUTPUT_DIR.iterdir():
file.unlink()
sticker_type = Prompt.ask(
'Do you want to create stickers or emoji or icon?', choices=['s', 'e', 'i']
)
match sticker_type:
case 's':
width = 512
size = 256
case 'e':
width = 100
size = 256
case 'i':
width = 100
size = 32
case _:
width = 512
size = 256
started_at = time()
files = list(INPUT_DIR.iterdir())
html = '\n'
for i, file in enumerate(files):
logger.info(f'Processing [{i+1}/{len(list(files))}] {file.name}...')
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.relative_to(PATH), file_output.relative_to(PATH), width, size)
html += f'\n'
logger.info(f'Finished in {round(time() - started_at, 2)} seconds')
html += '\n'
html += '\n'
with open(OUTPUT_DIR / 'index.html', 'w', encoding='utf-8') as f:
f.write(html)
if __name__ == '__main__':
configure_logging()
main()