import cv2
import os


def extract_frames_from_single_vid(input_vid, output_dir):

    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)

    vidcap = cv2.VideoCapture(input_vid)

    success = True

    count = 0

    while success:
        success, image = vidcap.read()
        #print('Read a new frame: ', success)
        cv2.imwrite(os.path.join(output_dir, '{:0>5}.jpg'.format(count)), image)  # save frame as JPEG file
        count += 10

    print('Converted {} frames successfully'.format(count))


def extract_frames_from_batch(base_in_dir, extension):

    for file in os.listdir(base_in_dir):
        if file.endswith(".mp4"):
            input_vid_file_path = os.path.join(base_in_dir, file)
            print("extracting frames for : " + file)

            filename, file_extension = os.path.splitext(file)
            output_dir = os.path.join(base_in_dir, "output/" + filename)

            extract_frames_from_single_vid(input_vid_file_path, output_dir)




if __name__ == '__main__':

    #####################################################
    # Use this to extract frames from single video file #
    #####################################################

    """
    base_dir = r'/home/omer/Work/Ford-Scenarios/BM1/Images+BBs/Train-Data/Car_Turning_Right_Train/'
    input_vid = os.path.join(base_dir, 'Boston_Massachusetts_USA_15_37-15_41.mp4')
    output_dir = os.path.join(base_dir, 'Boston_Massachusetts_USA_15_37-15_41')

    extract_frames_from_single_vid(input_vid, output_dir)
    """


    ##########################################################
    # Use this to extract frames from a batch of video files #
    ##########################################################

    base_dir = r'/home/user/Desktop/video_test'
    extension = ".MP4"

    extract_frames_from_batch(base_dir, extension)