优化首帧功能,如果没传入首帧,则提取首帧。

This commit is contained in:
liuwei
2025-05-29 16:36:51 +08:00
parent a01b9e444e
commit 31c062fa70
2 changed files with 56 additions and 3 deletions

48
utils/video_utils.py Normal file
View File

@@ -0,0 +1,48 @@
import os
import cv2
from loguru import logger
def _get_first_frame(video_path, output_path):
"""
提取视频的第一帧并保存为图片
:param video_path: 视频文件路径
:param output_path: 输出图片路径
:return: 输出图片的绝对路径如果失败则返回None
"""
cap = None
try:
logger.info(f"开始提取视频首帧: {video_path}")
# 打开视频文件
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
logger.error(f"无法打开视频: {video_path}")
return None
# 读取首帧
ret, frame = cap.read()
if not ret:
logger.error("无法读取视频帧")
cap.release()
return None
# 保存首帧为图片
try:
cv2.imwrite(output_path, frame)
logger.info(f"首帧已保存为: {output_path}")
except Exception as e:
logger.error(f"保存首帧图片失败: {e}")
cap.release()
return None
# 释放资源
cap.release()
return os.path.abspath(output_path)
except Exception as e:
logger.error(f"提取视频首帧时出错: {e}")
return None
finally:
cap.release()

View File

@@ -1,6 +1,7 @@
import asyncio import asyncio
import base64 import base64
import os import os
import time
from asyncio import Future from asyncio import Future
from asyncio import Queue, sleep from asyncio import Queue, sleep
from io import BytesIO from io import BytesIO
@@ -14,6 +15,7 @@ from pymediainfo import MediaInfo
import pysilk import pysilk
from pydub import AudioSegment from pydub import AudioSegment
from utils.video_utils import _get_first_frame
from wechat_ipad import UserLoggedOut from wechat_ipad import UserLoggedOut
from wechat_ipad.client.base import WechatAPIClientBase from wechat_ipad.client.base import WechatAPIClientBase
@@ -214,8 +216,7 @@ class MessageMixin(WechatAPIClientBase):
ValueError: 视频或图片参数都为空或都不为空时 ValueError: 视频或图片参数都为空或都不为空时
根据error_handler处理错误 根据error_handler处理错误
""" """
if not image:
image = Path(os.path.join(Path(__file__).resolve().parent, "fallback.png"))
# get video base64 and duration # get video base64 and duration
if isinstance(video, str): if isinstance(video, str):
vid_base64 = video vid_base64 = video
@@ -235,6 +236,9 @@ class MessageMixin(WechatAPIClientBase):
file_len = len(video_bytes) file_len = len(video_bytes)
vid_base64 = base64.b64encode(video_bytes).decode() vid_base64 = base64.b64encode(video_bytes).decode()
media_info = MediaInfo.parse(video_path) media_info = MediaInfo.parse(video_path)
first_frame = _get_first_frame(video_path, f"dify_frame_{int(time.time())}.jpg")
if first_frame:
image = first_frame
else: else:
raise ValueError("video should be str, bytes, or path") raise ValueError("video should be str, bytes, or path")
# 获取视频时长 # 获取视频时长
@@ -246,7 +250,8 @@ class MessageMixin(WechatAPIClientBase):
if duration is None: if duration is None:
duration = 1 duration = 1
self.logging.error(f"无法从视频文件获取时长: {video}") self.logging.error(f"无法从视频文件获取时长: {video}")
if not image:
image = Path(os.path.join(Path(__file__).resolve().parent, "fallback.png"))
# get image base64 # get image base64
if isinstance(image, str): if isinstance(image, str):
image_base64 = image image_base64 = image