forked from keplerlab/katna
-
Notifications
You must be signed in to change notification settings - Fork 0
/
example_video.py
82 lines (58 loc) · 1.92 KB
/
example_video.py
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
72
73
74
75
76
77
78
79
80
81
82
import os
import os.path
import cv2
import sys, getopt
from Katna.video import Video
from Katna.writer import KeyFrameDiskWriter
import multiprocessing
import ntpath
#from icecream import install
#install()
class CustomDiskWriter(KeyFrameDiskWriter):
"""
:param KeyFrameDiskWriter: Writer class to overwrite
:type KeyFrameDiskWriter: Writer
"""
def generate_output_filename(self, filepath, keyframe_number):
"""Custom output filename method
:param filepath: [description]
:type filepath: [type]
"""
filename = super().generate_output_filename(filepath, keyframe_number)
suffix = "keyframe"
return "_".join([filename, suffix])
def main_dir():
if len(sys.argv) ==1:
dir_path = os.path.join(".", "tests", "data")
else:
dir_path = sys.argv[1]
vd = Video()
no_of_frames_to_returned = 12
diskwriter = KeyFrameDiskWriter(location="selectedframes")
vd.extract_keyframes_from_videos_dir(
no_of_frames=no_of_frames_to_returned, dir_path=dir_path,
writer=diskwriter
)
def main():
# Extract specific number of key frames from video
# if os.name == 'nt':
# multiprocessing.freeze_support()
if len(sys.argv) ==1:
video_file_path = os.path.join(".", "tests", "data", "pos_video.mp4")
else:
video_file_path = sys.argv[1]
vd = Video()
# number of images to be returned
no_of_frames_to_returned = 12
diskwriter = KeyFrameDiskWriter(location="selectedframes")
# VIdeo file path
#video_file_path = os.path.join(".", "tests", "data", "pos_video.mp4")
print(f"Input video file path = {video_file_path}")
vd.extract_video_keyframes(
no_of_frames=no_of_frames_to_returned, file_path=video_file_path,
writer=diskwriter
)
if __name__ == "__main__":
multiprocessing.set_start_method("spawn")
main()
#main_dir()