import numpy as np
import cv2
import os
from skimage.feature import hog
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, classification_report
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import joblib
# 1. 数据准备和特征提取
def load_dataset_and_extract_features(data_dir):
"""
加载车型图像数据集并提取HOG特征
"""
features = []
labels = []
categories = ['0', '1', '2', '3'] # 车型分类
for label, category in enumerate(categories):
category_dir = os.path.join(data_dir, category)
if not os.path.exists(category_dir):
print(f"警告：类别目录 {category_dir} 不存在，已跳过")
continue
for img_name in os.listdir(category_dir):
if img_name.endswith('.jpg'):
# 读取图像并转为灰度图
img_path = os.path.join(category_dir, img_name)
try:
img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
if img is None:
print(f"警告：无法读取图像 {img_path}，已跳过")
continue
# 调整图像尺寸
img = cv2.resize(img, (128, 64)) # 标准尺寸用于HOG特征提取
# 提取HOG特征
hog_features = hog(img, orientations=9, pixels_per_cell=(8, 8),
cells_per_block=(2, 2), transform_sqrt=True,
block_norm='L2-Hys')
features.append(hog_features)
labels.append(label)
except Exception as e:
print(f"处理图像 {img_path} 时出错: {e}")
continue
if not features:
raise ValueError("未找到有效图像数据，请检查数据集路径和文件格式")
return np.array(features), np.array(labels)
# 2. 数据预处理
def preprocess_data(features, labels):
"""
数据标准化和训练/测试集划分
"""
# 特征标准化
scaler = StandardScaler()
features_scaled = scaler.fit_transform(features)
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
features_scaled, labels, test_size=0.2, random_state=42
)
return X_train, X_test, y_train, y_test, scaler
# 3. SVM模型训练
def train_svm_classifier(X_train, y_train):
"""
训练SVM分类器
"""
svm = SVC(kernel='rbf', C=1.0, gamma='scale', probability=True)
svm.fit(X_train, y_train)
return svm
# 4. 模型评估
def evaluate_model(model, X_test, y_test):
"""
评估模型性能
"""
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
report = classification_report(y_test, y_pred,
target_names=['0', '1', '2', '3'])
print(f"模型准确率: {accuracy:.2f}")
print("分类报告:\n", report)
# 可视化混淆矩阵
from sklearn.metrics import ConfusionMatrixDisplay
ConfusionMatrixDisplay.from_predictions(y_test, y_pred,
display_labels=['0', '1', '2', '3'])
plt.title('车型分类混淆矩阵')
plt.show()
# 主函数
def main():
try:
# 数据集路径 - 替换为实际路径
DATA_DIR = "C:/Users/Administrator/Desktop/jinwan"
if not os.path.exists(DATA_DIR):
raise FileNotFoundError(f"数据集目录 {DATA_DIR} 不存在")
# 执行流程
print("正在加载数据集并提取特征...")
features, labels = load_dataset_and_extract_features(DATA_DIR)
print("正在进行数据预处理...")
X_train, X_test, y_train, y_test, scaler = preprocess_data(features, labels)
print("正在训练SVM模型...")
model = train_svm_classifier(X_train, y_train)
print("正在评估模型性能...")
evaluate_model(model, X_test, y_test)
# 保存模型和标准化器到指定路径
SAVE_DIR = "C:/Users/Administrator/Desktop/ultralytics-main"
os.makedirs(SAVE_DIR, exist_ok=True) # 自动创建目录
model_path = os.path.join(SAVE_DIR, "vehicle_classifier_svm.pkl")
scaler_path = os.path.join(SAVE_DIR, "standard_scaler.pkl")
joblib.dump(model, model_path)
joblib.dump(scaler, scaler_path)
print(f"模型已保存到: {model_path}")
print(f"标准化器已保存到: {scaler_path}")
except Exception as e:
print(f"程序运行出错: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()为什么我训练完之后模型保存不了
本文解决了一个关于从sklearn.externals导入joblib时出现的错误。通过修改导入方式，即将from sklearn.externals import joblib更改为import joblib，成功解决了问题。
扫一扫
1601
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
