multithreaded_extractor.py

import cv2
import multiprocessing
import os
import time
from datetime import date
from concurrent.futures import ProcessPoolExecutor

TODAY = date.today().strftime("%d_%m")


def extractImages(vid, output):
    count = 0
    vidcap = cv2.VideoCapture(vid)
    success, image = vidcap.read()
    success = True
    while success:
        vidcap.set(cv2.CAP_PROP_POS_MSEC, (count * 1000))
        success, image = vidcap.read()
        if success:
            image = cv2.resize(image, (600, 400))
            cv2.imwrite(output + "/{}_{:d}.jpg".format(TODAY, count), image)
            count += 1
    return count


def video_to_frames(video_path, frames_dir):
    chunk_size = 1000
    video_path = os.path.normpath(video_path)
    frames_dir = os.path.normpath(frames_dir)
    video_dir, video_filename = os.path.split(video_path)
    video_filename = video_filename[:-4]
    os.makedirs(os.path.join(frames_dir), exist_ok=True)

    assert os.path.exists(video_path)

    capture = cv2.VideoCapture(video_path)
    total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    capture.release()

    if total < 1:
        print("Video has no frames. Check your OpenCV + ffmpeg installation")
        return None

    frame_chunks = [[i, i + chunk_size] for i in range(0, total, chunk_size)]
    frame_chunks[-1][-1] = min(frame_chunks[-1][-1], total - 1)

    print("Ripped workload with {} cpu".format(multiprocessing.cpu_count()))
    with ProcessPoolExecutor(max_workers=multiprocessing.cpu_count()) \
            as executor:
        for f in frame_chunks:
            executor.submit(
                extractImages,
                vid=video_path,
                output=os.path.join(frames_dir),
            )
    executor.shutdown(wait=False)
    return os.path.join(frames_dir)


if __name__ == "__main__":
    start = time.time()
    video_to_frames(video_path="ginchan.avi", frames_dir="output")
    print("Finished in: {} seconds".format(time.time() - start))