SVM图像分类模型构建

SVM图像分类模型构建

1. 将不同类型的图片分别存储在类型同名文件夹下。

2. 生成训练和测试数据集。

   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
   

3. 初始化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
   
   

4. 模型训练和测试.将训练好的模型存文.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")