DJI Goggles and O3 Air Unit Overlay Processing

run.py

#!/usr/bin/env python3
import os
import configparser
from glob import glob
import ffmpeg


### INSTRUCTIONS
# 1. Run using processing.duration=60 and processing.frame_nums=true 
#    Duration should be long enough so that you have a fast movement, 
#    where counting the frame difference is possible.
# 2. Watch the overlay in a media player and see which video is lagging 
#    behind and by how many frames. Focus on some movement in the center 
#    where the spot is discernible in both videos.
# 3. Remove previous config and set processing.skip_frames.{goggles,obu}=42


### CONFIG
CFG_FILE = 'config.ini'
cfg = configparser.ConfigParser()
if os.path.isfile(CFG_FILE):
    cfg.read(CFG_FILE)

FPS = cfg.getint('processing', 'fps', fallback = 120)
DURATION = cfg.getint('processing', 'duration', fallback = None)
FRAME_NUMS = cfg.getboolean('processing', 'frame_nums', fallback = False)
# skip frames have to be applied to the video lagging behind so they start at the same time
SKIP_FRAMES_GOGGLES = cfg.getint('processing.skip_frames', 'goggles', fallback = 0)
SKIP_FRAMES_OBU = cfg.getint('processing.skip_frames', 'obu', fallback = 0)
# higher value = smaller overlay
OVERLAY_SCALE = cfg.getfloat('overlay', 'scale', fallback = 4.5)

INPUT_GOGGLES = list(sorted(glob(cfg.get('input', 'goggles', fallback = 'goggles/*.MOV'))))
INPUT_OBU = list(sorted(glob(cfg.get('input', 'goggles', fallback = 'obu/*.MP4'))))

OUTPUT_CODEC_CONCAT = cfg.get('output.concat', 'codec', fallback = 'hevc_videotoolbox')
OUTPUT_QUALITY_CONCAT = cfg.getint('output.concat', 'quality', fallback = 60)
OUTPUT_GOGGLES_CONCAT = cfg.get('output.concat', 'goggles', fallback = 'goggles_concat.mp4')
OUTPUT_OBU_CONCAT = cfg.get('output.concat', 'obu', fallback = 'obu_concat.mp4')

OUTPUT_CODEC_PREPROC = cfg.get('output.preproc', 'codec', fallback = 'hevc_videotoolbox')
OUTPUT_QUALITY_PREPROC = cfg.getint('output.preproc', 'quality', fallback = 60)
OUTPUT_GOGGLES_PREPROC = cfg.get('output.preproc', 'goggles', fallback = 'goggles_preproc.mp4')
OUTPUT_OBU_PREPROC = cfg.get('output.preproc', 'obu', fallback = 'obu_preproc.mp4')

OUTPUT_CODEC_OVERLAY = cfg.get('output.overlay', 'codec', fallback = 'hevc_videotoolbox')
OUTPUT_QUALITY_OVERLAY = cfg.getint('output.overlay', 'quality', fallback = 60)
OUTPUT_OVERLAY = cfg.get('output', 'overlay', fallback = 'overlay.mp4')


### UTILITIES
def get_stream(file, stream_type = 'video'):
    for stream in ffmpeg.probe(file)['streams']:
        if stream_type == stream['codec_type']:
            return stream
    raise


### Stage 1: CONCAT
def concat(input_files, output_file, fps = None, duration = None):
    out = ffmpeg.concat(*[ ffmpeg.input(input_file) for input_file in input_files ])
    if fps:
        out = out.filter('fps', fps = fps)
    if duration:
        out = out.trim(duration = duration)
    out.output(output_file, c = OUTPUT_CODEC_CONCAT, q = OUTPUT_QUALITY_CONCAT).run()

concat(
    INPUT_GOGGLES, OUTPUT_GOGGLES_CONCAT,
    fps = FPS,
    duration = DURATION
)
concat(
    INPUT_OBU, OUTPUT_OBU_CONCAT,
    fps = FPS,
    duration = DURATION
)


### Stage 2: PREPROCESS
def preproc(input_file, output_file, fps = None, skip_frames = 0, frame_nums = True):
    out = ffmpeg.input(
        input_file,
        ss = skip_frames / fps,
        an = None # Remove audio
    )
    if frame_nums:
        out = out.filter('drawtext', text = '%{frame_num}', x = '(w-tw)/2', y = '(h-th)/2', fontcolor = 'white', fontsize = 100)
    if fps:
        out = out.filter('fps', fps = fps)
    out.output(output_file, c = OUTPUT_CODEC_PREPROC, q = OUTPUT_QUALITY_PREPROC).run()

preproc(
    OUTPUT_GOGGLES_CONCAT, OUTPUT_GOGGLES_PREPROC,
    fps = FPS,
    skip_frames = SKIP_FRAMES_GOGGLES,
    frame_nums = FRAME_NUMS
)
preproc(
    OUTPUT_OBU_CONCAT, OUTPUT_OBU_PREPROC,
    fps = FPS,
    skip_frames = SKIP_FRAMES_OBU,
    frame_nums = FRAME_NUMS
)


### Stage 2: OVERLAY
def overlay(obu_file, goggles_file, output_file, scale):
    stream = get_stream(obu_file)
    width = stream['width']
    height = stream['height']
    overlay_width = width / scale
    overlay_height = height / scale
    (
        ffmpeg
            .overlay(
                ffmpeg.input(obu_file),
                ffmpeg.input(goggles_file)
                    .filter('scale', overlay_width, overlay_height),
                x = (width - overlay_width),
                y = (height - overlay_height),
                shortest = 1,
            )
            .output(output_file, c = OUTPUT_CODEC_OVERLAY, q = OUTPUT_QUALITY_OVERLAY)
            .run()
    )

overlay(
    OUTPUT_OBU_PREPROC,
    OUTPUT_GOGGLES_PREPROC,
    OUTPUT_OVERLAY,
    OVERLAY_SCALE
)