general_purpose_python_scripts/ConvertMP4ToStackOfImages.py at main · stubbslab/general_purpose_python_scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import cv2
import time
import os
import numpy as np
import cantrips as can
import matplotlib.pyplot as plt

def video_to_frames(input_video_file, read_dir, save_dir = None, save_file_type = 'png', max_n_frames_to_convert = np.inf, color_save_suffixes = ['r','g','b']):
    """Function to extract frames from input video file
    and save them as separate frames in an output directory.
    Args:
        input_loc: Input video file.
        output_loc: Output directory to save the frames.
    Returns:
        None
    """
    if save_dir == None:
        save_dir = read_dir
    # Log the time
    time_start = time.time()
    # Start capturing the feed
    cap = cv2.VideoCapture(read_dir + input_video_file)
    # Find the number of frames
    video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 1
    print ("Number of frames: ", video_length)
    if max_n_frames_to_convert < video_length:
        print ('WARNING: There are more frames (' + str(video_length) + ') than the specified maximum number of frames (' + str(max_n_frames_to_convert) + ').  We will only save the first ' + str(max_n_frames_to_convert) + 'frames')
    count = 0
    print ("Converting video..\n")
    # Start converting the video
    while cap.isOpened():
        # Extract the frame
        ret, frame = cap.read()
        if not ret:
            continue
        # Write the results back to output location.
        save_name =  input_video_file.split('.')[0] + '_' + "%#05d."  % (count+1) + save_file_type
        if save_file_type.lower() == 'fits':
            #bw_frame = np.array(frame[:, :, 0]) #+ np.array(frame[:, :, 1]) + np.array(frame[:, :, 1])
            #bw_frame = np.flip(bw_frame, axis = 0)
            #bw_frame = bw_frame * 0.0 + 1
            #bw_frame = np.zeros((100, 100)) + 2
            #bw_frame = np.array([[bw_frame[i,j] * i + j for i in range(np.shape(bw_frame)[0])] for j in range(np.shape(bw_frame)[1])])
            for i in range(3):
                single_color_frame = np.flip(frame[:, :, i], axis  = 0)
                can.saveDataToFitsFile(np.transpose(single_color_frame), save_name.split('.')[0] + '_' + color_save_suffixes[i] + '.' + save_file_type, save_dir, header = 'default')
                plt.imshow(single_color_frame)
                plt.show()
        else:
            print ('save_name = ' + str(save_name))
            cv2.imwrite(save_dir + save_name, frame)
        count = count + 1
        # If there are no more frames left
        if (count > (video_length-1) or count > max_n_frames_to_convert):

            # Log the time again
            time_end = time.time()
            # Release the feed
            cap.release()
            # Print stats
            print ("Done extracting frames.\n%d frames extracted" % count)
            print ("It took %d seconds forconversion." % (time_end-time_start))
            break

if __name__=="__main__":

    input_file = 'C0007.MP4'
    input_file_dir = '/Users/elanaurbach/Desktop/SONY RX10 Test/'
    max_n_frames = 10 #np.inf
    save_file_type = 'fits' #png
    video_to_frames(input_file, input_file_dir, max_n_frames_to_convert = max_n_frames, save_file_type = save_file_type)