Skip to content

Commit

Permalink
feat: add image compression
Browse files Browse the repository at this point in the history
Add image compression feature to WechatComAppChannel to compress images larger than 10MB before uploading to WeChat server. The compression is done using the `compress_imgfile` function in `utils.py`. The `fsize` function is also added to `utils.py` to calculate the size of a file or buffer.
  • Loading branch information
lanvent committed Apr 21, 2023
1 parent 3e9e8d4 commit 6650017
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 0 deletions.
18 changes: 18 additions & 0 deletions channel/wechatcom/wechatcomapp_channel.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from channel.wechatcom.wechatcomapp_message import WechatComAppMessage
from common.log import logger
from common.singleton import singleton
from common.utils import compress_imgfile, fsize
from config import conf
from voice.audio_convert import any_to_amr

Expand Down Expand Up @@ -81,13 +82,22 @@ def send(self, reply: Reply, context: Context):
image_storage = io.BytesIO()
for block in pic_res.iter_content(1024):
image_storage.write(block)
if (sz := fsize(image_storage)) >= 10 * 1024 * 1024:
logger.info(
"[wechatcom] image too large, ready to compress, sz={}".format(sz)
)
image_storage = compress_imgfile(image_storage, 10 * 1024 * 1024 - 1)
logger.info(
"[wechatcom] image compressed, sz={}".format(fsize(image_storage))
)
image_storage.seek(0)
try:
response = self.client.media.upload("image", image_storage)
logger.debug("[wechatcom] upload image response: {}".format(response))
except WeChatClientException as e:
logger.error("[wechatcom] upload image failed: {}".format(e))
return

self.client.message.send_image(
self.agent_id, receiver, response["media_id"]
)
Expand All @@ -96,6 +106,14 @@ def send(self, reply: Reply, context: Context):
)
elif reply.type == ReplyType.IMAGE: # 从文件读取图片
image_storage = reply.content
if (sz := fsize(image_storage)) >= 10 * 1024 * 1024:
logger.info(
"[wechatcom] image too large, ready to compress, sz={}".format(sz)
)
image_storage = compress_imgfile(image_storage, 10 * 1024 * 1024 - 1)
logger.info(
"[wechatcom] image compressed, sz={}".format(fsize(image_storage))
)
image_storage.seek(0)
try:
response = self.client.media.upload("image", image_storage)
Expand Down
34 changes: 34 additions & 0 deletions common/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
import io
import os

from PIL import Image


def fsize(file):
if isinstance(file, io.BytesIO):
return file.getbuffer().nbytes
elif isinstance(file, str):
return os.path.getsize(file)
elif hasattr(file, "seek") and hasattr(file, "tell"):
pos = file.tell()
file.seek(0, os.SEEK_END)
size = file.tell()
file.seek(pos)
return size
else:
raise TypeError("Unsupported type")


def compress_imgfile(file, max_size):
if fsize(file) <= max_size:
return file
file.seek(0)
img = Image.open(file)
rgb_image = img.convert("RGB")
quality = 95
while True:
out_buf = io.BytesIO()
rgb_image.save(out_buf, "JPEG", quality=quality)
if fsize(out_buf) <= max_size:
return out_buf
quality -= 5

0 comments on commit 6650017

Please sign in to comment.