Testcomic / backend /keyframes /keyframes_simple.py
3v324v23's picture
Update Comic123 with local comic folder files
83e35a7
"""
Simplified Keyframe Extraction
Avoids infinite loops by using basic frame selection
"""
import os
import srt
import cv2
import numpy as np
from backend.keyframes.extract_frames import extract_frames
from backend.utils import copy_and_rename_file
def generate_keyframes_simple(video):
"""Generate keyframes using simplified method"""
print("🎯 Using simplified keyframe generation...")
# Read subtitle file
try:
with open("test1.srt") as f:
data = f.read()
subs = list(srt.parse(data))
except:
print("❌ Error reading subtitles")
return False
# Create final directory
final_dir = os.path.join("frames", "final")
if not os.path.exists(final_dir):
os.makedirs(final_dir)
print(f"Created directory: {final_dir}")
frame_counter = 1
total_subs = len(subs)
print(f"🎯 Processing {total_subs} subtitle segments...")
# Process segments with simplified logic
segments_to_process = min(16, total_subs) # Max 16 segments
for i, sub in enumerate(subs[:segments_to_process], 1):
print(f"πŸ“ Processing segment {i}/{segments_to_process}: {sub.content[:50]}...")
# Create segment directory
sub_dir = f"frames/sub{sub.index}"
if not os.path.exists(sub_dir):
os.makedirs(sub_dir)
try:
# Extract 3-5 frames per segment (reduced from 10)
frames = extract_frames(video, sub_dir,
sub.start.total_seconds(),
sub.end.total_seconds(),
3) # Only 3 frames per segment
if frames:
# Simple selection: pick middle frame or best quality frame
best_frame = _select_best_frame_simple(frames)
if best_frame and frame_counter <= 16:
# Copy to final directory
final_name = f"frame{frame_counter:03}.png"
copy_and_rename_file(best_frame, final_dir, final_name)
print(f"πŸ“– Frame {frame_counter}: {sub.content[:30]}...")
frame_counter += 1
except Exception as e:
print(f"⚠️ Error processing segment {i}: {e}")
continue
frames_generated = frame_counter - 1
print(f"βœ… Generated {frames_generated} frames using simplified method")
# If we don't have enough frames, duplicate some to reach 16
if frames_generated < 16:
print(f"πŸ”„ Duplicating frames to reach 16 total...")
for i in range(frames_generated + 1, 17):
# Duplicate existing frames
source_frame = f"frame{((i-1) % frames_generated) + 1:03}.png"
source_path = os.path.join(final_dir, source_frame)
target_path = os.path.join(final_dir, f"frame{i:03}.png")
if os.path.exists(source_path):
import shutil
shutil.copy2(source_path, target_path)
print(f"πŸ“‹ Duplicated frame{i:03}.png")
return True
def _select_best_frame_simple(frames):
"""Select best frame using simple criteria"""
if not frames:
return None
if len(frames) == 1:
return frames[0]
# Simple heuristic: pick frame with most color variance (usually more interesting)
best_frame = None
best_score = 0
for frame_path in frames:
try:
img = cv2.imread(frame_path)
if img is not None:
# Calculate color variance as a simple quality metric
variance = np.var(img)
if variance > best_score:
best_score = variance
best_frame = frame_path
except:
continue
# Fallback to middle frame if variance method fails
return best_frame if best_frame else frames[len(frames)//2]
if __name__ == "__main__":
# Test the simplified method
generate_keyframes_simple("video/IronMan.mp4")