make image frames from video

make_images.py

# cuts video to frames with defined stepsize

import numpy as np
import cv2
import os, sys, glob
from pathlib import Path

from argparse import ArgumentParser

DEFAULT_VIDEO_FNAME  = Path('_input/saved_utv_8_1.mkv')
FRAMES_STEP = 1 # write each FRAMES_STEP frame
DEFAULT_FPS = 50
DEFAULT_TARGET_FOLDER = '_output'

def make_empty_folder(out_path):
    # make path if not exists
    out_path.mkdir(parents=True, exist_ok=True) 
    # empty it if anything there
    flist = glob.glob(os.path.join(out_path, '*.png'))
    for f in flist:
        os.remove(f)


def get_video_info(fname):
    # gets video file properties (may differ from actual data, because it's read from header)
    capture = cv2.VideoCapture(fname)
    h = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    w = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    capture.set(cv2.CAP_PROP_POS_AVI_RATIO,1)
    d = float(capture.get(cv2.CAP_PROP_POS_MSEC) / 1000.0)
    fps = int(capture.get(cv2.CAP_PROP_FPS))
    frame_count = int(capture.get(cv2.CAP_PROP_POS_FRAMES))
    capture.release()
    print('fname = {} | h={} w={} duration={} fps={}'.format(fname,h,w,d,fps))
    return h,w,d,fps

def draw_text_box(img, text, x, y):
    font                   = cv2.FONT_HERSHEY_PLAIN
    fontScale              = 1
    fontColor              = (255,255,255)
    backColor              = (0,0,0)
    lineType               = 1
    # get the width and height of the text box
    (t_w, t_h) = cv2.getTextSize(text, font, fontScale=fontScale, thickness=lineType)[0]
    # make the coords of the box with a small padding of two pixels
    box_coords = ((x, y), (x + t_w - 2, y - t_h - 2))
    cv2.rectangle(img, box_coords[0], box_coords[1], backColor, cv2.FILLED)
    cv2.putText(img, text, (x, y), font, fontScale=fontScale, color=fontColor, thickness=lineType)
    return


##############################################################

# return filename without extension  (c:\temp\abcd.efg -> abcd)
def get_fname_wo_extension(filename):
    return os.path.splitext(os.path.basename(filename))[0]

# MAIN
def main():
##############################################################
    parser = ArgumentParser()
    parser.add_argument('-f', '--filename', dest='filename', help='select file')
    parser.add_argument('-t', '--targetfolder', dest='targetfolder', help='override target folder (default is {})'.format(DEFAULT_TARGET_FOLDER))
    parser.add_argument('-fps', '--framespersecond', dest='fps', help='override FPS')
    args = parser.parse_args()

    if (args.filename) is not None:
       VIDEO_FNAME = Path(args.filename.strip())
       print('Overriding fname to {}'.format(VIDEO_FNAME))
    else:
       VIDEO_FNAME = Path(DEFAULT_VIDEO_FNAME)
       print('Using hardcoded filename = {}'.format(VIDEO_FNAME))

    if (args.targetfolder) is not None:
       TARGET_FOLDER = args.targetfolder.strip()
       print('Overriding targetfolder to {}'.format(TARGET_FOLDER))
    else:
       TARGET_FOLDER = DEFAULT_TARGET_FOLDER
       print('Using hardcoded targetfolder = {}'.format(TARGET_FOLDER))


##############################################################

    fname=str(VIDEO_FNAME.absolute()) # resolve full path
    h,w,d,fps = get_video_info(fname)

    if (args.fps) is not None:
       FPS = int(args.fps.strip())
       print('Overriding fps to {}'.format(FPS))
    else:
       FPS = fps 
       print('Using fps from file {}'.format(fps))

    cap = cv2.VideoCapture(fname)
    PREFIX = get_fname_wo_extension(fname)
    cntr = 0
    OUTPUT_FOLDER = Path(str(Path.cwd())+'/'+TARGET_FOLDER)
    make_empty_folder(OUTPUT_FOLDER)
    print('Writing frames to {}'.format(OUTPUT_FOLDER))

    while(cap.isOpened()):
       try:
          ret, frame = cap.read()
          if frame is None:
             break
       except:
          break
 
#       gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
       gray = frame
       small = cv2.resize(gray, (0,0), fx=0.5, fy=0.5,  interpolation = cv2.INTER_AREA) 
       # draws text box in the lower left corner
       draw_text_box(small, '{}x{} {} s | {} of {} @ {} FPS | Press ESC to abort'.format(w,h,d, cntr, int(d * FPS), FPS ), 10, h//2 - 20)
       cv2.imshow('frame',small)
       if (cntr % FRAMES_STEP) ==0:
          cv2.imwrite(str(Path(str(OUTPUT_FOLDER.absolute())+'/'+PREFIX+'.{:08d}.png'.format(cntr))), gray)
       cntr = cntr + 1
       if cv2.waitKey(1) & 0xFF == 27: #ESC
          break

    cap.release()
    cv2.destroyAllWindows()
    print('SUCCESS')

if __name__ == "__main__":
    main()