-
Notifications
You must be signed in to change notification settings - Fork 0
/
trainSVM.py
32 lines (25 loc) · 855 Bytes
/
trainSVM.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import joblib
from sklearn import datasets
from skimage.feature import hog
from sklearn.svm import LinearSVC
import numpy as np
from collections import Counter
from sklearn.datasets import fetch_openml
# Load the dataset
dataset = fetch_openml('mnist_784')
# Extract the features and labels
features = np.array(dataset.data, 'int16')
labels = np.array(dataset.target, 'int')
# Extract the hog features
list_hog_fd = []
for feature in features:
fd = hog(feature.reshape((28, 28)), orientations=9, pixels_per_cell=(14, 14), cells_per_block=(1, 1))
list_hog_fd.append(fd)
hog_features = np.array(list_hog_fd, 'float64')
print("Count of digits in dataset", Counter(labels))
# Create an linear SVM object
clf = LinearSVC()
# Perform the training
clf.fit(hog_features, labels)
# Save the classifier
joblib.dump(clf, "digits_cls.pkl", compress=3)