优化IO问题,使用异步方案进行视频下载等操作。
This commit is contained in:
@@ -1,10 +1,10 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import requests
|
||||
import json
|
||||
import time
|
||||
import re # 添加re模块导入
|
||||
import asyncio
|
||||
from typing import Dict, Any, List, Optional, Tuple
|
||||
|
||||
from pathlib import Path
|
||||
@@ -209,7 +209,7 @@ class DifyPlugin(MessagePluginInterface):
|
||||
if file_type == 1:
|
||||
await bot.send_image_message(target, Path(response))
|
||||
elif file_type == 2:
|
||||
first_frame = self._get_first_frame(response, f"dify_frame_{int(time.time())}.jpg")
|
||||
first_frame = await self._get_first_frame(response, f"dify_frame_{int(time.time())}.jpg")
|
||||
await bot.send_video_message(target, Path(response), Path(first_frame))
|
||||
else:
|
||||
return False, "获取媒资失败"
|
||||
@@ -321,12 +321,12 @@ class DifyPlugin(MessagePluginInterface):
|
||||
if outputs["type"] == "image":
|
||||
downloader = MediaDownloader()
|
||||
image_url = outputs["result"]
|
||||
image_path = downloader.download_media(image_url)
|
||||
image_path = await downloader.download_media(image_url)
|
||||
answer = image_path
|
||||
if outputs["type"] == "video":
|
||||
downloader = MediaDownloader()
|
||||
image_url = outputs["result"]
|
||||
image_path = downloader.download_media(image_url)
|
||||
image_path = await downloader.download_media(image_url)
|
||||
answer = image_path
|
||||
# 处理文本类型返回
|
||||
elif "text" in outputs and isinstance(outputs["text"], str):
|
||||
@@ -446,40 +446,43 @@ class DifyPlugin(MessagePluginInterface):
|
||||
else:
|
||||
return 0
|
||||
|
||||
def _get_first_frame(self, video_path, output_path):
|
||||
async def _get_first_frame(self, video_path: str, output_path: str) -> Optional[str]:
|
||||
"""
|
||||
提取视频的第一帧并保存为图片
|
||||
异步提取视频的第一帧并保存为图片
|
||||
: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
|
||||
|
||||
# 使用线程池执行OpenCV操作
|
||||
def extract_frame():
|
||||
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
|
||||
|
||||
# 读取首帧
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
self.LOG.error("无法读取视频帧")
|
||||
cap.release()
|
||||
return None
|
||||
return os.path.abspath(output_path)
|
||||
|
||||
# 保存首帧为图片
|
||||
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)
|
||||
# 在线程池中执行OpenCV操作
|
||||
result = await asyncio.to_thread(extract_frame)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
self.LOG.error(f"提取视频首帧时出错: {e}")
|
||||
|
||||
Reference in New Issue
Block a user