quick_fft.py 4.27 KB
Newer Older
1 2
import matplotlib
matplotlib.use('TkAgg')
Robozman's avatar
Robozman committed
3 4 5
import time
import matplotlib.pyplot as plt
import numpy as np
Robozman's avatar
Robozman committed
6
import scipy
Robozman's avatar
Robozman committed
7
import scipy.io.wavfile as wavfile
Robozman's avatar
Robozman committed
8
from scipy import signal
Robozman's avatar
Robozman committed
9 10
import sounddevice
import time
Samuel Warfield's avatar
Samuel Warfield committed
11
import math
Jackson Garner's avatar
Jackson Garner committed
12 13
import msaf
import librosa
Robozman's avatar
Robozman committed
14
import sys
15
import pickle
16
import cv2
Robozman's avatar
Robozman committed
17

Robozman's avatar
Robozman committed
18
import librosaToWave
Robozman's avatar
Robozman committed
19
import generate_waves
Jackson Garner's avatar
Jackson Garner committed
20
import fftToImage
21
import ImageFunctions
Robozman's avatar
Robozman committed
22

Robozman's avatar
Robozman committed
23 24 25 26
load_from_file = False
file_name = None
wave_type = 0
fs=44100
Robozman's avatar
Robozman committed
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

def mkBounds(name):
    try:
        f = open(name + ".cache", "rb")
        return pickle.load(f)
    except:
        of = open(name + ".cache", "wb")
        result = msaf.process(name)
        pickle.dump(result, of)
        return result


def mkBeat(name):
    try:
        f = open(name + ".beat", "rb")
        return pickle.load(f)
    except:
        of = open(name + ".beat", "wb")
        audio = librosa.load(name, sr=None)[0]
        audio_harmonic, audio_percussive = librosa.effects.hpss(audio)
        tempo, frames = librosa.beat.beat_track(y=audio_harmonic,
                                                sr=fs, hop_length=1024)
        result = librosa.frames_to_time(frames, sr=fs,
                                        hop_length=1024)
        pickle.dump(result, of)
        return result

Robozman's avatar
Robozman committed
55 56 57 58 59
if len(sys.argv) == 1:
    print("Provide a file name or an integer as an argument")
    exit()
else:
    try:
60
        wave_type = int(sys.argv[1])
Robozman's avatar
Robozman committed
61 62 63
    except:
        load_from_file = True
        file_name = sys.argv[1]
Robozman's avatar
Robozman committed
64

Robozman's avatar
Robozman committed
65 66 67
if load_from_file == False:
    wav_data, rate = generate_waves.generate_waves(wave_type, 52)
elif load_from_file == True:
Robozman's avatar
Robozman committed
68
    #rate, wav_data = wavfile.read(sys.argv[1], mmap=True)
Robozman's avatar
Robozman committed
69
    # wav_data = samples
Robozman's avatar
Robozman committed
70
    wav_data, rate = librosaToWave.getWave(fileName=sys.argv[1])
71 72
    boundaries, labels = mkBounds(sys.argv[1])
    beat_times = mkBeat(sys.argv[1])
Robozman's avatar
Robozman committed
73 74 75

plt.ion()
fig = plt.figure()
76 77
fft_plot = fig.add_subplot(121)
#image = fig.add_subplot(122)
Robozman's avatar
Robozman committed
78

Jackson Garner's avatar
Jackson Garner committed
79
image = np.ones((720, 1280, 3))
Jackson Garner's avatar
Jackson Garner committed
80

Robozman's avatar
Robozman committed
81 82
sounddevice.play(wav_data, rate)

83
timestep = .01
84 85 86
    

timestep = 1/24
Robozman's avatar
Robozman committed
87
t = time.time()
Samuel Warfield's avatar
Samuel Warfield committed
88
origin = time.time()
Jackson Garner's avatar
Jackson Garner committed
89
nextBeat = 0
Robozman's avatar
Robozman committed
90 91
cap = cv2.VideoCapture('snow.webm')
is_beat = False
Samuel Warfield's avatar
Samuel Warfield committed
92
while t - origin < len(wav_data) / rate:
93
    ret, frame = cap.read()
