avi2png.py extract images from video

avi2png.py

"""
Created on Mon Nov 13 16:12:42 2017
@author: merwan
"""

import numpy as np
import cv2
import time
import os
import argparse


def Path2data():
    parser = argparse.ArgumentParser(description="Path 2 data")
    parser.add_argument("input_dir" , help="input directory.")
    parser.add_argument("avi_file",   help="The .avi file to load.")
    parser.add_argument("str_cam",    help="The camera number.")
    args = parser.parse_args()
    print "Read the video file %s, located in %s" % (args.avi_file, args.input_dir)
    return args.avi_file, args.input_dir, args.str_cam



# Function definition is here
def extract_png(avi_file, input_dir, str_cam):
    start_time = time.time()
    cap = cv2.VideoCapture(input_dir+'/'+avi_file)
    cmpt = 0;
    
    print(input_dir+'/'+avi_file)

    while(cap.isOpened()):
        ret, frame = cap.read()
        
        # print(cmpt)
        if ret == True:
            cmpt = cmpt + 1
            gray = frame # cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            cv2.imshow('frame',gray)
            # crreat the folder if this one doesn't exist
            newpath = input_dir+"/datacam/pic"+str_cam        
            if not os.path.exists(newpath):
                os.makedirs(newpath)
            # write the image in the corresponding folder
            cv2.imwrite(input_dir+"datacam/pic_"+ str_cam + "/image_" + str(cmpt) + ".png", gray)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        else:
            break
    
    cap.release()
    cv2.destroyAllWindows()
    # 
    print("data set "+ str_cam)
    print("--- %s images  ---" % (cmpt))
    print("--- %s seconds ---" % (time.time() - start_time))
    return;


if __name__ == '__main__':
    
    # Load file names, ...
    avi_file, input_dir, str_cam = Path2data()
    # Extract images from the video 
    extract_png(avi_file, input_dir, str_cam)