Browse Source

feat: add image compression

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.
master
lanvent 1 year ago
parent
commit
665001732b
2 changed files with 52 additions and 0 deletions
  1. +18
    -0
      channel/wechatcom/wechatcomapp_channel.py
  2. +34
    -0
      common/utils.py

+ 18
- 0
channel/wechatcom/wechatcomapp_channel.py View File

@@ -17,6 +17,7 @@ from channel.chat_channel import ChatChannel
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

@@ -81,6 +82,14 @@ class WechatComAppChannel(ChatChannel):
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)
@@ -88,6 +97,7 @@ class WechatComAppChannel(ChatChannel):
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"]
)
@@ -96,6 +106,14 @@ class WechatComAppChannel(ChatChannel):
)
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)


+ 34
- 0
common/utils.py View File

@@ -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

Loading…
Cancel
Save