Expectation Maximization Prediction Issues

edit

I was trying to segment a leaf from the background to be able to identify the type of disease that the leaf contains. First, I begin to work with a small dataset taken in a controlled environment, after I found out a new dataset which contains more images, but is not so controlled.

Controlled Environment

New Dataset Img Example

However, the code works very well for the first one, it is not good for the second as you can see below:

Results

def train_em(samples, n_clusters):
    print('[Start training EM]')

    em = cv2.ml.EM_create()
    em.setClustersNumber(n_clusters)
    em.setCovarianceMatrixType(cv2.ml.EM_COV_MAT_DIAGONAL)
    em.trainEM(samples)

    print('[Done training EM]')
    return em

def predict(em, image):
    pixelDict = {}
    mask = []
    height, width, _ = image.shape
    for h in range(0, height):
        for w in range(0, width):
            key = hash((image[h,w,0], image[h,w,1]))
            if key in pixelDict:
                mask.append(pixelDict[key])
            else:
                _, probs = em.predict(np.float32([[image[h,w,0], image[h,w,1]]]))
                if probs[0][0] >= probs[0][1]:
                    pixelDict[key] = 1
                else:
                    pixelDict[key] = 0
                mask.append(pixelDict[key])
    return np.array(mask).reshape(height, width).astype('uint8')

Could someone give me a direction?