Thong Nguyen commited on
Commit
3a74c08
·
1 Parent(s): 8effd64

add keyframe detector files

Browse files
video_keyframe_detector/KeyFrameDetector/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from video_keyframe_detector.KeyFrameDetector import key_frame_detector
video_keyframe_detector/KeyFrameDetector/key_frame_detector.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import csv
4
+ import numpy as np
5
+ import time
6
+ import peakutils
7
+ from video_keyframe_detector.KeyFrameDetector.utils import convert_frame_to_grayscale, prepare_dirs, plot_metrics
8
+
9
+ def keyframeDetection(source, dest, Thres, plotMetrics=False, verbose=False):
10
+
11
+ keyframePath = dest+'/keyFrames'
12
+ imageGridsPath = dest+'/imageGrids'
13
+ csvPath = dest+'/csvFile'
14
+ path2file = csvPath + '/output.csv'
15
+ prepare_dirs(keyframePath, imageGridsPath, csvPath)
16
+
17
+ cap = cv2.VideoCapture(source)
18
+ length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
19
+
20
+ if (cap.isOpened()== False):
21
+ print("Error opening video file")
22
+
23
+ lstfrm = []
24
+ lstdiffMag = []
25
+ timeSpans = []
26
+ images = []
27
+ full_color = []
28
+ lastFrame = None
29
+ Start_time = time.process_time()
30
+
31
+ # Read until video is completed
32
+ for i in range(length):
33
+ ret, frame = cap.read()
34
+ grayframe, blur_gray = convert_frame_to_grayscale(frame)
35
+
36
+ frame_number = cap.get(cv2.CAP_PROP_POS_FRAMES) - 1
37
+ lstfrm.append(frame_number)
38
+ images.append(grayframe)
39
+ full_color.append(frame)
40
+ if frame_number == 0:
41
+ lastFrame = blur_gray
42
+
43
+ diff = cv2.subtract(blur_gray, lastFrame)
44
+ diffMag = cv2.countNonZero(diff)
45
+ lstdiffMag.append(diffMag)
46
+ stop_time = time.process_time()
47
+ time_Span = stop_time-Start_time
48
+ timeSpans.append(time_Span)
49
+ lastFrame = blur_gray
50
+
51
+ cap.release()
52
+ y = np.array(lstdiffMag)
53
+ base = peakutils.baseline(y, 2)
54
+ indices = peakutils.indexes(y-base, Thres, min_dist=1)
55
+
56
+ ##plot to monitor the selected keyframe
57
+ if (plotMetrics):
58
+ plot_metrics(indices, lstfrm, lstdiffMag)
59
+
60
+ cnt = 1
61
+ for x in indices:
62
+ cv2.imwrite(os.path.join(keyframePath , 'keyframe'+ str(cnt) +'.jpg'), full_color[x])
63
+ cnt +=1
64
+ log_message = 'keyframe ' + str(cnt) + ' happened at ' + str(timeSpans[x]) + ' sec.'
65
+ if(verbose):
66
+ print(log_message)
67
+ with open(path2file, 'w') as csvFile:
68
+ writer = csv.writer(csvFile)
69
+ writer.writerows(log_message)
70
+ csvFile.close()
71
+
72
+ cv2.destroyAllWindows()
video_keyframe_detector/KeyFrameDetector/utils.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import matplotlib.pyplot as plt
4
+ from PIL import Image
5
+
6
+ def scale(img, xScale, yScale):
7
+ res = cv2.resize(img, None, fx=xScale, fy=yScale, interpolation=cv2.INTER_AREA)
8
+ return res
9
+
10
+
11
+ def crop(infile, height, width):
12
+ im = Image.open(infile)
13
+ imgwidth, imgheight = im.size
14
+ for i in range(imgheight // height):
15
+ for j in range(imgwidth // width):
16
+ box = (j * width, i * height, (j + 1) * width, (i + 1) * height)
17
+ yield im.crop(box)
18
+
19
+
20
+ def averagePixels(path):
21
+ r, g, b = 0, 0, 0
22
+ count = 0
23
+ pic = Image.open(path)
24
+ for x in range(pic.size[0]):
25
+ for y in range(pic.size[1]):
26
+ imgData = pic.load()
27
+ tempr, tempg, tempb = imgData[x, y]
28
+ r += tempr
29
+ g += tempg
30
+ b += tempb
31
+ count += 1
32
+ return (r / count), (g / count), (b / count), count
33
+
34
+ def convert_frame_to_grayscale(frame):
35
+ grayframe = None
36
+ gray = None
37
+ if frame is not None:
38
+ cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
39
+ gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
40
+ gray = scale(gray, 1, 1)
41
+ grayframe = scale(gray, 1, 1)
42
+ gray = cv2.GaussianBlur(gray, (9, 9), 0.0)
43
+ return grayframe, gray
44
+
45
+ def prepare_dirs(keyframePath, imageGridsPath, csvPath):
46
+ if not os.path.exists(keyframePath):
47
+ os.makedirs(keyframePath)
48
+ if not os.path.exists(imageGridsPath):
49
+ os.makedirs(imageGridsPath)
50
+ if not os.path.exists(csvPath):
51
+ os.makedirs(csvPath)
52
+
53
+
54
+ def plot_metrics(indices, lstfrm, lstdiffMag):
55
+ plt.plot(indices, y[indices], "x")
56
+ l = plt.plot(lstfrm, lstdiffMag, 'r-')
57
+ plt.xlabel('frames')
58
+ plt.ylabel('pixel difference')
59
+ plt.title("Pixel value differences from frame to frame and the peak values")
60
+ plt.show()
video_keyframe_detector/cli.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+
3
+ from video_keyframe_detector.KeyFrameDetector.key_frame_detector import keyframeDetection
4
+
5
+ def main():
6
+ parser = argparse.ArgumentParser()
7
+
8
+ parser.add_argument('-s','--source', help='source file', required=True)
9
+ parser.add_argument('-d', '--dest', help='destination folder', required=True)
10
+ parser.add_argument('-t','--Thres', help='Threshold of the image difference', default=0.3)
11
+
12
+ args = parser.parse_args()
13
+
14
+
15
+ keyframeDetection(args.source, args.dest, float(args.Thres))
16
+
17
+
18
+ if __name__ == '__main__':
19
+ main()
video_keyframe_detector/download_video.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ from pytube import YouTube
3
+
4
+
5
+ def main():
6
+ parser = argparse.ArgumentParser()
7
+ parser.add_argument("--video-url", type=str)
8
+ parser.add_argument("--save-file", type=str)
9
+ args = parser.parse_args()
10
+
11
+ youtube = YouTube(args.video_url)
12
+ youtube.streams.filter(file_extension='mp4')
13
+ youtube.streams.get_by_itag(18).download(filename=args.save_file)
14
+
15
+
16
+ if __name__ == '__main__':
17
+ main()
video_keyframe_detector/setup.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from setuptools import setup, find_packages
2
+
3
+ setup(
4
+ name="key-frame-detector",
5
+ version="1.0.0",
6
+ author="Joel Ibaceta",
7
+ author_email="mail@joelibaceta.com",
8
+ license='MIT',
9
+ description="It is a simple python tool to extract key frame images from a video file",
10
+ long_description="A simple tool to detect and extract key frame images from a video file",
11
+ url="https://github.com/joelibaceta/video-keyframe-extractor",
12
+ project_urls={
13
+ 'Source': 'https://github.com/joelibaceta/video-keyframe-extractor',
14
+ 'Tracker': 'https://github.com/joelibaceta/video-keyframe-extractor/issues'
15
+ },
16
+ packages=find_packages(),
17
+ include_package_data=True,
18
+ install_requires=['opencv-python', 'numpy', 'peakutils'],
19
+ classifiers=[
20
+ "Programming Language :: Python :: 3",
21
+ "Intended Audience :: Developers",
22
+ "License :: OSI Approved :: MIT License",
23
+ "Operating System :: OS Independent",
24
+ ],
25
+ keywords='video key-frame terminal opencv extractor',
26
+ entry_points={
27
+ "console_scripts": [
28
+ 'key-frames-detector=cli:main'
29
+ ]
30
+ }
31
+ )