1  # SPDX-FileCopyrightText: 2020 Carter Nelson for Adafruit Industries
 2  #
 3  # SPDX-License-Identifier: MIT
 4  
 5  import time
 6  import imageio
 7  from skimage.color import rgb2gray
 8  import matplotlib.pyplot as plt
 9  import numpy as np
10  
11  THRESH = 0.3
12  
13  RUN = int(input('Enter run number: '))
14  
15  vid = imageio.get_reader('run_{:03d}.mp4'.format(RUN), 'ffmpeg')
16  
17  #----------------
18  # MAIN PROCESSING
19  #----------------
20  frame_data = []
21  start = time.monotonic()
22  # go through video frame by frame
23  print("Processing", end='')
24  for frame in vid:
25      print('.', end='', flush=True)
26      frame_bin = rgb2gray(frame) > THRESH
27      frame_count = np.count_nonzero(frame_bin == True)
28      frame_percent = 100 * frame_count / (1920*1080)
29      frame_data.append((frame_count, frame_percent))
30  # overall stats
31  avg_count = sum([x[0] for x in frame_data]) / len(frame_data)
32  avg_percent = 100 * avg_count / (1920*1080)
33  
34  end = time.monotonic()
35  print("\nProcessing done in {} secs.".format(end - start))
36  print("Average Count = {}".format(avg_count))
37  print("Average Percent = {}".format(avg_percent))
38  
39  #-------------
40  # SAVE TO FILE
41  #-------------
42  print("Saving data to file...")
43  with open('run_{:03d}.csv'.format(RUN), 'w') as fp:
44      for frame, data in enumerate(frame_data):
45          fp.write('{},{},{}\n'.format(frame, data[0], data[1]))
46  
47  #---------
48  # PLOTTING
49  #---------
50  print("Generating plots...")
51  fig, ax = plt.subplots(1, figsize = (10,5))
52  ax.set_title("RUN {:03d}\nTHRESH = {}, AVG_CNT = {:4.2}, AVG_PER = {:.3}".format(RUN, THRESH,avg_count, avg_percent))
53  ax.set_xlabel("FRAME")
54  ax.set_ylabel("COUNT")
55  ax.plot([x[0] for x in frame_data])
56  fig.savefig('run_{:03d}_plot.png'.format(RUN))
57  
58  print("DONE.")