NarratoAI/app/utils/video_processor.py

import cv2
import numpy as np
from sklearn.cluster import KMeans
import os
import re
from typing import List, Tuple, Generator


class VideoProcessor:
    def __init__(self, video_path: str):
        """
        初始化视频处理器
        
        Args:
            video_path: 视频文件路径
        """
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"视频文件不存在: {video_path}")
        
        self.video_path = video_path
        self.cap = cv2.VideoCapture(video_path)
        
        if not self.cap.isOpened():
            raise RuntimeError(f"无法打开视频文件: {video_path}")
        
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))

    def __del__(self):
        """析构函数，确保视频资源被释放"""
        if hasattr(self, 'cap'):
            self.cap.release()

    def preprocess_video(self) -> Generator[np.ndarray, None, None]:
        """
        使用生成器方式读取视频帧
        
        Yields:
            np.ndarray: 视频帧
        """
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)  # 重置到视频开始
        while self.cap.isOpened():
            ret, frame = self.cap.read()
            if not ret:
                break
            yield frame

    def detect_shot_boundaries(self, frames: List[np.ndarray], threshold: int = 30) -> List[int]:
        """
        使用帧差法检测镜头边界
        
        Args:
            frames: 视频帧列表
            threshold: 差异阈值
            
        Returns:
            List[int]: 镜头边界帧的索引列表
        """
        shot_boundaries = []
        for i in range(1, len(frames)):
            prev_frame = cv2.cvtColor(frames[i - 1], cv2.COLOR_BGR2GRAY)
            curr_frame = cv2.cvtColor(frames[i], cv2.COLOR_BGR2GRAY)
            diff = np.mean(np.abs(curr_frame.astype(int) - prev_frame.astype(int)))
            if diff > threshold:
                shot_boundaries.append(i)
        return shot_boundaries

    def filter_keyframes_by_time(self, keyframes: List[np.ndarray], 
                               keyframe_indices: List[int]) -> Tuple[List[np.ndarray], List[int]]:
        """
        过滤关键帧，确保每秒最多只有一个关键帧
        
        Args:
            keyframes: 关键帧列表
            keyframe_indices: 关键帧索引列表
            
        Returns:
            Tuple[List[np.ndarray], List[int]]: 过滤后的关键帧列表和对应的帧索引
        """
        if not keyframes or not keyframe_indices:
            return keyframes, keyframe_indices
            
        filtered_frames = []
        filtered_indices = []
        last_second = -1
        
        for frame, idx in zip(keyframes, keyframe_indices):
            current_second = idx // self.fps
            if current_second != last_second:
                filtered_frames.append(frame)
                filtered_indices.append(idx)
                last_second = current_second
                
        return filtered_frames, filtered_indices

    def extract_keyframes(self, frames: List[np.ndarray], shot_boundaries: List[int]) -> Tuple[List[np.ndarray], List[int]]:
        """
        从每个镜头中提取关键帧，并确保每秒最多一个关键帧
        
        Args:
            frames: 视频帧列表
            shot_boundaries: 镜头边界列表
            
        Returns:
            Tuple[List[np.ndarray], List[int]]: 关键帧列表和对应的帧索引
        """
        keyframes = []
        keyframe_indices = []
        
        for i in range(len(shot_boundaries)):
            start = shot_boundaries[i - 1] if i > 0 else 0
            end = shot_boundaries[i]
            shot_frames = frames[start:end]

            frame_features = np.array([cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).flatten() 
                                     for frame in shot_frames])
            kmeans = KMeans(n_clusters=1, random_state=0).fit(frame_features)
            center_idx = np.argmin(np.sum((frame_features - kmeans.cluster_centers_[0]) ** 2, axis=1))

            keyframes.append(shot_frames[center_idx])
            keyframe_indices.append(start + center_idx)

        # 过滤每秒多余的关键帧
        filtered_keyframes, filtered_indices = self.filter_keyframes_by_time(keyframes, keyframe_indices)
        
        return filtered_keyframes, filtered_indices

    def save_keyframes(self, keyframes: List[np.ndarray], keyframe_indices: List[int], 
                      output_dir: str) -> None:
        """
        保存关键帧到指定目录，文件名格式为：keyframe_帧序号_时间戳.jpg
        
        Args:
            keyframes: 关键帧列表
            keyframe_indices: 关键帧索引列表
            output_dir: 输出目录
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        for keyframe, frame_idx in zip(keyframes, keyframe_indices):
            # 计算时间戳（秒）
            timestamp = frame_idx / self.fps
            # 将时间戳转换为 HH:MM:SS 格式
            hours = int(timestamp // 3600)
            minutes = int((timestamp % 3600) // 60)
            seconds = int(timestamp % 60)
            time_str = f"{hours:02d}{minutes:02d}{seconds:02d}"
            
            # 构建新的文件名格式：keyframe_帧序号_时间戳.jpg
            output_path = os.path.join(output_dir, 
                                     f'keyframe_{frame_idx:06d}_{time_str}.jpg')
            cv2.imwrite(output_path, keyframe)

        print(f"已保存 {len(keyframes)} 个关键帧到 {output_dir}")

    def extract_frames_by_numbers(self, frame_numbers: List[int], output_folder: str) -> None:
        """
        根据指定的帧号提取帧
        
        Args:
            frame_numbers: 要提取的帧号列表
            output_folder: 输出文件夹路径
        """
        if not frame_numbers:
            raise ValueError("未提供帧号列表")
        
        if any(fn >= self.total_frames or fn < 0 for fn in frame_numbers):
            raise ValueError("存在无效的帧号")

        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        for frame_number in frame_numbers:
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
            ret, frame = self.cap.read()

            if ret:
                # 计算时间戳
                timestamp = frame_number / self.fps
                hours = int(timestamp // 3600)
                minutes = int((timestamp % 3600) // 60)
                seconds = int(timestamp % 60)
                time_str = f"{hours:02d}{minutes:02d}{seconds:02d}"
                
                # 使用与关键帧相同的命名格式
                output_path = os.path.join(output_folder, 
                                         f"extracted_frame_{frame_number:06d}_{time_str}.jpg")
                cv2.imwrite(output_path, frame)
                print(f"已提取并保存帧 {frame_number}")
            else:
                print(f"无法读取帧 {frame_number}")

    @staticmethod
    def extract_numbers_from_folder(folder_path: str) -> List[int]:
        """
        从文件夹中提取帧号
        
        Args:
            folder_path: 关键帧文件夹路径
            
        Returns:
            List[int]: 排序后的帧号列表
        """
        files = [f for f in os.listdir(folder_path) if f.endswith('.jpg')]
        # 更新正则表达式以匹配新的文件名格式：keyframe_000123_010534.jpg
        pattern = re.compile(r'keyframe_(\d+)_\d+\.jpg$')
        numbers = []
        for f in files:
            match = pattern.search(f)
            if match:
                numbers.append(int(match.group(1)))
        return sorted(numbers)

    def process_video(self, output_dir: str, skip_seconds: float = 0) -> None:
        """
        处理视频并提取关键帧
        
        Args:
            output_dir: 输出目录
            skip_seconds: 跳过视频开头的秒数
        """
        # 计算要跳过的帧数
        skip_frames = int(skip_seconds * self.fps)
        
        # 获取所有帧
        frames = list(self.preprocess_video())
        
        # 跳过指定秒数的帧
        frames = frames[skip_frames:]
        
        if not frames:
            raise ValueError(f"跳过 {skip_seconds} 秒后没有剩余帧可以处理")
        
        shot_boundaries = self.detect_shot_boundaries(frames)
        keyframes, keyframe_indices = self.extract_keyframes(frames, shot_boundaries)
        
        # 调整关键帧索引，加上跳过的帧数
        adjusted_indices = [idx + skip_frames for idx in keyframe_indices]
        self.save_keyframes(keyframes, adjusted_indices, output_dir)