From 8da75ba3296db680720c3f45475bd21ad1c1d084 Mon Sep 17 00:00:00 2001
From: liuwei <liuwei@wdtrgf.com.cn>
Date: Thu, 29 May 2025 16:47:39 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E9=A6=96=E5=B8=A7=E5=8A=9F?=
 =?UTF-8?q?=E8=83=BD=EF=BC=8C=E5=A6=82=E6=9E=9C=E6=B2=A1=E4=BC=A0=E5=85=A5?=
 =?UTF-8?q?=E9=A6=96=E5=B8=A7=EF=BC=8C=E5=88=99=E6=8F=90=E5=8F=96=E9=A6=96?=
 =?UTF-8?q?=E5=B8=A7=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 utils/video_utils.py          | 23 +++++++++++++++++++++++
 wechat_ipad/client/message.py |  9 +++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/utils/video_utils.py b/utils/video_utils.py
index c593c93..94a90d2 100644
--- a/utils/video_utils.py
+++ b/utils/video_utils.py
@@ -1,4 +1,6 @@
+import io
 import os
+import tempfile
 
 import cv2
 
@@ -46,3 +48,24 @@ def _get_first_frame(video_path, output_path):
         return None
     finally:
         cap.release()
+
+
+def _get_first_frame_bytes(video_bytes, output_path):
+    """
+    提取视频的第一帧并保存为图片
+    :param video_bytes: 视频文件流
+    :param output_path: 输出图片路径
+    :return: 输出图片的绝对路径，如果失败则返回None
+    """
+    try:
+        logger.info(f"开始提取视频流首帧")
+        # 创建临时文件
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
+            # 将字节流写入临时文件
+            temp_file.write(video_bytes)
+            temp_file_path = temp_file.name  # 获取临时文件路径
+        # 打开视频文件
+        _get_first_frame(temp_file_path, output_path)
+    except Exception as e:
+        logger.error(f"提取视频首帧时出错: {e}")
+        return None
diff --git a/wechat_ipad/client/message.py b/wechat_ipad/client/message.py
index 7cb83e9..b5ef773 100644
--- a/wechat_ipad/client/message.py
+++ b/wechat_ipad/client/message.py
@@ -15,7 +15,7 @@ from pymediainfo import MediaInfo
 import pysilk
 from pydub import AudioSegment
 
-from utils.video_utils import _get_first_frame
+from utils.video_utils import _get_first_frame, _get_first_frame_bytes
 from wechat_ipad import UserLoggedOut
 from wechat_ipad.client.base import WechatAPIClientBase
 
@@ -231,6 +231,11 @@ class MessageMixin(WechatAPIClientBase):
             vid_base64 = base64.b64encode(video).decode()
             file_len = len(video)
             media_info = MediaInfo.parse(BytesIO(video))
+            # 如果没有传入首帧，则自己提取一次
+            if not has_image:
+                first_frame = _get_first_frame_bytes(video, f"frame_{int(time.time())}.jpg")
+                if first_frame:
+                    image = first_frame
         elif isinstance(video, os.PathLike):
             video_path = Path(video)
             if not video_path.exists():
@@ -242,7 +247,7 @@ class MessageMixin(WechatAPIClientBase):
             media_info = MediaInfo.parse(video_path)
             # 如果没有传入首帧，则自己提取一次
             if not has_image:
-                first_frame = _get_first_frame(video_path, f"dify_frame_{int(time.time())}.jpg")
+                first_frame = _get_first_frame(video_path, f"frame_{int(time.time())}.jpg")
                 if first_frame:
                     image = first_frame
         else: