Amazing-Python-Scripts

1
import os
2
from os import path
3
import math
4
import cv2
5
import argparse
6

7

8
class Extractor:
9
    """Extract frames from video file and saves them in a directory
10

11
    Arguments:
12
            vid_filename {str} -- Location of video file
13
            output_dir {str} -- Output directory to store frames
14

15

16
    """
17

18
    global args
19

20
    def __init__(self, vid_filename, output_dir, verbose=False):
21

22
        self.verbose = verbose
23
        self.output_dir = output_dir
24
        if not os.path.exists(self.output_dir):
25
            os.makedirs(self.output_dir)
26
        print(self.output_dir)
27
        if path.exists(vid_filename):
28
            self.vid_filename = vid_filename
29
        else:
30
            raise FileNotFoundError(f"File {vid_filename} not found")
31

32
        # initialize video capture
33
        self.cap = cv2.VideoCapture(self.vid_filename)
34

35
        # get video fps
36
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
37

38
        # Video length in frames
39
        self.vid_length = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
40

41
    def extractor(self):
42
        is_ok, frame = self.cap.read()
43
        count = 0
44
        print(self.fps)
45
        while is_ok:
46
            # is_ok, frame = self.cap.read()
47
            print(
48
                f"Extracting frame {count} of {self.vid_length-1}") if self.verbose else {}
49
            # save current frame
50
            cv2.imwrite(os.path.join(
51
                self.output_dir, f"{count}.jpg"), frame)
52
            print("COMPLETED") if count == self.vid_length-1 else {}
53

54
            is_ok, frame = self.cap.read()
55

56
            count += 1
57

58
    def extract_frames(file):
59
        global args
60

61
        # check whether file is greather than 0 bytes using path
62
        if os.stat(path.join(args.input_dir, file)).st_size > 0:
63
            # extract frames
64
            extractor = Extractor(path.join(args.input_dir, file), path.join(
65
                args.output_dir, file), args.sampling_rate)
66
            extractor.extractor()
67
        else:
68
            print(f"File {file} is empty")
69

70

71
def main():
72
    # setup parsers for CLI arguments
73
    parser = argparse.ArgumentParser("Frame Extractor")
74
    parser.add_argument("-i", "--video_file", type=str,
75
                        help="Location of video file", required=True)
76
    parser.add_argument("-o", "--output_dir", type=str,
77
                        help="Output directory to store frames", required=True)
78
    parser.add_argument("-v", "--verbose", type=bool,
79
                        help="Display details of extraction of frames", default=False)
80
    args = parser.parse_args()
81
    # print(args)
82
    # Extract frames from video file
83
    if args.video_file:
84
        name = args.video_file.split(".")[0]
85
        ext = args.video_file.split(".")[1]
86
        if ext != "mp4":
87
            raise argparse.ArgumentTypeError(
88
                f"File {args.video_file} is not a video(.mp4) file")
89
        output_dir = os.path.join(args.output_dir, name)
90
        extract = Extractor(args.video_file, args.output_dir,  args.verbose)
91
        extract.extractor()
92

93

94
if __name__ == '__main__':
95
    main()
96