機器學習-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.