mi@2: import os, sys
mi@2: import numpy as np
mi@2: import scipy as sp
mi@2: from scipy.ndimage.filters import *
mi@2: from vampy import *
mi@2: from copy import deepcopy
mi@2: 
mi@2: class HPspectrogram: 
mi@2: 	
mi@2: 	def __init__(self, inputSampleRate):
mi@2: 
mi@2: 		# flags:
mi@2: 		self.vampy_flags = vf_DEBUG | vf_ARRAY | vf_REALTIME
mi@2: 		self.inputSampleRate = inputSampleRate
mi@2: 		self.stepSize = 1024
mi@2: 		self.blockSize = 2048
mi@2: 		self.channels = 1
mi@2: 		self.harmonic_medwin = 8 	
mi@2: 		self.percussive_medwin = 8 	
mi@2: 		self.harmonic_thresh = 0.5
mi@2: 		self.percussive_thresh = 0.5
mi@2: 		self.whitening_on = False
mi@2: 		self.whitenRelaxCoeff = 0.9997
mi@2: 		self.whitenFloor = 0.01
mi@2: 		self.magPeaks = None
mi@2: 		
mi@2: 	def initialise(self, channels, stepSize, blockSize):
mi@2: 		self.channels = channels
mi@2: 		self.stepSize = stepSize		
mi@2: 		self.blockSize = blockSize
mi@2: 		self.maglist = []
mi@2: 		self.timestamp = []
mi@2: 		return True
mi@2: 		
mi@2: 	def reset(self):
mi@2: 		# reset any initial conditions
mi@2: 		self.maglist = []
mi@2: 		self.timestamp = []
mi@2: 		return None
mi@2: 	
mi@2: 	def getMaker(self):
mi@2: 		return 'Mi Tian'
mi@2: 	
mi@2: 	def getName(self):
mi@2: 		return 'Harmonic percussive spectrogram'
mi@2: 		
mi@2: 	def getIdentifier(self):
mi@2: 		return 'hps'
mi@2: 
mi@2: 	def getDescription(self):
mi@2: 		return 'Harmonic and percussive emphasised spectrograms.'
mi@2: 	
mi@2: 	def getCopyright(self):
mi@2: 		return ''
mi@2: 
mi@2: 	def getMaxChannelCount(self):
mi@2: 		return 1
mi@2: 		
mi@2: 	def getInputDomain(self):
mi@2: 		return FrequencyDomain
mi@2: 			
mi@2: 	def getOutputDescriptors(self):
mi@2: 
mi@2: 		#Generic values are the same for all
mi@2: 		Generic = OutputDescriptor()
mi@2: 		Generic.hasFixedBinCount=False
mi@2: 		Generic.hasKnownExtents=False
mi@2: 		Generic.isQuantized=False
mi@2: 		Generic.sampleType = OneSamplePerStep
mi@2: 		Generic.unit = 'Hz'		
mi@2: 		
mi@2: 		HSPEC = OutputDescriptor(Generic)
mi@2: 		HSPEC.identifier = 'gt'
mi@2: 		HSPEC.name = 'Harmonic emphasised spectrogram'
mi@2: 		HSPEC.description = 'Harmonic emphasised spectrogram'
mi@2: 		HSPEC.binCount = self.nBands
mi@2: 		
mi@2: 		PSPEC = OutputDescriptor(Generic)
mi@2: 		PSPEC.identifier = 'gt'
mi@2: 		PSPEC.name = 'Percussive emphasised spectrogram'
mi@2: 		PSPEC.description = 'Percussive emphasised spectrogram'
mi@2: 		
mi@2: 		return OutputList(HSPEC, PSPEC)
mi@2: 
mi@2: 	def getParameterDescriptors(self):
mi@2: 		
mi@2: 		harmonic_thresh = ParameterDescriptor()
mi@2: 		harmonic_thresh.identifier='percussive_thresh'
mi@2: 		harmonic_thresh.name='percussive_thresh'
mi@2: 		harmonic_thresh.description='Binarisation threshold for harmonic emphasised magnitude spectrogram'
mi@2: 		harmonic_thresh.unit='v'
mi@2: 		harmonic_thresh.minValue=0
mi@2: 		harmonic_thresh.maxValue=1
mi@2: 		harmonic_thresh.defaultValue=0.5
mi@2: 		harmonic_thresh.isQuantized=False
mi@2: 
mi@2: 		percussive_thresh = ParameterDescriptor(harmonic_thresh)
mi@2: 		percussive_thresh.identifier='percussive_thresh'
mi@2: 		percussive_thresh.name='percussive_thresh'
mi@2: 		percussive_thresh.description='Binarisation threshold for percussive emphasised magnitude spectrogram'
mi@2: 		
mi@2: 		harmonic_medwin = ParameterDescriptor()
mi@2: 		harmonic_medwin.identifier = 'harmonic_medwin'
mi@2: 		harmonic_medwin.description = 'Median window length for harmonic part.'
mi@2: 		harmonic_medwin.name = 'Median window length for harmonic part.'
mi@2: 		harmonic_medwin.minValue = 1
mi@2: 		harmonic_medwin.defaultValue = 8
mi@2: 		harmonic_medwin.maxValue = 15
mi@2: 		harmonic_medwin.isQuantized = True
mi@2: 
mi@2: 		percussive_medwin = ParameterDescriptor(harmonic_medwin)
mi@2: 		percussive_medwin.identifier = 'percussive_medwin'
mi@2: 		percussive_medwin.description = 'Median window length for percussive part.'
mi@2: 		percussive_medwin.name = 'Median window length for percussive part.'
mi@2: 		
mi@2: 		boolDescriptor = ParameterDescriptor()
mi@2: 		boolDescriptor.isQuantized = True
mi@2: 		boolDescriptor.minValue= 0
mi@2: 		boolDescriptor.maxValue= 1
mi@2: 		boolDescriptor.quantizeStep = 1
mi@2: 		
mi@2: 		whitening = ParameterDescriptor(boolDescriptor) 
mi@2: 		whitening.identifier='whitening'
mi@2: 		whitening.name='Adaptive whitening'
mi@2: 		whitening.description='Turn adaptive whitening on or off'
mi@2: 		whitening.defaultValue = False
mi@2: 		
mi@2: 		return ParameterList(harmonic_thresh, percussive_thresh, harmonic_medwin, percussive_medwin, whitening)
mi@2: 
mi@2: 	def setParameter(self, paramid, newval):
mi@2: 		if paramid == 'percussive_thresh' :
mi@2: 			self.percussive_thresh = newval
mi@2: 		if paramid == 'harmonic_medwin' :
mi@2: 			self.harmonic_medwin = newval
mi@2: 		if paramid == 'harmonic_medwin' :
mi@2: 			self.harmonic_medwin = newval
mi@2: 		if paramid == 'percussive_medwin' :
mi@2: 			self.percussive_medwin = newval
mi@2: 		if paramid == 'whitening' :
mi@2: 			self.whitening_on = newval == 1.0
mi@2: 			
mi@2: 		return None
mi@2: 		
mi@2: 	def getParameter(self, paramid):
mi@2: 		if paramid == 'percussive_thresh' :
mi@2: 			return self.percussive_thresh
mi@2: 		if paramid == 'harmonic_medwin' :
mi@2: 			return self.harmonic_medwin
mi@2: 		if paramid == 'percussive_medwin' :
mi@2: 			return self.percussive_medwin
mi@2: 		if paramid == 'harmonic_medwin' :
mi@2: 			return self.harmonic_medwin	
mi@2: 		if paramid == 'whitening' :
mi@2: 			if self.whitening_on :
mi@2: 				return 1.0
mi@2: 			else :
mi@2: 				return 0.0
mi@2: 		else:
mi@2: 			return 0.0
mi@2: 
mi@2: 	def whiten(self, magnitudeSpectrogram):
mi@2: 		'''This function reproduces adaptive whitening as described in Dan Stowell's paper.'''
mi@2: 	
mi@2: 		half_length = self.blockSize * 0.5 + 1.0
mi@2: 		nFrames = magnitudeSpectrogram.shape[0]
mi@2: 		whitened_ms = np.zeros_like(magnitudeSpectrogram)
mi@2: 		
mi@2: 		if self.magPeaks is None :
mi@2: 			self.magPeaks = np.zeros(half_length, dtype = float32)
mi@2: 	
mi@2: 		for i in xrange(nFrames):
mi@2: 			m = magnitudeSpectrogram[i, :]	
mi@2: 			idx = m < self.magPeaks
mi@2: 			m[idx] += (self.magPeaks[idx] - m[idx]) * self.whitenRelaxCoeff
mi@2: 			m[m < self.whitenFloor] = self.whitenFloor
mi@2: 			self.magPeaks = m
mi@2: 		
mi@2: 			whitened_ms[i, :] = magnitudeSpectrogram[i, :] / m
mi@2: 		
mi@2: 		return whitened_ms
mi@2: 			
mi@2: 	def process(self, inputbuffers, timestamp):
mi@2: 
mi@2: 		output_featureSet = FeatureSet()
mi@2: 		self.timestamp.append(timestamp)
mi@2: 		complexSpectrum = inputbuffers[0]		
mi@2: 		magnitudeSpectrum = abs(complexSpectrum) 
mi@2: 		self.maglist.append(magnitudeSpectrum)
mi@2: 
mi@2: 		return output_featureSet
mi@2: 	
mi@2: 	def getRemainingFeatures(self):
mi@2: 		
mi@2: 		output_featureSet = FeatureSet() 
mi@2: 		
mi@2: 		nFrames = len(self.maglist)
mi@2: 		magnitudeSpectrogram = np.array(self.maglist)
mi@2: 		
mi@2: 		if self.whitening_on:
mi@2: 			magnitudeSpectrogram = self.whiten(magnitudeSpectrogram)
mi@2: 		else:
mi@2: 			magnitudeSpectrogram = (magnitudeSpectrogram - np.min(magnitudeSpectrogram)) / (np.max(magnitudeSpectrogram) - np.min(magnitudeSpectrogram))
mi@2: 		
mi@2: 		harmonic_ma = median_filter(magnitudeSpectrogram, size=(self.harmonic_medwin, 1))
mi@2: 		percussive_ma = median_filter(magnitudeSpectrogram, size=(1, self.percussive_medwin))
mi@2: 		
mi@2: 		harmonicSpectrogram = magnitudeSpectrogram * (harmonic_ma>=percussive_ma).astype(float)
mi@2: 		percussiveSpectrogram = magnitudeSpectrogram * (harmonic_ma<percussive_ma).astype(float)
mi@2: 
mi@2: 		output_featureSet[0] = flist0 = FeatureList()
mi@2: 		output_featureSet[1] = flist1 = FeatureList()
mi@2: 		for frame in xrange(nFrames):
mi@2: 			f = Feature(timestamp = self.timestamp[frame], values = ma_harmonicSpectrogram[frame, :])
mi@2: 			flist0.append(f)
mi@2: 			f = Feature(timestamp = self.timestamp[frame], values = ma_percussiveSpectrogram[frame, :])
mi@2: 			flist1.append(f)
mi@2: 		
mi@2: 		return output_featureSet