优化首帧功能,如果没传入首帧,则提取首帧。
This commit is contained in:
48
utils/video_utils.py
Normal file
48
utils/video_utils.py
Normal 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()
|
||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user