-
Notifications
You must be signed in to change notification settings - Fork 0
/
lightning.py
122 lines (89 loc) · 3.76 KB
/
lightning.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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import os
import datetime
import argparse
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from progress.spinner import Spinner
import ntpath
def make_safe(unsafe_string):
return "".join([c for c in unsafe_string if c.isalpha() or c.isdigit()]).rstrip()
parser = argparse.ArgumentParser(description='Extract lightning frames from a video.')
parser.add_argument('video_file_name', type=str,
help='The file with the lightning in it')
parser.add_argument('--threshold', dest='threshold', action='store',
default=10,
help='Use a non-default (default is 10) threshold for detirmining what a lightning flash is.')
parser.add_argument('--outfolder', dest='outfolder', action='store',
help='Specify a folder for frames & data to be saved to.')
args = parser.parse_args()
THRESHOLD = float(args.threshold)
VIDEO_FILE_NAME = args.video_file_name
print(f"Using Threshold: {THRESHOLD}")
print(f"Using Video File Name: {VIDEO_FILE_NAME}")
if __name__ == '__main__':
if not os.path.isfile(VIDEO_FILE_NAME):
print(f"File not found: {VIDEO_FILE_NAME}")
print("Exiting...")
exit(404)
if not args.outfolder:
OUTFOLDER = f"{make_safe(ntpath.basename(VIDEO_FILE_NAME))}__OUTPUT"
if ntpath.dirname(VIDEO_FILE_NAME):
OUTFOLDER = f"{ntpath.dirname(VIDEO_FILE_NAME)}/{OUTFOLDER}"
else:
OUTFOLDER = args.outfolder
print(f"Output going to folder: {OUTFOLDER}")
print(f"Starting at: {datetime.datetime.now().isoformat()}")
if not os.path.isdir(OUTFOLDER):
os.makedirs(OUTFOLDER, exist_ok=True)
cap = cv2.VideoCapture(VIDEO_FILE_NAME)
frame_num = 0
frame_data = []
spinner = Spinner('Processing ')
last_mean_brightness = None
frames_to_keep = {}
temp_rec_count = 0
temp_rec = False
while (cap.isOpened()):
ret, frame = cap.read()
if not ret:
print(f"Looks like we are done at Frame number: {frame_num}")
break
mean_brightness = np.mean(frame)
if last_mean_brightness is not None:
diff_brightness = abs(mean_brightness-last_mean_brightness)
percent_diff = (diff_brightness/mean_brightness)*100
if percent_diff > THRESHOLD:
temp_rec_count = 4
temp_rec = True
if temp_rec or temp_rec_count > 0:
frames_to_keep[frame_num] = frame
print(f"Added frame number: {frame_num}, % diff={percent_diff}, temp_rec_count={temp_rec_count}")
temp_rec_count -= 1
temp_rec = False
frame_data.append([frame_num, mean_brightness, percent_diff])
frame_num += 1
last_mean_brightness = mean_brightness
last_frame_num = frame_num
last_frame = frame
spinner.next()
for frame_id in frames_to_keep:
cv2.imwrite(f"{OUTFOLDER}/frame_{str(frame_id)}.jpg", frames_to_keep[frame_id])
cap.release()
cv2.destroyAllWindows()
print(f"Ending at: {datetime.datetime.now().isoformat()}")
if len(frame_data) == 0:
print(f"No data was found, this file location ok?:{VIDEO_FILE_NAME}")
exit(400)
df = pd.DataFrame(frame_data, columns=[
"frame_num",
"brightness",
"percent_diff"])
print(df)
df.to_csv(f"{OUTFOLDER}/frame_brighness_data.csv", columns=["frame_num",
"brightness",
"percent_diff"
], index=False)
df.plot(x="frame_num", y="brightness")
plt.savefig(f"{OUTFOLDER}/frame_brighness_data.pdf")