优化首帧功能,如果没传入首帧,则提取首帧。
This commit is contained in:
@@ -1,4 +1,6 @@
|
|||||||
|
import io
|
||||||
import os
|
import os
|
||||||
|
import tempfile
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
|
|
||||||
@@ -46,3 +48,24 @@ def _get_first_frame(video_path, output_path):
|
|||||||
return None
|
return None
|
||||||
finally:
|
finally:
|
||||||
cap.release()
|
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
|
||||||
|
|||||||
@@ -15,7 +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 utils.video_utils import _get_first_frame, _get_first_frame_bytes
|
||||||
from wechat_ipad import UserLoggedOut
|
from wechat_ipad import UserLoggedOut
|
||||||
from wechat_ipad.client.base import WechatAPIClientBase
|
from wechat_ipad.client.base import WechatAPIClientBase
|
||||||
|
|
||||||
@@ -231,6 +231,11 @@ class MessageMixin(WechatAPIClientBase):
|
|||||||
vid_base64 = base64.b64encode(video).decode()
|
vid_base64 = base64.b64encode(video).decode()
|
||||||
file_len = len(video)
|
file_len = len(video)
|
||||||
media_info = MediaInfo.parse(BytesIO(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):
|
elif isinstance(video, os.PathLike):
|
||||||
video_path = Path(video)
|
video_path = Path(video)
|
||||||
if not video_path.exists():
|
if not video_path.exists():
|
||||||
@@ -242,7 +247,7 @@ class MessageMixin(WechatAPIClientBase):
|
|||||||
media_info = MediaInfo.parse(video_path)
|
media_info = MediaInfo.parse(video_path)
|
||||||
# 如果没有传入首帧,则自己提取一次
|
# 如果没有传入首帧,则自己提取一次
|
||||||
if not has_image:
|
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:
|
if first_frame:
|
||||||
image = first_frame
|
image = first_frame
|
||||||
else:
|
else:
|
||||||
|
|||||||
Reference in New Issue
Block a user