OpenCV python utils

opencv_utils_video.py

import cv2
from os import path
from math import floor
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
#from tqdm import tqdm_notebook as tqdm # for jupyter notebooks

# Video to Frame
def get_frame(video, second, save_folder=None, prefix="frame", resize=None, to_rgb=True, prefix_padding=4):
    """
    get_frame save a frame
    """
    video.set(cv2.CAP_PROP_POS_MSEC, second*1000)
    fps, _, _, _ = get_video_data(video)
    frame_count = floor(second*fps)
    ret, frame = video.read()
    if ret:
        if resize is not None and ():
            h_old, w_old,_ = frame.shape
            h_new, w_new = resize
            interpolation = cv2.INTER_AREA if (h_old > h_new) and (w_old > w_new) else cv2.INTER_CUBIC
            frame = cv2.resize(frame, resize, interpolation=interpolation)
        if save_folder is not None:
            save_as = f"{prefix}_{str(frame_count).zfill(prefix_padding)}.jpg"
            cv2.imwrite(path.join(save_folder, save_as), frame)
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) if to_rgb else frame
    return ret, frame

def get_frames(video, start, end=-1, save_folder=None, prefix="frame", resize=None, to_rgb=True, prefix_padding=4, return_frames=False):
    """
    get_frame save frames from a range of seconds
    """
    second = float(start)
    fps, sec_per_frame, _, duration = get_video_data(video)
    frame_count = floor(second*fps)
    end = duration if end == -1 else end
    frames = []
    total_frames = int((end - start)*fps)
    pbar = tqdm(total=total_frames)
    while (second < end):
        video.set(cv2.CAP_PROP_POS_MSEC, second*1000)
        ret, frame = video.read()
        if ret:
            if resize is not None and ():
                h_old, w_old,_ = frame.shape
                h_new, w_new = resize
                interpolation = cv2.INTER_AREA if (h_old > h_new) and (w_old > w_new) else cv2.INTER_CUBIC
                frame = cv2.resize(frame, resize, interpolation=interpolation)
            if save_folder is not None:
                save_as = path.join(save_folder, f"{prefix}_{str(frame_count).zfill(prefix_padding)}.jpg")
                cv2.imwrite(save_as, frame)
            if return_frames:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) if to_rgb else frame
                frames.append(frame)
        else:
            break
        second += sec_per_frame
        frame_count += 1
        pbar.update(1)
    if return_frames:
        return frame_count, np.array(frames)
    else:
        return frame_count

def get_video_data(video):
    """
    get_video_data extract data from video.
    """
    fps = video.get(cv2.CAP_PROP_FPS)
    sec_per_frame = 1/fps
    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count * sec_per_frame
    return fps, sec_per_frame, frame_count, duration


def plot_frame(frame, title="frame", figsize=(8,12)):
    plt.figure(figsize=figsize)
    plt.imshow(frame)
    plt.yticks([])
    plt.xticks([])
    plt.grid(False)
    plt.show()