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_())