VideoAgent/_videoutil/split.py

import os
import time
import shutil
import subprocess
import numpy as np
from tqdm import tqdm
from moviepy.video import fx as vfx
from moviepy.video.io.VideoFileClip import VideoFileClip
import logging
import multiprocessing

def _get_ffmpeg_bin():
    """返回支持 libx264 的 ffmpeg 路径"""
    env_bin = os.getenv("FFMPEG_BIN")
    if env_bin:
        return env_bin
    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except ImportError:
        return "ffmpeg"

_FFMPEG_BIN = _get_ffmpeg_bin()

logger = logging.getLogger(__name__)


def preprocess_video(
    video_path,
    target_width=640,
    target_height=360,
    target_fps=24,
    video_output_format='mp4',
):
    """将输入视频调整为固定分辨率和帧率，存储到视频文件上级目录的同级 processed 目录下。

    如果目标文件已存在则直接复用，返回预处理后的视频路径。
    """
    video_name = os.path.basename(video_path).split('.')[0]
    processed_dir = os.path.join(os.path.dirname(os.path.dirname(video_path)), 'processed')
    os.makedirs(processed_dir, exist_ok=True)
    output_path = os.path.join(processed_dir, f"{video_name}.{video_output_format}")


    if os.path.exists(output_path):
        logger.info(f"Preprocessed video already exists: {output_path}")
        return output_path

    logger.info(f"Preprocessing video {video_name}: {target_width}x{target_height} @ {target_fps}fps -> {output_path}")
    cmd = [
        _FFMPEG_BIN, "-y",
        "-i", video_path,
        "-vf", f"scale={target_width}:{target_height}",
        "-r", str(target_fps),
        "-c:v", "libx264",
        "-c:a", "aac",
        "-b:a", "128k",
        "-loglevel", "error",
        output_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg preprocessing failed for {video_path}:\n{result.stderr}")
    logger.info(f"Preprocessed video saved to {output_path}")
    return output_path

def split_video(
    video_path,
    working_dir,
    segment_length,
    num_frames_per_segment,
    audio_output_format='mp3',
    ):  
    unique_timestamp = str(int(time.time() * 1000))
    video_name = os.path.basename(video_path).split('.')[0]
    video_segment_cache_path = os.path.join(working_dir, '_cache', video_name)
    if os.path.exists(video_segment_cache_path):
        shutil.rmtree(video_segment_cache_path)
    os.makedirs(video_segment_cache_path, exist_ok=False)
    
    segment_index = 0
    segment_index2name, segment_times_info = {}, {}
    with VideoFileClip(video_path) as video:
    
        total_video_length = int(video.duration)
        start_times = list(range(0, total_video_length, segment_length))
        # if the last segment is shorter than 5 seconds, we merged it to the last segment
        if len(start_times) > 1 and (total_video_length - start_times[-1]) < 5:
            start_times = start_times[:-1]
        
        for start in tqdm(start_times, desc=f"Spliting Video {video_name}"):
            if start != start_times[-1]:
                end = min(start + segment_length, total_video_length)
            else:
                end = total_video_length
            
            subvideo = video.subclip(start, end)
            subvideo_length = subvideo.duration
            frame_times = np.linspace(0, subvideo_length, num_frames_per_segment, endpoint=False)
            frame_times += start
            
            segment_index2name[f"{segment_index}"] = f"{unique_timestamp}-{segment_index}-{start}-{end}"
            segment_times_info[f"{segment_index}"] = {"frame_times": frame_times, "timestamp": (start, end)}
            
            # save audio
            audio_file_base_name = segment_index2name[f"{segment_index}"]
            audio_file = f'{audio_file_base_name}.{audio_output_format}'
            try:
                subaudio = subvideo.audio
                subaudio.write_audiofile(os.path.join(video_segment_cache_path, audio_file), codec='pcm_s16le', fps=16000,  nbytes=2,  verbose=False, logger=None)
            except Exception as e:
                logger.warning(f"Warning: Failed to extract audio for video {video_name} ({start}-{end}). Probably due to lack of audio track.")

            segment_index += 1

    return segment_index2name, segment_times_info

def saving_video_segments(
    video_name,
    video_path,
    working_dir,
    segment_index2name,
    segment_times_info,
    error_queue,
    video_output_format='mp4',
    ):
    try:
        with VideoFileClip(video_path) as video:
            video_segment_cache_path = os.path.join(working_dir, '_cache', video_name)
            for index in tqdm(segment_index2name, desc=f"Saving Video Segments {video_name}"):
                start, end = segment_times_info[index]["timestamp"][0], segment_times_info[index]["timestamp"][1]
                video_file = f'{segment_index2name[index]}.{video_output_format}'
                # print("video_file ", video_file )
                subvideo = video.subclip(start, end)
                subvideo.write_videofile(os.path.join(video_segment_cache_path, video_file), codec='libx264', verbose=False, logger=None)
    except Exception as e:
        error_queue.put(f"Error in saving_video_segments:\n {str(e)}")
        raise RuntimeError