From 56a0175cc290bdf9328e7d73369a111719f68588 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Tue, 12 Dec 2023 11:26:44 +0800 Subject: [PATCH 1/4] fix: move the img save out of loop Signed-off-by: Frost Ming --- hy_tg.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/hy_tg.py b/hy_tg.py index e2b96dc..e4a2e58 100644 --- a/hy_tg.py +++ b/hy_tg.py @@ -70,9 +70,7 @@ def compress_image(input_path, output_path, target_size): (int(width * factor), int(height * factor)), PIL.Image.Resampling.LANCZOS, ) - if sizeof_image(img) <= target_bytes: - img.save(output_path, format="JPEG", quality=quality) - return + img.save(output_path, format="JPEG", quality=quality) def draw_pretty_map(location, file_name, style): From 3007ed67a9874ef2ef9c17b5656b3f4df77c36a1 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Tue, 12 Dec 2023 11:50:19 +0800 Subject: [PATCH 2/4] fix: don't use local files Signed-off-by: Frost Ming --- hy_tg.py | 44 ++++++++++++++++++++++++++------------------ 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/hy_tg.py b/hy_tg.py index e4a2e58..b9f2ed5 100644 --- a/hy_tg.py +++ b/hy_tg.py @@ -2,7 +2,9 @@ import argparse import gc import io import random +import shutil import subprocess +import traceback import numpy as np import PIL @@ -16,8 +18,6 @@ from telebot import TeleBot # type: ignore from telebot.types import Message # type: ignore PIL.Image.MAX_IMAGE_PIXELS = 933120000 -file_in = "map.jpg" -file_out = "map_out.jpg" class Plot(PrettyPlot): @@ -57,28 +57,34 @@ def sizeof_image(image): return buff.tell() -def compress_image(input_path, output_path, target_size): +def compress_image(input_image, target_size): quality = 95 factor = 1.0 - with Image.open(input_path) as img: - target_bytes = 10 * 1024 * 1024 - - while sizeof_image(img) > target_bytes: + input_image.seek(0) + output = io.BytesIO() + with Image.open(input_image) as img: + while sizeof_image(img) > target_size: factor -= 0.05 width, height = img.size img = img.resize( (int(width * factor), int(height * factor)), PIL.Image.Resampling.LANCZOS, ) - img.save(output_path, format="JPEG", quality=quality) + img.save(output, format="JPEG", quality=quality) + output.seek(0) + return output -def draw_pretty_map(location, file_name, style): +def draw_pretty_map(location, style): aoi = get_aoi(address=location, radius=1100, rectangular=True) df = get_osm_geometries(aoi=aoi) fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all() - fig.savefig(file_name) - compress_image(file_in, file_out, 9) # telegram tog need png less than 10MB + buffer = io.BytesIO() + fig.savefig(buffer, format="jpeg") + return compress_image( + buffer, + 10 * 1024 * 1024, # telegram tog need png less than 10MB + ) def main(): @@ -144,16 +150,18 @@ def main(): style = random.choice(styles_list) try: # TODO why this memory leak? - draw_pretty_map(location, file_in, style) + out_image = draw_pretty_map(location, style) # tg can only send image less than 10MB - with open(file_out, "rb") as photo: - bot.send_photo( - message.chat.id, photo, reply_to_message_id=message.message_id - ) + with open("map_out.jpg", "wb") as f: + shutil.copyfileobj(out_image, f) + out_image.seek(0) + bot.send_photo( + message.chat.id, out_image, reply_to_message_id=message.message_id + ) - except Exception as e: + except Exception: + traceback.print_exc() bot.reply_to(message, "Something wrong please check") - print(str(e)) bot.delete_message(reply_message.chat.id, reply_message.message_id) # we need this, fuck it gc.collect() From 23484f5b7cffa5185500af1c7e20daa8a2d356ea Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Tue, 12 Dec 2023 12:07:48 +0800 Subject: [PATCH 3/4] fix: use spooled temp file to reduce memory usage Signed-off-by: Frost Ming --- hy_tg.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/hy_tg.py b/hy_tg.py index b9f2ed5..16c91d2 100644 --- a/hy_tg.py +++ b/hy_tg.py @@ -1,10 +1,10 @@ import argparse import gc -import io import random import shutil import subprocess import traceback +from tempfile import SpooledTemporaryFile import numpy as np import PIL @@ -18,6 +18,7 @@ from telebot import TeleBot # type: ignore from telebot.types import Message # type: ignore PIL.Image.MAX_IMAGE_PIXELS = 933120000 +MAX_IN_MEMORY = 10 * 1024 * 1024 # 10MiB class Plot(PrettyPlot): @@ -52,16 +53,15 @@ class Plot(PrettyPlot): def sizeof_image(image): - with io.BytesIO() as buff: - image.save(buff, format="JPEG", quality=95) - return buff.tell() + with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as f: + image.save(f, format="JPEG", quality=95) + return f.tell() def compress_image(input_image, target_size): quality = 95 factor = 1.0 - input_image.seek(0) - output = io.BytesIO() + output = SpooledTemporaryFile(max_size=MAX_IN_MEMORY) with Image.open(input_image) as img: while sizeof_image(img) > target_size: factor -= 0.05 @@ -79,12 +79,13 @@ def draw_pretty_map(location, style): aoi = get_aoi(address=location, radius=1100, rectangular=True) df = get_osm_geometries(aoi=aoi) fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all() - buffer = io.BytesIO() - fig.savefig(buffer, format="jpeg") - return compress_image( - buffer, - 10 * 1024 * 1024, # telegram tog need png less than 10MB - ) + with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as buffer: + fig.savefig(buffer, format="jpeg") + buffer.seek(0) + return compress_image( + buffer, + 10 * 1024 * 1024, # telegram tog need png less than 10MB + ) def main(): @@ -158,6 +159,7 @@ def main(): bot.send_photo( message.chat.id, out_image, reply_to_message_id=message.message_id ) + out_image.close() except Exception: traceback.print_exc() From 82b209fb229663d9ff9bc20ec207677cf56ea1fd Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Tue, 12 Dec 2023 14:33:50 +0800 Subject: [PATCH 4/4] fix: create tempfile in the outer scope Signed-off-by: Frost Ming --- hy_tg.py | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/hy_tg.py b/hy_tg.py index 16c91d2..4d28a13 100644 --- a/hy_tg.py +++ b/hy_tg.py @@ -58,10 +58,9 @@ def sizeof_image(image): return f.tell() -def compress_image(input_image, target_size): +def compress_image(input_image, output_image, target_size): quality = 95 factor = 1.0 - output = SpooledTemporaryFile(max_size=MAX_IN_MEMORY) with Image.open(input_image) as img: while sizeof_image(img) > target_size: factor -= 0.05 @@ -70,20 +69,20 @@ def compress_image(input_image, target_size): (int(width * factor), int(height * factor)), PIL.Image.Resampling.LANCZOS, ) - img.save(output, format="JPEG", quality=quality) - output.seek(0) - return output + img.save(output_image, format="JPEG", quality=quality) + output_image.seek(0) -def draw_pretty_map(location, style): +def draw_pretty_map(location, style, output_file): aoi = get_aoi(address=location, radius=1100, rectangular=True) df = get_osm_geometries(aoi=aoi) fig = Plot(df=df, aoi_bounds=aoi.bounds, draw_settings=STYLES[style]).plot_all() with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as buffer: fig.savefig(buffer, format="jpeg") buffer.seek(0) - return compress_image( + compress_image( buffer, + output_file, 10 * 1024 * 1024, # telegram tog need png less than 10MB ) @@ -151,15 +150,15 @@ def main(): style = random.choice(styles_list) try: # TODO why this memory leak? - out_image = draw_pretty_map(location, style) - # tg can only send image less than 10MB - with open("map_out.jpg", "wb") as f: - shutil.copyfileobj(out_image, f) - out_image.seek(0) - bot.send_photo( - message.chat.id, out_image, reply_to_message_id=message.message_id - ) - out_image.close() + with SpooledTemporaryFile(max_size=MAX_IN_MEMORY) as out_image: + draw_pretty_map(location, style, out_image) + # tg can only send image less than 10MB + with open("map_out.jpg", "wb") as f: # for debug + shutil.copyfileobj(out_image, f) + out_image.seek(0) + bot.send_photo( + message.chat.id, out_image, reply_to_message_id=message.message_id + ) except Exception: traceback.print_exc()