加入首帧提取功能

plugins/dify/main.py

@@ -1,4 +1,6 @@
 import os
+
+import cv2
 import requests
 import json
 import time
@@ -159,7 +161,8 @@ class DifyPlugin(MessagePluginInterface):
                         if file_type == 1:
                             await bot.send_image_message((roomid if roomid else sender), Path(response))
                         elif file_type == 2:
-                            await bot.send_video_message((roomid if roomid else sender), Path(response))
+                            first_farme = self._get_first_frame(response, f"dify_frame_{int(time.time())}.jpg")
+                            await bot.send_video_message((roomid if roomid else sender), Path(response), Path(first_farme))
                         else:
                             return False, "获取媒资失败"
                     else:
@@ -233,7 +236,9 @@ class DifyPlugin(MessagePluginInterface):
                     if file_type == 1:
                         await bot.send_image_message((roomid if roomid else sender), Path(response))
                     elif file_type == 2:
-                        await bot.send_video_message((roomid if roomid else sender), Path(response))
+                        # 提取首帧
+                        first_farme = self._get_first_frame(response, f"dify_frame_{int(time.time())}.jpg")
+                        await bot.send_video_message((roomid if roomid else sender), Path(response), Path(first_farme))
                     else:
                         return False, "获取媒资失败"
                 else:
@@ -486,3 +491,42 @@ class DifyPlugin(MessagePluginInterface):
                 return 0
         else:
             return 0
+
+    def _get_first_frame(self, video_path, output_path):
+        """
+        提取视频的第一帧并保存为图片
+        :param video_path: 视频文件路径
+        :param output_path: 输出图片路径
+        :return: 输出图片的绝对路径，如果失败则返回None
+        """
+        try:
+            self.LOG.info(f"开始提取视频首帧: {video_path}")
+            # 打开视频文件
+            cap = cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                self.LOG.error(f"无法打开视频: {video_path}")
+                return None
+
+            # 读取首帧
+            ret, frame = cap.read()
+            if not ret:
+                self.LOG.error("无法读取视频帧")
+                cap.release()
+                return None
+
+            # 保存首帧为图片
+            try:
+                cv2.imwrite(output_path, frame)
+                self.LOG.info(f"首帧已保存为: {output_path}")
+            except Exception as e:
+                self.LOG.error(f"保存首帧图片失败: {e}")
+                cap.release()
+                return None
+
+            # 释放资源
+            cap.release()
+            return os.path.abspath(output_path)
+
+        except Exception as e:
+            self.LOG.error(f"提取视频首帧时出错: {e}")
+            return None