Record sound and do spectrogram to the sound signal

spectrogram.py

import numpy as np
from scipy.io.wavfile import read,write
import matplotlib.pyplot as plt 
from matplotlib.pylab import *

def omega(N,k,n):
	return np.exp(-2*np.pi/N*k*n*1j)

def hanWindow(N):
	diag=np.array([np.sin(np.pi*i/(N-1))**2 for i in range(N)])
	H=np.diag(diag)
	return H

def createF(N):
	F=np.zeros((N,N),dtype=complex)
	for i in range(N):
		for j in range(N):
			F[i,j]=omega(N,j,i)
	return F

def createA(N,input_length,hop):
	FH=np.dot(createF(N),hanWindow(N))
	A=np.zeros(((2*input_length-1)*N,input_length*N),dtype=complex)

	for i in range(2*input_length-1):
		A[i*N:(i+1)*N,i*hop:i*hop+N]=FH
	return A

def DFT(N,hop,input_data):
	'''
	N: DFT window
	hop: hop size
	input_data: signal
	'''
	input_length=int(len(input_data)/N)
	print "Using BIG matrix to do spectrogram analysis"
	print "The size of input data is:",input_length*N
	A=createA(N,input_length,hop)
	plot_matrix(A)

	DFT_coeff=np.dot(A,input_data[:input_length*N])
	
	spectrogram=[]
	for i in range(2*input_length-1):
		spectrogram.append(np.abs(DFT_coeff[i*N:i*N+N/2]))
	
	spectrogram=np.array(spectrogram)
	return spectrogram

def DFT_sliding(N,hop,input_data):
	'''
	N: DFT window
	hop: hop size
	input_data: signal
	'''
	input_length=int(len(input_data)/N)
	print "Using sliding method to do spectrogram analysis"
	print "The size of input data is:",input_length*N
	spectrogram=[]
	H=hanWindow(N)
	FH=np.dot(createF(N),H)

	for i in range((2*input_length-1)):
		signal=input_data[i*N/2:i*N/2+N]
		DFT_coeff=np.dot(FH,signal)
		# Using numpy as validation
		DFT_coeff_numpy=np.fft.fft(np.dot(H,signal))
		print "The value should be small:(validation with numpy.fft.fft)",sum(abs(DFT_coeff-DFT_coeff_numpy))
		spectrogram.append(np.abs(DFT_coeff)[:512])
	
	spectrogram=np.array(spectrogram)
	return spectrogram

def record(time):
	# This record function is from http://people.csail.mit.edu/hubert/pyaudio/
	import pyaudio
	import wave

	CHUNK = 1024
	FORMAT = pyaudio.paInt16
	CHANNELS = 2
	RATE = 44100
	RECORD_SECONDS = time
	WAVE_OUTPUT_FILENAME = "output.wav"

	p = pyaudio.PyAudio()

	stream = p.open(format=FORMAT,
	                channels=CHANNELS,
	                rate=RATE,
	                input=True,
	                frames_per_buffer=CHUNK)

	print("* recording")

	frames = []

	for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
	    data = stream.read(CHUNK)
	    frames.append(data)

	print("* done recording")

	stream.stop_stream()
	stream.close()
	p.terminate()

	wf = wave.open(WAVE_OUTPUT_FILENAME, 'wb')
	wf.setnchannels(CHANNELS)
	wf.setsampwidth(p.get_sample_size(FORMAT))
	wf.setframerate(RATE)
	wf.writeframes(b''.join(frames))
	wf.close()

def plot_matrix(matrix):
	
	# print "matrix size,",matrix.shape
	matrix_abs=np.abs(matrix)
	print np.amax(matrix_abs)
	# matrix_abs=np.random.rand(1024,1024)
	plt.figure()
	# matshow(matrix_abs)
	# print matrix_abs
	# plt.pcolor(matrix_abs)
	plt.imshow(matrix_abs)
	plt.xlabel("Column")
	plt.ylabel("Row")
	plt.colorbar()
	plt.savefig("A_matrix.png")

	print "A matrix plotted"
	# print matrix_abs


def plot_spectrogram(spectrogram_coeff,filname):
	plt.figure()
	plt.pcolor(np.transpose(spectrogram_coeff))
	plt.xlabel("Time (* 512)")
	plt.ylabel("Frequency")
	plt.colorbar()
	plt.savefig(filname+".png")

	# matshow(spectrogram_coeff)
def main():
	Fs = 44100;  # sampling rate
	rate,data=read("output.wav")
	y=data[:,1]
	lungime=len(y)
	timp=len(y)/44100.
	t=np.linspace(0,timp,len(y))
	y=np.array(y)
	
	### Using big matrix to do analysis
	spectrogram=DFT(1024,512,y[:3*1024])
	plot_spectrogram(spectrogram,"method1")

	### Using sliding method to do analysis
	spectrogram=DFT_sliding(1024,512,y[:40*1024])
	plot_spectrogram(spectrogram,"method2")


if __name__ == '__main__':
	record(4)
	main()