import sys import sqlite3 from PyQt5.QtSql import QSqlDatabase, QSqlQuery from PyQt5.QtCore import QDateTime from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QMessageBox from PyQt5.QtGui import QPixmap, QImage import numpy as np import cv2 import os import time import tensorflow as tf from object_detection.utils import label_map_util from object_detection.utils import visualization_utils as vis_utils from PyQt5.QtCore import QDateTime, QTimer #忽略 import os os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' class MainWindow(QMainWindow): def __init__(self): super(MainWindow, self).__init__() # 设置窗口大小 self.setGeometry(100, 100, 800, 700) # 设置窗口名称 self.setWindowTitle("object_detection") # 创建一个标签用于显示图像 self.label = QLabel(self) self.label.setGeometry(10, 10, 640, 480) # 创建一个按钮用于打开摄像头 self.open_camera_button = QPushButton('Open Camera', self) self.open_camera_button.setGeometry(10, 550, 200, 60) self.open_camera_button.clicked.connect(self.on_open_camera_button_clicked) # 创建一个按钮用于保存识别结果到数据库 self.save_to_db_button = QPushButton('Save to Database', self) self.save_to_db_button.setGeometry(550, 500, 200, 30) self.save_to_db_button.clicked.connect(self.on_save_to_db_button_clicked) # 创建一个按钮用于保存识别结果到数据库 self.save_to_db_button = QPushButton('Save to Database', self) self.save_to_db_button.setGeometry(550, 500, 200, 30) self.save_to_db_button.clicked.connect(self.on_save_to_db_button_clicked) # 创建一个按钮用于保存图像 self.save_image_button = QPushButton('Save Image', self) self.save_image_button.setGeometry(280, 550, 200, 60) self.save_image_button.clicked.connect(self.on_save_image_button_clicked) # 创建一个按钮用于关闭应用程序 self.close_button = QPushButton('Close', self) self.close_button.setGeometry(550, 550, 200, 60) self.close_button.clicked.connect(self.close) # 创建一个标签用于显示日期和时间 self.datetime_label = QLabel(self) self.datetime_label.setGeometry(10, 500, 300, 30) # 更新日期和时间 self.update_datetime() # 创建一个定时器,每秒钟更新一次日期和时间 self.timer = QTimer(self) self.timer.timeout.connect(self.update_datetime) self.timer.start(1000) # 初始化摄像头 self.cap = cv2.VideoCapture(0) # 加载模型 self.MODEL_NAME = 'ssdlite_mobilenet_v2_coco_2018_05_09' self.PATH_TO_CKPT = "ssdlite_mobilenet_v2_coco_2018_05_09/frozen_inference_graph.pb" self.PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt') self.NUM_CLASSES = 90 self.detection_graph = tf.Graph() with self.detection_graph.as_default(): od_graph_def = tf.compat.v1.GraphDef() with open(self.PATH_TO_CKPT, 'rb') as fid: serialized_graph = fid.read() od_graph_def.ParseFromString(serialized_graph) tf.import_graph_def(od_graph_def, name='') label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS) categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=self.NUM_CLASSES, use_display_name=True) self.category_index = label_map_util.create_category_index(categories) # 开始计时 self.t_start = time.time() self.fps = 0 def on_save_to_db_button_clicked(self): # 获取识别结果 class_names = self.class_names scores = self.scores # 获取当前时间 # 连接到数据库 db = QSqlDatabase.addDatabase('QSQLITE') db.setDatabaseName('object_detection.db') if not db.open(): QMessageBox.critical(self, "Error", "Failed to connect to database.") return # 创建表格 query = QSqlQuery() query.exec_( "CREATE TABLE IF NOT EXISTS detection_results (id INTEGER PRIMARY KEY AUTOINCREMENT, class_name TEXT, score REAL)") # 将识别结果插入数据库 for i in range(len(class_names)): query.prepare("INSERT INTO detection_results (class_name, score) VALUES (?, ?)") query.bindValue(0, class_names[i]) query.bindValue(1, scores[i]) if not query.exec_(): QMessageBox.critical(self, "Error", "Failed to insert data into database.") # 查询数据库中的数据 query.exec_("SELECT * FROM detection_results") # 在窗口中显示数据库内容 result = "" while query.next(): id = query.value(0) class_name = query.value(1) score = query.value(2) result += f"ID: {id}, Name: {class_name}\n" QMessageBox.information(self, "Save to Database", "Data saved to database successfully.") # 显示数据库内容 QMessageBox.information(self, "Detection Results", result) def on_open_camera_button_clicked(self): while True: ret, frame = self.cap.read() self.display_image(frame) self.fps += 1 mfps = self.fps / (time.time() - self.t_start) cv2.putText(frame, "FPS " + str(int(mfps)), (10, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) # 检查是否按下了ESC键 k = cv2.waitKey(30) & 0xff if k == 27: # 按下ESC退出循环 break def on_save_image_button_clicked(self): # 保存图像 sess = tf.compat.v1.Session(graph=self.detection_graph) image = self.label.pixmap().toImage() image.save('temp.jpg') frame=cv2.imread('temp.jpg') MODEL_NAME = 'ssdlite_mobilenet_v2_coco_2018_05_09' # fast PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb' print(PATH_TO_CKPT) PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt') NUM_CLASSES = 90 IMAGE_SIZE = (12, 8) fileAlreadyExists = os.path.isfile(PATH_TO_CKPT) if not fileAlreadyExists: print('Model does not exsist !') exit # LOAD GRAPH print('Loading...') detection_graph = tf.Graph() with detection_graph.as_default(): od_graph_def = tf.compat.v1.GraphDef() with tf.io.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: serialized_graph = fid.read() od_graph_def.ParseFromString(serialized_graph) tf.import_graph_def(od_graph_def, name='') label_map = label_map_util.load_labelmap(PATH_TO_LABELS) categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) category_index = label_map_util.create_category_index(categories) print('Finish Load Graph..') # Main t_start = time.time() fps = 0 with detection_graph.as_default(): with tf.compat.v1.Session(graph=detection_graph) as sess: frame = cv2.imread('temp.jpg') ############## image_np_expanded = np.expand_dims(frame, axis=0) image_tensor = detection_graph.get_tensor_by_name('image_tensor:0') detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0') detection_scores = detection_graph.get_tensor_by_name('detection_scores:0') detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') num_detections = detection_graph.get_tensor_by_name('num_detections:0') print('Running detection..') (boxes, scores, classes, num) = sess.run( [detection_boxes, detection_scores, detection_classes, num_detections], feed_dict={image_tensor: image_np_expanded}) print('Done. Visualizing..') vis_utils.visualize_boxes_and_labels_on_image_array( frame, np.squeeze(boxes), np.squeeze(classes).astype(np.int32), np.squeeze(scores), category_index, use_normalized_coordinates=True, line_thickness=8) class_names = [self.category_index[class_id]['name'] for class_id in classes[0]] scores = scores[0] classs=[] scoress=[] for index, score in enumerate(scores): if score >= 0.6: classs.append(class_names[index]) scoress.append(score) self.class_names = classs self.scores = scoress cv2.imwrite('zhoudetest.jpg',frame) QMessageBox.information(self, "Save Image", "Image saved successfully.") #print(class_names[0]) #print(scores[0]) def display_image(self, frame): # 将OpenCV图像转换为Qt图像 h, w, ch = frame.shape bytes_per_line = ch * w qt_image = QImage(frame.data, w, h, bytes_per_line, QImage.Format_RGB888).rgbSwapped() # 在标签中显示图像 pixmap = QPixmap.fromImage(qt_image) self.label.setPixmap(pixmap) def update_datetime(self): # 获取当前日期和时间 current_datetime = QDateTime.currentDateTime() # 格式化日期和时间为字符串 datetime_str = current_datetime.toString("yyyy-MM-dd hh:mm:ss") # 在标签上显示日期和时间 self.datetime_label.setText(datetime_str) if __name__ == '__main__': app = QApplication(sys.argv) window = MainWindow() window.show() sys.exit(app.exec_())
zhoude
最新推荐文章于 2024-08-19 10:00:46 发布