zhoude

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