基于Python Katna包自定义关键帧命名关联视频帧时间

问题背景 #

使用Katna提取视频关键帧时，希望自定义文件名，让每个关键帧名称包含对应帧在视频中的精确时间戳。尝试修改Katna源码未成功，单独用OpenCV提取时间戳效率太低，需要高效的集成方案。

可行实现方案 #

方案一：扩展Katna组件（无需修改核心源码，推荐） #

通过扩展FrameExtractor和自定义KeyFrameDiskWriter，让提取过程同时记录帧时间戳，并在保存时生成带时间的文件名。

1. 修改FrameExtractor类，添加时间戳记录

修改Katna的frame_extractor.py中的对应方法，让提取帧时同时保存时间戳：

# 修改__extract_all_frames_from_video__方法
def __extract_all_frames_from_video__(self, videopath):
    cap = cv2.VideoCapture(str(videopath))
    ret, frame = cap.read()
    i = 1
    chunk_no = 0
    
    while ret:
        curr_frame = None
        prev_frame = None
        frame_diffs = []
        frames = []
        timestamps = []  # 新增：存储帧时间戳
        for _ in range(0, self.max_frames_in_chunk):
            if ret:
                # 获取当前帧的时间戳（单位：毫秒）
                timestamp = cap.get(cv2.CAP_PROP_POS_MSEC)
                prev_frame, curr_frame = self.__process_frame(frame, prev_frame, frame_diffs, frames)
                timestamps.append(timestamp)  # 保存时间戳
                i = i + 1
                ret, frame = cap.read()
            else:
                cap.release()
                break
        chunk_no = chunk_no + 1
        yield frames, frame_diffs, timestamps  # 修改返回值，加入时间戳
    cap.release()

# 修改__get_frames_in_local_maxima__方法，返回帧和对应时间戳
def __get_frames_in_local_maxima__(self, frames, frame_diffs, timestamps):
    extracted_key_frames = []
    extracted_timestamps = []
    diff_array = np.array(frame_diffs)
    sm_diff_array = self.__smooth__(diff_array, self.len_window)
    frame_indexes = np.asarray(argrelextrema(sm_diff_array, np.greater))[0]

    for frame_index in frame_indexes:
        extracted_key_frames.append(frames[frame_index - 1])
        extracted_timestamps.append(timestamps[frame_index - 1])  # 对应时间戳
    del frames[:]
    del sm_diff_array
    del diff_array
    del frame_diffs[:]
    del timestamps[:]
    return extracted_key_frames, extracted_timestamps  # 修改返回值

# 修改extract_candidate_frames方法，返回带时间戳的列表
def extract_candidate_frames(self, videopath):
    extracted_candidate_key_frames = []
    extracted_candidate_timestamps = []
    frame_extractor_from_video_generator = self.__extract_all_frames_from_video__(videopath)

    for frames, frame_diffs, timestamps in frame_extractor_from_video_generator:
        if self.USE_LOCAL_MAXIMA:
            frames_chunk, timestamps_chunk = self.__get_frames_in_local_maxima__(frames, frame_diffs, timestamps)
            extracted_candidate_key_frames.extend(frames_chunk)
            extracted_candidate_timestamps.extend(timestamps_chunk)
    return extracted_candidate_key_frames, extracted_candidate_timestamps  # 修改返回值

2. 自定义Writer类，支持时间戳命名

创建自定义的Writer，替代默认的KeyFrameDiskWriter：

from Katna.writer import KeyFrameWriter
import cv2
import os
import datetime

class TimestampKeyFrameWriter(KeyFrameWriter):
    def __init__(self, location):
        super().__init__()
        self.location = location
        os.makedirs(self.location, exist_ok=True)

    def write(self, frames, timestamps):
        for idx, (frame, ts) in enumerate(zip(frames, timestamps)):
            # 将毫秒时间戳转为HH_MM_SS_sss格式（避免文件名含冒号）
            dt = datetime.datetime.fromtimestamp(ts/1000)
            time_str = dt.strftime("%H_%M_%S_%f")[:-3]  # 保留三位毫秒
            filename = f"keyframe_{time_str}.jpg"
            file_path = os.path.join(self.location, filename)
            cv2.imwrite(file_path, frame)

3. 修改调用代码

from Katna.video import Video
import os

if __name__ == "__main__":
    vd = Video()
    no_of_frames_to_returned = 96
    # 使用自定义Writer
    diskwriter = TimestampKeyFrameWriter(location="selectedframes")
    video_file_path = os.path.join("/home/plumpi/Desktop/test_katna", "exp.mp4")

    print(f"Input video file path = {video_file_path}")
    # 需要同步修改Video类的extract_video_keyframes方法：
    # 找到调用frame_extractor.extract_candidate_frames的地方，改为接收帧和时间戳
    # 筛选出最终关键帧后，调用writer.write(selected_frames, selected_timestamps)
    vd.extract_video_keyframes(
         no_of_frames=no_of_frames_to_returned, file_path=video_file_path,
         writer=diskwriter
    )

方案二：直接修改Katna核心Writer（快速实现） #

如果愿意修改Katna的KeyFrameDiskWriter，可以在获取关键帧时同步传递时间戳，修改writer.py中的write方法：

# 修改Katna/writer.py的KeyFrameDiskWriter
def write(self, frames, timestamps):  # 添加timestamps参数
    for idx, (frame, ts) in enumerate(zip(frames, timestamps)):
        dt = datetime.datetime.fromtimestamp(ts/1000)
        time_str = dt.strftime("%H_%M_%S_%f")[:-3]
        filename = f"keyframe_{time_str}.jpg"
        file_path = os.path.join(self.location, filename)
        cv2.imwrite(file_path, frame)

同时需要修改Video类中调用Writer的逻辑，传递时间戳参数。

注意事项 #

• 时间戳格式可根据需求调整，比如用秒数（保留两位小数）命名，避免特殊字符
• 修改Katna源码后，确保本地环境引用的是修改后的版本，而非pip安装的原始版本
• 若不想修改源码，也可以在提取关键帧后，通过帧索引和视频帧率计算时间戳：时间戳 = 帧索引 / 帧率，但效率略低于集成方案

内容的提问来源于stack exchange，提问作者ali noori