mitian@16: #!/usr/bin/env python
mitian@16: # encoding: utf-8
mitian@16: """
mitian@16: snmf.py
mitian@16: -- A modified version of C-NMF by Nieto. Input is no longer feature matrix but SSM.
mitian@16: 
mitian@16: C-NMF method for segmentation, modified from here:
mitian@16: 
mitian@16: Nieto, O., Jehan, T., Convex Non-negative Matrix Factorization For Automatic
mitian@16: Music Structure Identification. Proc. of the 38th IEEE International Conference
mitian@16: on Acoustics, Speech, and Signal Processing (ICASSP). Vancouver, Canada, 2013.
mitian@16: 
mitian@16: """
mitian@16: 
mitian@16: import sys
mitian@16: import os
mitian@16: import pymf
mitian@16: from utils import SegUtil
mitian@16: from utils.SegUtil import median_filter
mitian@16: import numpy as np
mitian@16: 
mitian@16: # Algorithm params
mitian@16: h = 8				# Size of median filter for SSMs in C-NMF
mitian@16: R = 16				# Size of the median filter for the activation matrix C-NMF
mitian@16: rank = 3			# Rank of decomposition for the boundaries
mitian@16: rank_labels = 16	 # Rank of decomposition for the labels
mitian@16: R_labels = 4		# Size of the median filter for the labels
mitian@16: 
mitian@16: 
mitian@16: def cnmf(S, rank, niter=500):
mitian@16: 	"""(Convex) Non-Negative Matrix Factorization.
mitian@16: 
mitian@16: 	Parameters
mitian@16: 	----------
mitian@16: 	S: np.array(N, N)
mitian@16: 	   SSM.
mitian@16: 	rank: int
mitian@16: 		Rank of decomposition
mitian@16: 	niter: int
mitian@16: 		Number of iterations to be used
mitian@16: 
mitian@16: 	Returns
mitian@16: 	-------
mitian@16: 	F: np.array
mitian@16: 		Cluster matrix (decomposed matrix)
mitian@16: 	G: np.array
mitian@16: 		Activation matrix (decomposed matrix)
mitian@16: 		(s.t. S ~= F * G)
mitian@16: 	"""
mitian@16: 	nmf_mdl = pymf.CNMF(S, num_bases=rank)
mitian@16: 	nmf_mdl.factorize(niter=niter)
mitian@16: 	F = np.asarray(nmf_mdl.W)
mitian@16: 	G = np.asarray(nmf_mdl.H)
mitian@16: 	return F, G
mitian@16: 
mitian@16: def nmf(S, rank, nither=500):
mitian@16: 	nmf_mdl = pymf.NMF(S, num_bases=rank, niter=nither)
mitian@16: 	nmf_mdl.factorize()
mitian@16: 	F = np.asarray(nmf_mdl.W)
mitian@16: 	G = np.asarray(nmf_mdl.H)
mitian@16: 	return F, G
mitian@16: 	
mitian@16: 	
mitian@16: def most_frequent(x):
mitian@16: 	"""Returns the most frequent value in x."""
mitian@16: 	return np.argmax(np.bincount(x))
mitian@16: 
mitian@16: 
mitian@16: def compute_labels(X, rank, R, bound_idxs, niter=300):
mitian@16: 	"""Computes the labels using the bounds."""
mitian@16: 
mitian@16: 	X = X.T
mitian@16: 	try:
mitian@16: 		F, G = cnmf(X, rank, niter=niter)
mitian@16: 	except:
mitian@16: 		return [1]
mitian@16: 
mitian@16: 	label_frames = filter_activation_matrix(G.T, R)
mitian@16: 	label_frames = np.asarray(label_frames, dtype=int)
mitian@16: 
mitian@16: 	# Get labels from the label frames
mitian@16: 	labels = []
mitian@16: 	bound_inters = zip(bound_idxs[:-1], bound_idxs[1:])
mitian@16: 	for bound_inter in bound_inters:
mitian@16: 		if bound_inter[1] - bound_inter[0] <= 0:
mitian@16: 			labels.append(np.max(label_frames) + 1)
mitian@16: 		else:
mitian@16: 			labels.append(most_frequent(
mitian@16: 				label_frames[bound_inter[0]:bound_inter[1]]))
mitian@16: 
mitian@16: 	return labels
mitian@16: 
mitian@16: 
mitian@16: def filter_activation_matrix(G, R):
mitian@16: 	"""Filters the activation matrix G, and returns a flattened copy."""
mitian@16: 	idx = np.argmax(G, axis=1)
mitian@16: 	max_idx = np.arange(G.shape[0])
mitian@16: 	max_idx = (max_idx, idx.flatten())
mitian@16: 	G[:, :] = 0
mitian@16: 	G[max_idx] = idx + 1
mitian@16: 	G = np.sum(G, axis=1)
mitian@16: 	G = SegUtil.median_filter(G[:, np.newaxis], R)
mitian@16: 	return G.flatten()
mitian@16: 
mitian@16: 
mitian@16: def segmentation(X, rank=4, R=15, h=8, niter=300, CNMF=True):
mitian@16: 	"""
mitian@16: 	Gets the segmentation (boundaries and labels) from the factorization
mitian@16: 	matrices.
mitian@16: 
mitian@16: 	Parameters
mitian@16: 	----------
mitian@16: 	X: np.array()
mitian@16: 		Features matrix (e.g. chromagram)
mitian@16: 	rank: int
mitian@16: 		Rank of decomposition
mitian@16: 	R: int
mitian@16: 		Size of the median filter for activation matrix
mitian@16: 	niter: int
mitian@16: 		Number of iterations for k-means
mitian@16: 	bound_idxs : list
mitian@16: 		Use previously found boundaries (None to detect them)
mitian@16: 	CNMF : bool
mitian@16: 		If True, use CNMF; otherwise use NMF
mitian@16: 		
mitian@16: 	Returns
mitian@16: 	-------
mitian@16: 	bounds_idx: np.array
mitian@16: 		Bound indeces found
mitian@16: 	labels: np.array
mitian@16: 		Indeces of the labels representing the similarity between segments.
mitian@16: 	"""
mitian@16: 
mitian@16: 	# Filter
mitian@16: 	X = median_filter(X, M=h)
mitian@16: 	X = X.T
mitian@16: 
mitian@16: 	# Find non filtered boundaries
mitian@16: 	bound_idxs = None
mitian@16: 	while True:
mitian@16: 		if bound_idxs is None:
mitian@16: 			try:
mitian@16: 				if CNMF: F, G = cnmf(X, rank, niter=niter)
mitian@16: 				else: F, G = nmf(X, rank, niter=niter)
mitian@16: 			except:
mitian@16: 				return np.empty(0), [1]
mitian@16: 
mitian@16: 			# Filter G
mitian@16: 			G = filter_activation_matrix(G.T, R)
mitian@16: 			if bound_idxs is None:
mitian@16: 				bound_idxs = np.where(np.diff(G) != 0)[0] + 1
mitian@16: 
mitian@16: 		if len(np.unique(bound_idxs)) <= 2:
mitian@16: 			rank += 1
mitian@16: 			bound_idxs = None
mitian@16: 		else:
mitian@16: 			break
mitian@16: 
mitian@16: 	return G, bound_idxs