File size: 4,143 Bytes
83e35a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
"""
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") |