如何使用 python 从 Kmeans 的输出中保存特定的彩色图像

How to save specific color image from the output of Kmeans using python

我正在使用下面的 使用 K-means 进行基于颜色的分割。在此代码中,每个簇都保存到一个图像中。在我的情况下,要求有点不同。我只想保存蓝色图像。你能帮我看看我怎样才能只保存蓝色图像吗?

import numpy as np
import cv2
import pdb
from matplotlib import pyplot as plt
img = cv2.imread('a.png')
Z = np.float32(img.reshape((-1,3)))

criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
K = 4
_,labels,centers = cv2.kmeans(Z, K, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
labels = labels.reshape((img.shape[:-1]))
reduced = np.uint8(centers)[labels]

result = [np.hstack([img, reduced])]
for i, c in enumerate(centers):
    mask = cv2.inRange(labels, i, i)
    mask = np.dstack([mask]*3) # Make it 3 channel
    ex_img = cv2.bitwise_and(img, mask)
    ex_reduced = cv2.bitwise_and(reduced, mask)
    result.append(np.hstack([ex_img, ex_reduced]))
    pdb.set_trace()

cv2.imwrite('watermelon_out.jpg', np.vstack(result))

原图

使用此代码后,我得到以下结果 link:

预期结果:

这应该只打印蓝色图像。首先找到最接近蓝色的中心,然后仅在该中心代表的簇中绘制点

import numpy as np
import cv2
import pdb
from matplotlib import pyplot as plt
img = cv2.imread('a.png')
Z = np.float32(img.reshape((-1,3)))

criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
K = 4
_,labels,centers = cv2.kmeans(Z, K, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
labels = labels.reshape((img.shape[:-1]))
reduced = np.uint8(centers)[labels]


blue_dis = 99999999
blue_center = -1
b = (255, 50 , 0)
for i, c in enumerate(centers):
    dis = (c[0]-b[0])**2 + (c[1]-b[1])**2 + (c[1]-b[1])**2
    if dis < blue_dis:
        blue_center = i
        blue_dis = dis


    

result = [np.hstack([img, reduced])]
for i, c in enumerate(centers):
    if i!=blue_center:
        continue
    mask = cv2.inRange(labels, i, i)
    mask = np.dstack([mask]*3) # Make it 3 channel
    ex_img = cv2.bitwise_and(img, mask)
    ex_reduced = cv2.bitwise_and(reduced, mask)
    result.append(np.hstack([ex_img, ex_reduced]))
    pdb.set_trace()

cv2.imwrite('watermelon_out.jpg', np.vstack(result))