205 lines
8.2 KiB
Python
205 lines
8.2 KiB
Python
import random
|
||
import sys
|
||
import cv2
|
||
import numpy as np
|
||
import random
|
||
import math
|
||
from PyQt5.QtCore import Qt
|
||
from sklearn.ensemble import RandomForestClassifier
|
||
from sklearn.metrics import accuracy_score
|
||
from PyQt5.QtGui import QImage, QPixmap
|
||
from PyQt5.QtWidgets import (QWidget, QApplication, QLabel, QVBoxLayout, QHBoxLayout, QPushButton,
|
||
QFileDialog,QComboBox,QMessageBox,QTextBrowser)
|
||
|
||
|
||
class Example(QWidget):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.image = None
|
||
self.textbrowser = QTextBrowser()
|
||
|
||
self.btn_open = QPushButton('Изображения folder1')
|
||
self.btn_open.clicked.connect(self.openImages1)
|
||
|
||
self.btn_open1 = QPushButton('Изображения folder2')
|
||
self.btn_open1.clicked.connect(self.openImages2)
|
||
|
||
self.btn_ = QPushButton('Обучить и создать матрицы')
|
||
self.btn_.clicked.connect(self.matrix_and_train)
|
||
|
||
self.btn1 = QPushButton('Детектировать тестовые изображения')
|
||
self.btn1.clicked.connect(self.detect_image)
|
||
|
||
self.btn2 = QPushButton('Сформировать лес и натренировать полученную модель')
|
||
self.btn2.clicked.connect(self.form_les_and_train)
|
||
|
||
self.btn_.setVisible(False)
|
||
self.btn1.setVisible(False)
|
||
self.btn2.setVisible(False)
|
||
|
||
top_bar = QHBoxLayout()
|
||
top_bar.addWidget(self.btn_open)
|
||
top_bar.addWidget(self.btn_open1)
|
||
top_bar.addWidget(self.btn_)
|
||
top_bar.addWidget(self.btn1)
|
||
top_bar.addWidget(self.btn2)
|
||
root = QVBoxLayout(self)
|
||
root.addLayout(top_bar)
|
||
root.addWidget(self.textbrowser)
|
||
|
||
self.spisok = list()
|
||
self.spisok2 = list()
|
||
self.spisok3 = list()
|
||
|
||
self.train = list()
|
||
self.matrix = list()
|
||
self.resize(540, 574)
|
||
self.setWindowTitle('ST_4')
|
||
self.show()
|
||
|
||
def openImages1(self):
|
||
filenames1 = QFileDialog.getOpenFileNames(None, 'Открыть изображения', '.',
|
||
'Image Files (*.png *.jpg *.jpeg *.bmp)')
|
||
lk = filenames1[0]
|
||
self.erroropened(lk, "1")
|
||
self.mass(lk, 1)
|
||
lk.clear()
|
||
|
||
def openImages2(self):
|
||
filenames2 = QFileDialog.getOpenFileNames(None, 'Открыть изображения', '.',
|
||
'Image Files (*.png *.jpg *.jpeg *.bmp)')
|
||
lk = filenames2[0]
|
||
self.erroropened(lk, "2")
|
||
self.mass(lk, 2)
|
||
lk.clear()
|
||
|
||
def erroropened(self, s, s1):
|
||
if len(s) != 0:
|
||
q = QMessageBox.information(self, "Информация", "Изображения из " + s1 + " папки получены!")
|
||
else:
|
||
q = QMessageBox.information(self, "Информация", "Вы не выбрали изображения!")
|
||
|
||
def mass(self, p, s1):
|
||
if s1 == 1:
|
||
self.spisok3.clear()
|
||
for v in range(len(p)):
|
||
self.spisok.append(str(p[v]))
|
||
self.spisok3.append(str(p[v]) + "SKT")
|
||
self.appendos("Тестовый набор картинок", self.spisok)
|
||
if s1 == 2:
|
||
for v in range(len(p)):
|
||
self.spisok2.append(str(p[v]))
|
||
self.spisok3.append(str(p[v]) + "NO")
|
||
self.spisok3 = list(set(self.spisok3))
|
||
self.appendos("Набор картинок для тренировки", self.spisok2)
|
||
self.btn_open.setVisible(False)
|
||
self.btn_open1.setVisible(False)
|
||
self.btn_.setVisible(True)
|
||
q = QMessageBox.information(self, "Информация", "Количество изображений тестовой категории: " + str(
|
||
len(self.spisok)) + "\nКоличество изображений основной категории: " + str(
|
||
len(self.spisok2)) + "\nОбщее количество изображений: " + str(
|
||
int(len(self.spisok) + int(len(self.spisok2)))))
|
||
|
||
def matrix_and_train(self):
|
||
for v in range(len(self.spisok3)):
|
||
s = str(self.spisok3[v])
|
||
if s.endswith("SKT"):
|
||
self.train.append(round(float(0.5), 1))
|
||
self.matrix.append(round(float(1.0), 1))
|
||
if s.endswith("NO"):
|
||
q = round(float(random.uniform(1.0, 3.0)), 1)
|
||
self.train.append(q)
|
||
self.matrix.append(round(float(-1.0), 1))
|
||
|
||
self.appendos("Ваши тренировочные данные", self.train)
|
||
self.appendos("Ваша матрица", self.matrix)
|
||
train = np.array([self.train], dtype=int)
|
||
labels = np.array(self.matrix, dtype=int)
|
||
self.svm = cv2.ml.SVM_create()
|
||
self.svm.train(train, cv2.ml.COL_SAMPLE, labels)
|
||
self.svm.save("1.yml")
|
||
self.textbrowser.append("Модель сохранена!")
|
||
close = QMessageBox.question(self, "Поздравляем!", "Ваша модель натренирована!",
|
||
QMessageBox.Yes | QMessageBox.No)
|
||
if close == QMessageBox.Yes:
|
||
pass
|
||
self.btn_.setVisible(False)
|
||
self.btn1.setVisible(True)
|
||
|
||
def appendos(self, s1, s2):
|
||
self.textbrowser.append(s1)
|
||
for v in range(len(s2)):
|
||
self.textbrowser.append(str(s2[v]))
|
||
|
||
def detect_image(self):
|
||
self.model = RandomForestClassifier(n_estimators=100, bootstrap=True, max_features='sqrt')
|
||
self.spisok3.clear()
|
||
for v in range(len(self.spisok)):
|
||
lkst = list()
|
||
img = cv2.imread(str(self.spisok[v]))
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
sift = cv2.SIFT_create()
|
||
kp = sift.detect(gray, None)
|
||
for keyPoint in kp:
|
||
# self.spisok3.append(keyPoint.pt[0])
|
||
# self.spisok3.append(keyPoint.pt[1])
|
||
# print(keyPoint.pt[0])
|
||
# print(keyPoint.pt[1])
|
||
lkst.append(keyPoint.pt[0])
|
||
lkst.append(keyPoint.pt[1])
|
||
|
||
# self.spisok3.append(lkst)
|
||
train = np.array([lkst], dtype=int)
|
||
labels = np.array([lkst], dtype=int)
|
||
self.model.fit(train, labels)
|
||
self.textbrowser.append("Модель натренирована на рисунке " + str(v))
|
||
# print(len(self.spisok3))
|
||
self.textbrowser.append("Выявлены контурные точки на тестовых изображениях!")
|
||
|
||
# s = keyPoint.size
|
||
# print(x,y,s)
|
||
self.btn1.setVisible(False)
|
||
self.btn2.setVisible(True)
|
||
|
||
def form_les_and_train(self):
|
||
|
||
# print(self.spisok3)
|
||
for v in range(len(self.spisok2)):
|
||
lkst = list()
|
||
img = cv2.imread(str(self.spisok2[v]))
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
sift = cv2.SIFT_create()
|
||
kp = sift.detect(gray, None)
|
||
for keyPoint in kp:
|
||
# self.spisok3.append(keyPoint.pt[0])
|
||
# self.spisok3.append(keyPoint.pt[1])
|
||
# print(keyPoint.pt[0])
|
||
# print(keyPoint.pt[1])
|
||
lkst.append(keyPoint.pt[0])
|
||
lkst.append(keyPoint.pt[1])
|
||
train = np.array([lkst], dtype=int)
|
||
self.model.predict(train)
|
||
# self.model.
|
||
print("RED")
|
||
# train = np.array([self.spisok3],dtype=int)
|
||
# labels = np.array([self.spisok3],dtype=int)
|
||
# model = RandomForestClassifier(n_estimators=100,bootstrap = True,max_features = 'sqrt')
|
||
# model.fit(train, labels)
|
||
# tree = DecisionTreeClassifier()
|
||
# tree.fit(self.spisok3, self.spisok3)
|
||
# accuracy = accuracy_score(len(self.spisok3), len(self.spisok3))
|
||
# print('Model Accuracy:',accuracy)
|
||
|
||
def closeEvent(self, event):
|
||
close = QMessageBox.question(self, "Выход", "Вы хотите завершить работу?", QMessageBox.Yes | QMessageBox.No)
|
||
if close == QMessageBox.Yes:
|
||
event.accept()
|
||
else:
|
||
event.ignore()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
app = QApplication(sys.argv)
|
||
win = Example()
|
||
sys.exit(app.exec_())
|