機器學習-EM演算法(Expectation maximization algorithm)(四)Experiment-Image Segmentation Using Gaussian Mixture Model with PyTorch.

First, use the K-means algorithm to find K central pixels. Second, use Expectation maximization (EM) algorithm to optimize the parameters of the model.

Equipment:
cpu: i7-9750h
gpu: GTX 1660 Ti

原圖original image

K-mean code

This code will first compute k (user defined) mean values of the input image and output the same image except the pixel values are scale to the nearest mean.

GMM code

Using the means from k-means as initial value for GMM, after perform EM algorithm, output the image scale to the most probable mean value.

Run K-means and GMM for K = 5, 7, 10, and 15, then show the resulting images.

It can be easily seen from the above resulting images that when K increases, it will be closer to the original image.