FrameProcessor/graph/steps/frame_selection.py

# FrameProcessor/graph/steps/frame_selection.py

import os
import shutil
import json
from KeyFrameSelection.FeatureExtraction import process_video, save_records
from KeyFrameSelection.Similarties import hash_filter, clip_filter
from FrameProcessor.utils.io_utils import get_frames_from_folder, save_description_to_csv
from FrameProcessor.processor.multi_frame import process_frames
from config.paths import output_csv_file, output_json_file

def extract_and_process_frames_node(state):
    video_path = state["frame_path"]

    if os.path.exists("outputs"):
        shutil.rmtree("outputs")
    os.makedirs("outputs/final_output", exist_ok=True)

    keyframe_dir = "outputs/keyframes"
    csv_path = "outputs/keyframes.csv"

    # Step 1: Extract frames
    records, fps = process_video(video_path, interval_sec=10)

    # Step 2: Filter
    min_frames = 10
    max_iterations = 20
    hash_threshold = 5
    ssim_threshold = 0.95
    clip_threshold = 0.90
    iteration = 0
    filtered = records

    while len(filtered) >= min_frames and iteration < max_iterations:
        filtered = hash_filter(filtered, hash_threshold, ssim_threshold, 5)
        filtered = clip_filter(filtered, clip_threshold, 5)

        hash_threshold = max(1, hash_threshold - 1)
        ssim_threshold = max(0.5, ssim_threshold - 0.05)
        clip_threshold = min(0.99, clip_threshold + 0.03)
        iteration += 1

    # Step 3: Save
    save_records(filtered, keyframe_dir, csv_path, fps)
    frame_paths = get_frames_from_folder(keyframe_dir)

    # Step 4: Process
    results = process_frames(frame_paths)
    important_frames = [r for r in results if r["importance"] == "important"]

    for result in important_frames:
        save_description_to_csv(result, output_csv_file)

    with open(output_json_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    # Update state
    state["frame_paths"] = frame_paths
    state["results"] = results
    state["important_frames"] = important_frames
    return state