Przeglądaj źródła

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.
lanvent 3 lat temu
rodzic
commit
665001732b
2 zmienionych plików z 52 dodań i 0 usunięć
  1. 18 0
      channel/wechatcom/wechatcomapp_channel.py
  2. 34 0
      common/utils.py

+ 18 - 0
channel/wechatcom/wechatcomapp_channel.py

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

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