import numpy as np
from keras import Model
from keras.applications import VGG16
from keras.layers import GlobalAveragePooling2D
from keras.utils import to_categorical
from lazypredict.Supervised import LazyClassifier
from util import save_file, get_config_value, load_file

# 加载训练集和测试集
X_train, X_test, y_train, y_test = load_file("X_train, X_test, y_train, y_test", f'{get_config_value("Xy_root")}/Xy')

# 将灰度图像转换为RGB图像并缩放到0-1之间
X_train_scaled = np.stack((X_train,) * 3, axis=-1) / 255.0
X_test_scaled = np.stack((X_test,) * 3, axis=-1) / 255.0

# 调整输入形状
X_train_scaled = X_train_scaled.reshape(-1, 100, 100, 3)
X_test_scaled = X_test_scaled.reshape(-1, 100, 100, 3)

# 对标签进行独热编码
num_classes = np.max(y_train) + 1
y_train_encoded = to_categorical(y_train, num_classes)
y_test_encoded = to_categorical(y_test, num_classes)

# 加载预训练的VGG16模型并添加全局平均池化层
base_model = VGG16(weights='imagenet', include_top=False, input_shape=(100, 100, 3))
x = base_model.output
x = GlobalAveragePooling2D()(x)

# 创建新的模型
model = Model(inputs=base_model.input, outputs=x)

# 提取训练集和测试集的特征
X_train_features = model.predict(X_train_scaled)
X_test_features = model.predict(X_test_scaled)

# 使用LazyClassifier进行模型指标对比
print("开始评估所有的模型:")
clf = LazyClassifier()

# 获取每个模型的预测对比结果
scores, _ = clf.fit(X_train_features, X_test_features, y_train, y_test)
print(scores)

# 获取F1分数最高的模型
best_model_name = scores['F1 Score'].idxmax()
print("F1分数最高的模型是: ", best_model_name)

# 从模型字典中获取最佳模型对象
best_model = clf.models[best_model_name]

# 序列化最佳模型
save_file(best_model, "最好的F1分数的模型", f'{get_config_value("model_root")}/best_model')