Robozman's avatar
Robozman committed
94
    if load_from_file == True:
95 96 97 98
        try:
            if time.time() - origin >= beat_times[nextBeat]:
                nextBeat += 1
                print(f"Beat {t - origin}")
Robozman's avatar
Robozman committed
99
                is_beat = True
100 101 102
        except:
            print("Last beat encountered")

Samuel Warfield's avatar
Samuel Warfield committed
103 104

    while (time.time() - t < timestep):
Robozman's avatar
Robozman committed
105
        pass
Samuel Warfield's avatar
Samuel Warfield committed
106

Robozman's avatar
Robozman committed
107
    t = time.time()
Samuel Warfield's avatar
Samuel Warfield committed
108 109 110

    second = wav_data[math.floor((t - origin) * rate) : math.floor((t - origin + timestep) * rate)]

Jackson Garner's avatar
Jackson Garner committed
111
    # print(second)
Samuel Warfield's avatar
Samuel Warfield committed
112

113 114
    transform = np.abs(np.fft.fft(second[:]))
    plot_transform = np.fft.fftshift(transform)
115
    print(ImageFunctions.getBaseWeight(transform))
116 117 118 119 120
    max_val = np.amax(transform)
    hsv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv_frame)
    #v = np.where(True, 255 * (350 - max_val) / 350, 255)
    #v *= math.floor((350 - max_val) / 350)
Robozman's avatar
Robozman committed
121
    np.multiply(v, (((350 - max_val) / 350) + 0.5*5)/6, out=v, casting='unsafe')
122 123 124 125 126
    final_hsv = cv2.merge((h, s, v))
    img = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
    #image.xticks([])
    #image.ytacks([])
    #image.show()
Jackson Garner's avatar
Jackson Garner committed
127
    a = np.abs(np.fft.fftshift(np.fft.fft(second[:])))
Robozman's avatar
Robozman committed
128 129 130
    #fftToImage.fftToImage(a, image)
    #cv2.imshow("t", (img * (1/255) + image).clip(0, 1))
    preV = (img * (1/255)).clip(0, 1)
Jackson Garner's avatar
Jackson Garner committed
131 132
    row,cols,rgb = img.shape

Robozman's avatar
Robozman committed
133 134
    a = cv2.getGaussianKernel(cols,300.0)
    b = cv2.getGaussianKernel(row,300.0)
Jackson Garner's avatar
Jackson Garner committed
135 136 137
    c = b*a.T
    d = c/c.max()
    d = np.array([d, d, d]).transpose(1, 2, 0)
Robozman's avatar
Robozman committed
138 139
    preV = preV *  ( d * 10)
    #print(preV)
Jackson Garner's avatar
Jackson Garner committed
140
    cv2.imshow('frame', preV)
141
    cv2.waitKey(1)
Jackson Garner's avatar
Jackson Garner committed
142 143
    # cv2.imshow('frame', img)
    # cv2.waitKey(1)
144 145 146

    peaks, _ = signal.find_peaks(transform, prominence=20)

Robozman's avatar
Robozman committed
147 148
    is_beat = False

149
    plot = False
150 151

    if plot is True:
152
        print(peaks)
153 154 155 156 157
        fft_plot.clear()
        if load_from_file == True:
            fft_plot.plot(plot_transform) # np.abs(np.fft.fftshift(np.fft.fft(second[:])))
        else:
            fft_plot.plot(plot_transform) # np.abs(np.fft.fftshift(np.fft.fft(second[:])))
Samuel Warfield's avatar
Samuel Warfield committed
158

159 160
        for x in peaks:
            fft_plot.plot(x,plot_transform[x], marker='x')
Samuel Warfield's avatar
Samuel Warfield committed
161

162
        fft_plot.set_ylim([0,300])
Samuel Warfield's avatar
Samuel Warfield committed
163

164 165
        fig.canvas.draw()
        fig.canvas.flush_events()