cdryzun/tg_bot_collections
1
0
Fork 0
mirror of https://github.com/cdryzun/tg_bot_collections.git synced 2025-11-04 08:46:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: use spooled temp file to reduce memory usage

Browse Source 
Signed-off-by: Frost Ming <me@frostming.com>
This commit is contained in:
Frost Ming
2023-12-12 12:07:48 +08:00
parent 3007ed67a9
commit 23484f5b7c
1 changed files with 14 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
hy_tg.py
View File

@ -1,10 +1,10 @@
import argparse import argparse
import gc import gc
import io
import random import random
import shutil import shutil
import subprocess import subprocess
import traceback import traceback
from tempfile import SpooledTemporaryFile
import numpy as np import numpy as np
import PIL import PIL
@ -18,6 +18,7 @@ from telebot import TeleBot # type: ignore
from telebot.types import Message # type: ignore from telebot.types import Message # type: ignore
PIL.Image.MAX_IMAGE_PIXELS = 933120000 PIL.Image.MAX_IMAGE_PIXELS = 933120000
MAX_IN_MEMORY = 10 * 1024 * 1024 # 10MiB
class Plot(PrettyPlot): class Plot(PrettyPlot):
@ -52,16 +53,15 @@ class Plot(PrettyPlot):
def sizeof_image(image): def sizeof_image(image):
with io.BytesIO() as buff: with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as f:
image.save(buff, format="JPEG", quality=95) image.save(f, format="JPEG", quality=95)
return buff.tell() return f.tell()
def compress_image(input_image, target_size): def compress_image(input_image, target_size):
quality = 95 quality = 95
factor = 1.0 factor = 1.0
input_image.seek(0) output = SpooledTemporaryFile(max_size=MAX_IN_MEMORY)
output = io.BytesIO()
with Image.open(input_image) as img: with Image.open(input_image) as img:
while sizeof_image(img) > target_size: while sizeof_image(img) > target_size:
factor -= 0.05 factor -= 0.05
@ -79,8 +79,9 @@ def draw_pretty_map(location, style):
aoi = get_aoi(address=location, radius=1100, rectangular=True) aoi = get_aoi(address=location, radius=1100, rectangular=True)
df = get_osm_geometries(aoi=aoi) df = get_osm_geometries(aoi=aoi)
fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all() fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all()
buffer = io.BytesIO() with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as buffer:
fig.savefig(buffer, format="jpeg") fig.savefig(buffer, format="jpeg")
buffer.seek(0)
return compress_image( return compress_image(
buffer, buffer,
10 * 1024 * 1024, # telegram tog need png less than 10MB 10 * 1024 * 1024, # telegram tog need png less than 10MB
@ -158,6 +159,7 @@ def main():
bot.send_photo( bot.send_photo(
message.chat.id, out_image, reply_to_message_id=message.message_id message.chat.id, out_image, reply_to_message_id=message.message_id
) )
out_image.close()
except Exception: except Exception:
traceback.print_exc() traceback.print_exc()