Code for calculating FFT and radial profile of an image

fft.py

#!/usr/bin/env python3
#
# Code for calculating FFT and radial profile of an image
#
# Implements a method similar to:
# Katsavounidis, I., Aaron, A., & Ronca, D. (2015, October). Native resolution detection
# of video sequences. In SMPTE 2015 annual technical conference and exhibition (pp. 1-20).
# SMPTE.
#
# Author: Werner Robitza
#
# Based on work from Julian Zebelein, Steve Göring, TU Ilmenau,
# and various Internet examples
#
# Requirements: pip3 install opencv-python matplotlib tqdm numpy

import cv2
import numpy as np
from matplotlib import pyplot as plt
import sys
from argparse import ArgumentParser
import json
import os
from tqdm import tqdm

# copied from stackoverflow
def radial_profile(data, center):
    y, x = np.indices((data.shape))
    r = np.sqrt((x - center[0]) ** 2 + (y - center[1]) ** 2)
    r = r.astype(np.int)

    tbin = np.bincount(r.ravel(), data.ravel())
    nr = np.bincount(r.ravel())
    radialprofile = tbin / nr
    return radialprofile


def calc_fft(input_file, frame_limit=None, output_first_frame=False, output_avg=False):
    cap = cv2.VideoCapture(input_file)

    number_of_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if frame_limit:
        frame_limit = min(frame_limit, number_of_frames)
    else:
        frame_limit = number_of_frames

    # Generate empty list fft sum values
    high_frequencies = []
    avg_magnitude_spectrum = None

    # start video
    frames_left = frame_limit

    image_saved = False
    fig_prefix = os.path.splitext(os.path.basename(input_file))[0]

    pbar = tqdm(total=frame_limit, file=sys.stderr)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # prefinal = cv2.resize(gray, (3840, 2160))
        # final = cv2.bilateralFilter(prefinal, 9, 75, 75)

        f = np.fft.fft2(gray)
        fshift = np.fft.fftshift(f)
        magnitude_spectrum = 20 * np.log(np.abs(fshift))

        if avg_magnitude_spectrum is None:
            avg_magnitude_spectrum = magnitude_spectrum
        else:
            avg_magnitude_spectrum = np.mean(np.array([avg_magnitude_spectrum, magnitude_spectrum]), axis=0)

        file_width = int(cap.get(3))
        file_height = int(cap.get(4))

        if output_first_frame and not image_saved:
            plt.figure(figsize=(file_width / 100, file_height / 100), dpi=72)
            plt.imshow(magnitude_spectrum, cmap="gray")
            plt.savefig(fig_prefix + "_first-frame.png", bbox_inches='tight')
            image_saved = True

        # file_height, file_width = magnitude_spectrum.shape
        center = (file_width / 2, file_height / 2)

        # Calculate the azimuthally averaged 1D power spectrum
        psf_1d = radial_profile(magnitude_spectrum, center)
        low_freq_ind = int((len(psf_1d) / 2))

        sum_of_high_frequencies = sum(psf_1d[low_freq_ind:])

        high_frequencies.append(sum_of_high_frequencies)

        frames_left -= 1
        if frames_left == 0:
            break

        k = cv2.waitKey(30) & 0xFF
        if k == 27:
            break

        pbar.update()

    pbar.update()
    pbar.close()

    if output_avg:
        plt.figure(figsize=(file_width / 100, file_height / 100), dpi=72)
        plt.imshow(avg_magnitude_spectrum, cmap="gray")
        plt.savefig(fig_prefix + "-avg.png", bbox_inches='tight')

    ret = {
        'input_file': input_file,
        'average_fft': int((sum(int(i) for i in high_frequencies)) / len(high_frequencies)),
        'max_fft': int(max(high_frequencies)),
        'min_fft': int(min(high_frequencies)),
        'median_fft': int(np.median(high_frequencies)),
        'percentile_five': int(high_frequencies[int((5 / 100) * len(high_frequencies))]),
        'percentile_nine_five': int(high_frequencies[int((95 / 100) * len(high_frequencies))])
    }

    print(json.dumps(ret, indent=4))


def main():
    parser = ArgumentParser()
    parser.add_argument("input")
    parser.add_argument("-l", "--frame-limit", type=int)
    parser.add_argument("-of", "--output-first-frame", action="store_true")
    parser.add_argument("-oa", "--output-avg", action="store_true")
    cli_args = parser.parse_args()

    calc_fft(
        cli_args.input,
        frame_limit=cli_args.frame_limit,
        output_first_frame=cli_args.output_first_frame,
        output_avg=cli_args.output_avg
    )


if __name__ == "__main__":
    main()