# 引入必要的库
import cv2
import pickle
import numpy as np
import matplotlib.pyplot as plt

# 加载保存的模型
with open('clf_model.pkl', 'rb') as f:
    clf = pickle.load(f)

# 读取砂岩截面图2及其对应分区图  
original_image = cv2.imread('Sandstone_2.tif', 0)  
segment = cv2.imread('Sandstone_2_segment.tif', 0)  
  
# 获取图像数据和标签  
X = original_image.reshape(-1,1)  

# 对新图像进行预测并显示结果  
y_pred = clf.predict(X)
segmentation = y_pred.reshape(original_image.shape[:2])

# 计算准确率
accuracy = np.mean(segmentation == segment)

# 使用matplotlib显示原分割图和转换后的图像  
fig, ax = plt.subplots(1, 3, figsize=(16, 6))    
ax[0].imshow(original_image,cmap='gray')    
ax[0].set_title("original image", fontsize=15)    
ax[1].imshow(segment,cmap='gray')    
ax[1].set_title("segment", fontsize=15)    
ax[2].imshow(segmentation,cmap='gray')    
ax[2].set_title("segmentation(acc:{:.3f})".format(accuracy), fontsize=15)    
plt.tight_layout()    
plt.show()