51 lines
1.3 KiB
Python
Executable File
51 lines
1.3 KiB
Python
Executable File
from typing import Callable
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
from tqdm import tqdm, trange
|
|
|
|
|
|
def compare(frame1, frame2):
|
|
absdiff = cv2.absdiff(frame1, frame2)
|
|
return np.mean(absdiff)
|
|
|
|
|
|
def load_frames(cap: cv2.VideoCapture, frame_interval: int, progress: Callable) -> list:
|
|
frames = []
|
|
num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
i = 0
|
|
res = cap.grab()
|
|
while i < num_frames:
|
|
if res:
|
|
frames.append(cap.retrieve()[1])
|
|
for _ in range(frame_interval):
|
|
res = cap.grab()
|
|
i += frame_interval
|
|
progress(i)
|
|
|
|
return frames
|
|
|
|
|
|
def differences(frames: list, progress: Callable) -> list[float]:
|
|
diffs = []
|
|
|
|
for index, (frame1, frame2) in enumerate(zip(frames[:-1], frames[1:]), start=1):
|
|
diff = compare(frame1, frame2)
|
|
diffs.append(diff)
|
|
progress(index)
|
|
|
|
return diffs
|
|
|
|
|
|
def select_frames(frames: list, diffs: list, threshold: float) -> list[Image.Image]:
|
|
selected_frames = [frames[0]]
|
|
|
|
for frame, diff in zip(frames[1:], diffs, strict=True):
|
|
if diff > threshold:
|
|
selected_frames.append(frame)
|
|
|
|
# Convert to Images
|
|
return [Image.fromarray(frame[:, :, ::-1]) for frame in selected_frames]
|