130 lines
3.5 KiB
Python

import cv2
import numpy as np
import matplotlib.pyplot as plt
from skimage.feature import hog
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from sklearn import svm
from sklearn import tree
import pickle
import os
def lbp(image):
height, width = image.shape[:2]
lbp_image = np.zeros_like(image, dtype=np.uint8)
for y in range(1, height-1):
for x in range(1, width-1):
center = image[y, x]
code = 0
# Compare the pixel values of the neighborhood with the center pixel
code |= (image[y-1, x-1] > center) << 7
code |= (image[y-1, x] > center) << 6
code |= (image[y-1, x+1] > center) << 5
code |= (image[y, x+1] > center) << 4
code |= (image[y+1, x+1] > center) << 3
code |= (image[y+1, x] > center) << 2
code |= (image[y+1, x-1] > center) << 1
code |= (image[y, x-1] > center) << 0
lbp_image[y, x] = code
return lbp_image
def load_data(pictures, marks):
pictures = np.load(pictures)
marks = np.load(marks)
return pictures, marks
def get_data(folder_path):
img_data = []
marks = []
i = 0
for root, dirs, files in os.walk(folder_path):
for file in files:
file_path = os.path.join(root, file)
if 'c' in file_path:
marks.append(1)
else:
marks.append(0)
img = cv2.imread(file_path)
img = cv2.resize(img, (100, 100))
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
lbpHist = cv2.calcHist([lbp(gray)], [0], None, [256], [0, 256])
lbpHist = lbpHist.astype(np.float64)
lbpHist = lbpHist.flatten()
hogFeatures = hog(gray, orientations=9, pixels_per_cell=(10, 10), cells_per_block=(2, 2), block_norm='L2-Hys')
features = np.concatenate((lbpHist, hogFeatures), axis=0)
img_data.append(features)
print(i)
i += 1
img_data = np.array(img_data)
marks = np.array(marks)
return img_data, marks
def save_data(img_data, marks, img_data_path, marks_data_path):
np.save(img_data_path, img_data)
np.save(marks_data_path, marks)
def train_and_save_model(img_data, marks):
#---------------------KNN-------------------
knn = KNeighborsClassifier(n_neighbors=5)
knn.fit(img_data, marks)
knnPickle = open('Models\\knn.bin', 'wb')
pickle.dump(knn, knnPickle)
knnPickle.close()
#------------------SVM----------------------
svmM = svm.SVC()
svmM.fit(img_data, marks)
svmPickle = open('Models\\svm.bin', 'wb')
pickle.dump(svmM, svmPickle)
svmPickle.close()
#------------------DT----------------------
dt = tree.DecisionTreeClassifier()
dt.fit(img_data, marks)
dtPickle = open('Models\\dt.bin', 'wb')
pickle.dump(dt, dtPickle)
dtPickle.close()
def load_model():
#---------------------KNN-------------------
knnPickle = open('Models\\knn.bin', 'rb')
knn = pickle.load(knnPickle)
knnPickle.close()
#------------------SVM----------------------
svmPickle = open('Models\\svm.bin', 'rb')
svm = pickle.load(svmPickle)
svmPickle.close()
#------------------DT----------------------
dtPickle = open('Models\\dt.bin', 'rb')
dt = pickle.load(dtPickle)
dtPickle.close()
return knn, svm, dt