Files
TelegramStickers/main.py
2025-01-31 22:14:59 +03:00

136 lines
3.6 KiB
Python

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) * 0.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 = '<html><head><style>body { background-color: gray; }</style></head><body>\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'<video width="250" height="250" src="{file.name.replace(file.suffix, '.webm')}" autoplay loop></video>\n'
logger.info(f'Finished in {round(time() - started_at, 2)} seconds')
html += '</body>\n'
html += '</html>\n'
with open(OUTPUT_DIR / 'index.html', 'w', encoding='utf-8') as f:
f.write(html)
if __name__ == '__main__':
configure_logging()
main()