extract frames from videos (OpenCV)

vid2img.py

"""
turn videos in source directory into images in target directory
"""
import sys
import os
import cv2

def save_images(cap, path, num_frames=1):
    """save frames from clip into image (equidistant frames)"""
    # path should be like target/videoname
    def write(k, frame):
        cv2.imwrite(f"{path}{k}.png", frame)

    if num_frames == 1:
        ret, frame = cap.read()
        write(0, frame)
    else:
        frames = []
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret: # end of video
                break
            frames.append(frame)
        step = int(len(frames)/num_frames)
        f = 0
        chosen = 0
        while chosen < num_frames:
            write(chosen, frames[f])
            f += step
            chosen += 1
    cap.release()


# get source and target dirs and append '/' to them if it's not already there
source = sys.argv[1]
if source[-1] != '/':
    source += '/'
target = sys.argv[2]
if target[-1] != '/':
    target += '/'

print(f"extracting frames from {source} to {target}")
files = os.listdir(source)
print(f"found {len(files)} files")

file_count = 0
for fname in files:
    file_count += 1
    # print(f"extracting from {fname} - {file_count}/{len(files)}")
    print("*", end='')
    cap = cv2.VideoCapture(source + fname) # VideoCapture(source/file.mp4)
    path = target + fname.split('.')[0] # target/file
    save_images(cap, path, 5)
print()
print("done")