SVM图像分类模型构建
-
将不同类型的图片分别存储在类型同名文件夹下。
-
生成训练和测试数据集。
import pandas as pd import os from skimage.transform import resize from skimage.io import imread from sklearn.model_selection import train_test_split Categories = ["zx", "lz", "rz"] # 图像类别列表 # datadir文件夹为包括所有图片文件夹父文件夹 def generateData(datadir): flat_data_arr = [] # input array target_arr = [] # output array for i in Categories: print(f'loading... category : {i}') path = os.path.join(datadir, i) for img in os.listdir(path): img_array = imread(os.path.join(path, img)) img_resized = resize(img_array, (150, 150, 3)) flat_data_arr.append(img_resized.flatten()) target_arr.append(Categories.index(i)) print(f'loaded category:{i} successfully') flat_data = np.array(flat_data_arr) target = np.array(target_arr) df = pd.DataFrame(flat_data) # dataframe df['Target'] = target x = df.iloc[:, :-1] # input data y = df.iloc[:, -1] # output data return train_test_split(x, y, test_size=0.20, random_state=77, stratify=y) pass
-
初始化SVM模型
from sklearn.model_selection import GridSearchCV from sklearn import svm def generateSVC(): param_grid = {'C': [0.1, 1, 10, 100], 'gamma': [0.0001, 0.001, 0.1, 1], 'kernel': ['rbf', 'poly']} svc = svm.SVC(probability=True) return GridSearchCV(svc, param_grid) pass
-
模型训练和测试.将训练好的模型存文.sav格式文件,输出测试精度。
from sklearn.metrics import accuracy_score import pickle def train_test(datadir, filename='svm/svm1.sav'): model = generateSVC() x_train, x_test, y_train, y_test = generateData(datadir) print('Splitted Successfully') model.fit(x_train, y_train) print('The Model is trained well with the given images') # save svm model pickle.dump(model, open(filename, 'wb')) # load the model from disk and test model = pickle.load(open(filename, 'rb')) y_pred = model.predict(x_test) print("The predicted Data is :") print(y_pred) print("The actual data is:") print(np.array(y_test)) print(f"The model is {accuracy_score(y_pred, y_test) * 100}% accurate")