Example of EM Maximization using Gaussian Mixture Model for Unsupervised Segmentation of Image

GaussianMixtureModel.py

import cv2
import numpy as np
from sklearn.mixture import GaussianMixture

def preprocess(x): 
   return (x - x.mean(axis=(0,1,2), keepdims=True)) / x.std(axis=(0,1,2), keepdims=True)

# EM hyper parameters
epsilon = 1e-4 # stopping criterion
R = 10 # number of re-runs
N = 2 # number of components
max_iter = 300 # stopping criterion max iterations

# read in the example image
img = cv2.imread('./data/images/test/253027.jpg')
orig = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = preprocess(np.copy(orig))
x = img.reshape(img.shape[0]*img.shape[1],-1)

# use EM to fit N Gaussians
MoG = GaussianMixture(n_init=R,init_params='random',max_iter=max_iter,n_components=N,tol=epsilon,verbose=1)
MoG.fit(x)

# create segmentation map
clustering = MoG.predict(x).reshape(img.shape[0],img.shape[1])

# show results
plt.imshow(clustering)
plt.show()