# 环境要求:Python 3.8+, scikit-learn, OpenCV, TensorFlow/Keras, numpy, pandas, matplotlib
# 安装命令:pip install scikit-learn opencv-python tensorflow numpy pandas matplotlib scikit-image
import application
import os
import numpy as np
import cv2
import tensorflow as tf
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.pipeline import make_pipeline
import matplotlib.pyplot as plt
import seaborn as sns
from tensorflow.keras.applications import VGG16
from tensorflow.keras.models import Model
from skimage.feature import graycomatrix, graycoprops
import joblib
import matplotlib as mpl
# 设置中文字体支持(解决中文显示问题)
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei']
plt.rcParams['axes.unicode_minus'] = False
# ======================
# 1. 数据准备与预处理
# ======================
def load_and_preprocess_data(data_dir, styles, img_size=(224, 224), max_per_class=400):
"""
加载并预处理图像数据
:param data_dir: 数据集目录
:param styles: 艺术风格列表
:param img_size: 目标图像尺寸
:param max_per_class: 每个风格最多加载的图像数
:return: 图像数据和标签
"""
images = []
labels = []
print("开始加载和预处理图像...")
for style_idx, style in enumerate(styles):
style_dir = os.path.join(data_dir, style)
if not os.path.exists(style_dir):
print(f"警告: {style_dir} 目录不存在,跳过")
continue
print(f"处理风格: {style}")
count = 0
for img_name in os.listdir(style_dir):
if max_per_class and count >= max_per_class:
break
img_path = os.path.join(style_dir, img_name)
try:
# 读取并调整大小
img = cv2.imread(img_path)
if img is None:
continue
# 调整尺寸和颜色空间转换
img = cv2.resize(img, img_size)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
images.append(img)
labels.append(style_idx)
count += 1
except Exception as e:
print(f"处理图像 {img_path} 时出错: {str(e)}")
print(f"已加载 {count} 张 {style} 风格的图像")
return np.array(images), np.array(labels)
# 设置数据集路径和风格类别
DATA_DIR = application.MY_DATA_DIR # 需替换为实际数据集路径
ART_STYLES = application.MY_ART_STYLES # 示例风格
# 加载数据(每类最多加载1000张图像)
images, labels = load_and_preprocess_data(DATA_DIR, ART_STYLES, max_per_class=400)
if len(images) == 0:
print("错误: 未加载到任何图像数据,请检查路径和文件")
exit()
print(f"成功加载 {len(images)} 张图像,{len(ART_STYLES)} 种风格")
# 划分训练集和测试集 (70%训练, 30%测试)
X_train, X_test, y_train, y_test = train_test_split(
images, labels, test_size=0.3, random_state=42, stratify=labels
)
# ======================
# 2. 特征工程(添加内存优化和错误处理)
# ======================
def extract_traditional_features(images):
"""
提取传统图像特征(优化内存使用并添加错误处理)
:param images: 图像数据
:return: 特征矩阵
"""
traditional_features = []
print("提取传统图像特征...")
for i, img in enumerate(images):
if i % 50 == 0: # 每处理50张图像输出一次进度
print(f"处理进度: {i + 1}/{len(images)}")
features = []
try:
# 1. 颜色直方图 (RGB)
for j in range(3):
hist = cv2.calcHist([img], [j], None, [256], [0, 256])
features.extend(hist.flatten())
# 2. 纹理特征 (灰度共生矩阵)
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
glcm = graycomatrix(gray, distances=[1], angles=[0], levels=256, symmetric=True, normed=True)
# 计算纹理属性
contrast = graycoprops(glcm, 'contrast')[0, 0]
dissimilarity = graycoprops(glcm, 'dissimilarity')[0, 0]
homogeneity = graycoprops(glcm, 'homogeneity')[0, 0]
energy = graycoprops(glcm, 'energy')[0, 0]
correlation = graycoprops(glcm, 'correlation')[0, 0]
features.extend([contrast, dissimilarity, homogeneity, energy, correlation])
# 3. Hu矩形状特征
_, thresh = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY)
moments = cv2.moments(thresh)
hu_moments = cv2.HuMoments(moments).flatten()
features.extend(hu_moments)
except Exception as e:
print(f"处理图像时出错: {str(e)},使用默认特征值")
# 提供默认特征值
features.extend([0.0] * (256 * 3 + 5 + 7))
traditional_features.append(features)
return np.array(traditional_features)
def extract_deep_features(images):
"""
使用预训练VGG16提取深度特征(添加内存优化)
:param images: 图像数据
:return: 深度特征矩阵
"""
print("提取深度特征...")
# 使用本地权重文件(确保您已下载)
weights_path = 'vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'
# 加载预训练VGG16模型
base_model = VGG16(weights=weights_path, include_top=False,
input_shape=(224, 224, 3))
# 创建特征提取模型
model = Model(inputs=base_model.input,
outputs=base_model.get_layer('block5_pool').output)
# 预处理图像 (VGG16要求)
preprocessed_imgs = tf.keras.applications.vgg16.preprocess_input(
images.astype('float32'))
# 分批提取特征以节省内存
batch_size = 8 # 减小batch_size以适应内存限制
deep_features = []
for i in range(0, len(preprocessed_imgs), batch_size):
batch = preprocessed_imgs[i:i + batch_size]
batch_features = model.predict(batch, verbose=0)
batch_features = batch_features.reshape(batch_features.shape[0], -1)
deep_features.append(batch_features)
if i % (4 * batch_size) == 0: # 每处理4个批次输出一次进度
print(f"深度特征提取进度: {min(i + batch_size, len(preprocessed_imgs))}/{len(preprocessed_imgs)}")
deep_features = np.vstack(deep_features)
return deep_features
print("开始提取传统特征...")
train_trad_features = extract_traditional_features(X_train)
test_trad_features = extract_traditional_features(X_test)
print("开始提取深度特征...")
train_deep_features = extract_deep_features(X_train)
test_deep_features = extract_deep_features(X_test)
# 合并特征
print("合并传统特征和深度特征...")
X_train_features = np.hstack([train_trad_features, train_deep_features])
X_test_features = np.hstack([test_trad_features, test_deep_features])
# 释放不再需要的大对象以节省内存
del train_trad_features, test_trad_features, train_deep_features, test_deep_features
# 特征标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_features)
X_test_scaled = scaler.transform(X_test_features)
# ======================
# 3. 模型构建与训练 (优先使用内存效率高的模型)
# ======================
print("构建和训练分类模型...")
# 优先使用随机森林,内存消耗更小
rf_model = make_pipeline(RandomForestClassifier(n_estimators=100, random_state=42))
rf_model.fit(X_train_scaled, y_train)
# # 在内存允许的情况下训练SVM
# try:
# svm_model = make_pipeline(SVC(probability=True, random_state=42))
# svm_model.fit(X_train_scaled, y_train)
# use_svm = True
# except MemoryError:
# print("内存不足,跳过SVM训练")
# use_svm = False
# ======================
# 4. 模型评估
# ======================
def evaluate_model(model, X_test, y_test, model_name):
"""
评估模型性能并可视化结果
"""
print(f"\n{model_name} 性能评估:")
y_pred = model.predict(X_test)
# 打印分类报告
print(classification_report(y_test, y_pred, target_names=ART_STYLES))
# 绘制混淆矩阵
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
xticklabels=ART_STYLES, yticklabels=ART_STYLES)
plt.title(f'{model_name} 混淆矩阵')
plt.xlabel('预测标签')
plt.ylabel('真实标签')
plt.savefig(f'{model_name}_confusion_matrix.png') # 保存图像
plt.show()
# 评估随机森林
evaluate_model(rf_model, X_test_scaled, y_test, "随机森林(Random Forest)")
# # 评估SVM(如果可用)
# if use_svm:
# evaluate_model(svm_model, X_test_scaled, y_test, "支持向量机(SVM)")
# ======================
# 5. 模型优化 (优先优化随机森林)
# ======================
print("\n开始随机森林模型优化...")
# 定义参数网格
param_grid = {
'randomforestclassifier__n_estimators': [50, 100, 150],
'randomforestclassifier__max_depth': [None, 10, 20],
'randomforestclassifier__min_samples_split': [2, 5, 10]
}
# 网格搜索
grid_search = GridSearchCV(
rf_model,
param_grid,
cv=3,
scoring='accuracy',
n_jobs=1, # 减少并行以避免内存问题
verbose=1
)
grid_search.fit(X_train_scaled, y_train)
# 输出最佳参数
print(f"最佳参数: {grid_search.best_params_}")
print(f"最佳准确率: {grid_search.best_score_:.4f}")
# 评估优化后模型
best_model = grid_search.best_estimator_
evaluate_model(best_model, X_test_scaled, y_test, "优化后随机森林")
# ======================
# 6. 可视化预测结果
# ======================
def visualize_predictions(model, X_test, y_test, num_samples=5):
"""
可视化模型预测结果
"""
plt.figure(figsize=(15, 10))
indices = np.random.choice(range(len(X_test)), num_samples)
for i, idx in enumerate(indices):
# 获取图像和预测
image = X_test[idx]
true_label = ART_STYLES[y_test[idx]]
features = X_test_features[idx].reshape(1, -1)
features_scaled = scaler.transform(features)
# 预测概率
probs = model.predict_proba(features_scaled)[0]
pred_label = ART_STYLES[np.argmax(probs)]
# 绘制图像和预测结果
plt.subplot(2, num_samples, i + 1)
plt.imshow(image)
plt.title(f"真实: {true_label}\n预测: {pred_label}")
plt.axis('off')
# 绘制概率分布
plt.subplot(2, num_samples, i + num_samples + 1)
plt.barh(ART_STYLES, probs)
plt.xlabel('概率')
plt.xlim(0, 1)
plt.tight_layout()
plt.savefig('prediction_visualization.png') # 保存图像
plt.show()
# 可视化优化后模型的预测
visualize_predictions(best_model, X_test, y_test, num_samples=5)
# ======================
# 7. 保存模型
# ======================
# 设置保存目录
save_dir = application.My_save_dir
os.makedirs(save_dir, exist_ok=True)
# 保存最佳模型
model_path = os.path.join(save_dir, 'art_style_classifier.pkl')
scaler_path = os.path.join(save_dir, 'feature_scaler.pkl')
joblib.dump(best_model, model_path)
joblib.dump(scaler, scaler_path)
print(f"模型已保存为 {model_path}")
print(f"特征标准化器已保存为 {scaler_path}")
# ======================
# 8. 模型调用示例
# ======================
def predict_new_image(image_path, model, scaler):
"""
使用训练好的模型预测新图像的风格
"""
# 加载特征提取模型(与训练时相同)
weights_path = 'vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'
base_model = VGG16(weights=weights_path, include_top=False,
input_shape=(224, 224, 3))
feature_extractor = Model(inputs=base_model.input,
outputs=base_model.get_layer('block5_pool').output)
# 加载并预处理新图像
img = cv2.imread(image_path)
if img is None:
raise FileNotFoundError(f"无法加载图像: {image_path}")
img = cv2.resize(img, (224, 224))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_array = np.expand_dims(img, axis=0) # 添加批次维度
# 提取传统特征
trad_features = extract_traditional_features([img])
# 提取深度特征
preprocessed_img = tf.keras.applications.vgg16.preprocess_input(
img_array.astype('float32'))
deep_features = feature_extractor.predict(preprocessed_img)
deep_features = deep_features.reshape(deep_features.shape[0], -1)
# 合并特征
features = np.hstack([trad_features, deep_features])
# 标准化特征
features_scaled = scaler.transform(features)
# 预测
probs = model.predict_proba(features_scaled)[0]
pred_idx = np.argmax(probs)
pred_style = ART_STYLES[pred_idx]
# 创建概率字典
prob_dict = {style: prob for style, prob in zip(ART_STYLES, probs)}
return pred_style, prob_dict, img
print("程序执行完成!")优化一个这个代码,使它的准确率提高,只使用传统机器学习算法