mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-06-02 00:57:41 +08:00
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.
This commit is contained in:
@@ -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
common/utils.py
Normal file
34
common/utils.py
Normal 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
|
||||
Reference in New Issue
Block a user