import cv2
import numpy as np
import os
from PIL import Image

class KeyFrameExtractor:
    def __init__(self, threshold=0.4, min_interval=10):
        """
        初始化關(guān)鍵幀提取器
        :param threshold: 幀差異閾值，超過此值則視為關(guān)鍵幀
        :param min_interval: 關(guān)鍵幀之間的最小間隔（幀數(shù)）
        """
        self.threshold = threshold
        self.min_interval = min_interval
        self.last_keyframe = None
        self.last_keyframe_idx = -min_interval
        
    def _calculate_frame_difference(self, frame1, frame2):
        """計(jì)算兩幀之間的差異度"""
        # 轉(zhuǎn)換為灰度圖
        gray1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
        gray2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
        
        # 計(jì)算直方圖
        hist1 = cv2.calcHist([gray1], [0], None, [256], [0, 256])
        hist2 = cv2.calcHist([gray2], [0], None, [256], [0, 256])
        
        # 歸一化直方圖
        hist1 = cv2.normalize(hist1, hist1).flatten()
        hist2 = cv2.normalize(hist2, hist2).flatten()
        
        # 計(jì)算直方圖相似度（相關(guān)性）
        similarity = cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)
        
        # 返回差異度（1 - 相似度）
        return 1 - similarity
    
    def extract_keyframes(self, video_path, output_dir="keyframes"):
        """
        從視頻中提取關(guān)鍵幀
        :param video_path: 視頻文件路徑
        :param output_dir: 關(guān)鍵幀保存目錄
        :return: 提取的關(guān)鍵幀列表及其幀索引
        """
        # 創(chuàng)建輸出目錄
        os.makedirs(output_dir, exist_ok=True)
        
        # 打開視頻文件
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"無法打開視頻文件: {video_path}")
        
        keyframes = []
        frame_idx = 0
        ret, prev_frame = cap.read()
        
        if not ret:
            raise ValueError("無法讀取視頻幀")
        
        # 將第一幀作為初始關(guān)鍵幀
        self.last_keyframe = prev_frame
        self.last_keyframe_idx = 0
        keyframes.append((0, prev_frame))
        self._save_keyframe(prev_frame, 0, output_dir)
        
        while True:
            ret, curr_frame = cap.read()
            frame_idx += 1
            
            if not ret:
                break  # 視頻結(jié)束
            
            # 計(jì)算當(dāng)前幀與上一關(guān)鍵幀的差異
            diff = self._calculate_frame_difference(self.last_keyframe, curr_frame)
            
            # 檢查是否滿足關(guān)鍵幀條件
            if diff > self.threshold and (frame_idx - self.last_keyframe_idx) > self.min_interval:
                keyframes.append((frame_idx, curr_frame))
                self.last_keyframe = curr_frame
                self.last_keyframe_idx = frame_idx
                self._save_keyframe(curr_frame, frame_idx, output_dir)
        
        cap.release()
        print(f"提取完成，共提取 {len(keyframes)} 個關(guān)鍵幀，保存至 {output_dir} 目錄")
        return keyframes
    
    def _save_keyframe(self, frame, frame_idx, output_dir):
        """保存關(guān)鍵幀為圖片文件"""
        # 轉(zhuǎn)換BGR為RGB（OpenCV默認(rèn)BGR格式）
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(frame_rgb)
        output_path = os.path.join(output_dir, f"keyframe_{frame_idx:06d}.jpg")
        img.save(output_path)

if __name__ == "__main__":
    # 使用示例
    import argparse
    
    parser = argparse.ArgumentParser(description="視頻關(guān)鍵幀提取工具")
    parser.add_argument("video_path", help="視頻文件路徑")
    parser.add_argument("--threshold", type=float, default=0.4, 
                      help="幀差異閾值，值越大提取的關(guān)鍵幀越少")
    parser.add_argument("--min_interval", type=int, default=10, 
                      help="關(guān)鍵幀之間的最小間隔（幀數(shù)）")
    parser.add_argument("--output_dir", default="keyframes", 
                      help="關(guān)鍵幀保存目錄")
    
    args = parser.parse_args()
    
    # 初始化提取器并提取關(guān)鍵幀
    extractor = KeyFrameExtractor(threshold=args.threshold, min_interval=args.min_interval)
    extractor.extract_keyframes(args.video_path, args.output_dir)