From 9757ed2c7eeafa3fa92271809ac7bed87a54822d Mon Sep 17 00:00:00 2001 From: Zhichao Date: Mon, 5 Aug 2024 11:31:09 +0800 Subject: [PATCH] recover to use moviepy to save video --- vision_agent/tools/tools.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/vision_agent/tools/tools.py b/vision_agent/tools/tools.py index 322a7f25..38b8e31c 100644 --- a/vision_agent/tools/tools.py +++ b/vision_agent/tools/tools.py @@ -9,6 +9,7 @@ import cv2 import numpy as np import requests +from moviepy.editor import ImageSequenceClip from PIL import Image, ImageDraw, ImageFont from pillow_heif import register_heif_opener # type: ignore from pytube import YouTube # type: ignore @@ -1063,20 +1064,15 @@ def save_video( _LOGGER.warning(f"Invalid fps value: {fps}. Setting fps to 4 (default value).") fps = 4 - if not output_video_path: - output_video_path = tempfile.NamedTemporaryFile( - suffix=".mp4", delete=False - ).name - - height, width, layers = frames[0].shape if frames else (0, 0, 0) - fourcc = cv2.VideoWriter_fourcc(*"X264") # type: ignore - video = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height)) - for frame in frames: - video.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) - video.release() - - _save_video_to_result(output_video_path) - return output_video_path + with ImageSequenceClip(frames, fps=fps) as video: + if output_video_path: + f = open(output_video_path, "wb") + else: + f = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) # type: ignore + video.write_videofile(f.name, codec="libx264") + f.close() + _save_video_to_result(f.name) + return f.name def _save_video_to_result(video_uri: str) -> None: