import numpy as np
import cv2
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pickle

# 加载训练图像和对应的分区图像
image = cv2.imread("Sandstone_1.tif", 0)
segmentation = cv2.imread("Sandstone_1_segment.tif", 0)

# 输出图像形状
print("图像形状:", image.shape)

# 将图像和分区数据展平
X = image.reshape(-1, 1)
y = segmentation.reshape(-1)
print("完成从砂岩截面图1及其对应分区中获取X和y")
# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
print("完成train_test_split")
# 创建随机森林分类器模型
clf = RandomForestClassifier()

# 训练模型
clf.fit(X_train, y_train)
print("完成随机森林模型clf的训练")
# 在测试集上进行预测
accuracy = clf.score(X_test, y_test)
print("准确率:", accuracy)
print("准确率:", accuracy)
# 保存模型到硬盘
with open("clf.pkl", "wb") as file:
    pickle.dump(clf, file)
print("已保存随机森林模型clf到硬盘")