Produces a timelapse video of a thunderstorm with an emphasis on lightning

thunderstorm_timelapse.py

#!/usr/bin/env python
import cv2
import os.path
import sys
import tqdm

def frames(reader, indices=None, desc=None):
    frame_count = int(reader.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT))
    if indices is None:
        indices = xrange(0, frame_count)

    for frame_index in tqdm.tqdm(indices, desc=desc):
        reader.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, frame_index)
        flag, frame = reader.read()
        if flag:
            yield frame
        else:
            yield None

def brightness(img):
    if img is None:
        return 0
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    return (gray >= 250).sum()

reader = cv2.VideoCapture(sys.argv[1])

bright_pixels = []
if os.path.exists('bright_pixels.txt'):
    with open('bright_pixels.txt', 'r') as fin:
        bright_pixels = map(int, fin.readlines())
else:
    bright_pixels = map(brightness, frames(reader, desc='Processing video'))
    with open('bright_pixels.txt', 'w') as fout:
        fout.write("\n".join(map(str, bright_pixels)))

# Include frames when either:
#   1. they are in our timelapse modulo, or
#   2. a local window surrounding them contains a frame brighter than some threshold.
#
frames_to_keep = []
window_size = 10
frame_skip = 15
brightness_thresh = 1000
for index in xrange(len(bright_pixels)):
    window_start = max(0, index - window_size)
    window_end = index + window_size
    if index % frame_skip == 0 or max(bright_pixels[window_start:window_end]) >= brightness_thresh:
        frames_to_keep.append(index)

# the avc1 code is kind of slow once we hit a lot of frames, so split it up into
# a bunch of smaller videos and concatenate with ffmpeg later:
#
#   $ ffmpeg -f concat -safe 0 -i <(for x in $PWD/video*.mp4; do echo "file '$x'"; done) -c copy output.mp4
#
w = int(reader.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH) + 0.5)
h = int(reader.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT) + 0.5)
fourcc = cv2.cv.CV_FOURCC(*'avc1')

print 'Writing', len(frames_to_keep), 'frames'

num_per_video = 500
for x in xrange(0, len(frames_to_keep), num_per_video):
    indices = frames_to_keep[x:x + num_per_video]
    writer = cv2.VideoWriter('video%02d.mp4' % (x / num_per_video), fourcc, 30, (w, h))
    for frame in frames(reader, indices=indices, desc='Writing bright frames'):
        writer.write(frame)