CV、聚类实现图片双色翻转(python)

有一次需要翻转一个图片的颜色，不想用PS，顺手用Python搓了个代码实现双色翻转。

使用的时候，需要图片是二元色调。当然多色图片也能翻转成两种主要颜色或者黑白色调。效果如下：

 

很简单的原理：

1. 把图片读取为RGB三元色矩阵，数据格式为[None,None,3]。
2. 把三元色数组，提出来当三维向量，调用Kmeans聚类算法，找到两个质心
3. 按几何距离重新给三元色矩阵赋值
4. 将三元色矩阵输出为新图片

当然还有许多其他方法。可以设置使用CV库的轮廓计算，根据轮廓直接填充。可以使用神经网络算法（这个稍显复杂）。这上面的算法也可优化，比如不使用几何距离评判重新定义一个函数来使过渡更自然。

代码如下：

import cv2
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
img_path="target.jpg"#原图片路径
img = cv2.imread(img_path)#读取图片
width,height=img.shape[0],img.shape[1]
img_li=img.reshape(img.shape[0]*img.shape[1],3)
estimator = KMeans(n_clusters=2, max_iter=4000, init='k-means++', n_init=50)#
estimator.fit(img_li) 
cluster_centers=estimator.cluster_centers_

new_img_li=img_li.copy()
def get_distance(li1,li2):#计算rgb几何距离
    out=0
    for i in range(len(li1)):
        out+=abs((li1[i]-li2[i])**2)
    out=out**(1/2)
    return out
def turn_to_reverse(li):#计算翻转后的颜色矩阵
    distance1=get_distance(li,cluster_centers[0])
    distance2=get_distance(li,cluster_centers[1])
    if distance1 >= distance2:
        #return [255,255,255] #该设置可生成黑白化图
        return cluster_centers[0]
    else:
        #return [0,0,0] #该设置可生成黑白化图需要与上方同时启用
        return cluster_centers[1]
for i in range(len(new_img_li)):
    new_img_li[i]=turn_to_reverse(new_img_li[i])
new_img_matrix=new_img_li.reshape(width,height,3)
plt.imshow(new_img_matrix)
cv2.imwrite("out.jpg", new_img_matrix)#输出图片