【学习笔记】pyQt5学习笔记(3）——第一个图像识别demoV2.0

上一个程序（下文说V1.0版本）是批量处理图片文件夹中的所有图片，识别后将图片批量保存。（其实就是Google object detection API的jupyter note book的程序……心虚）。2.0版本中，对程序做出了相应改动，是手动选取需要识别的图片，然后对单张图片进行识别。在识别过程中，原图显示在界面左侧，识别后的图片显示在界面右侧。

因为个人水平所限，程序还存在很多bug，但是可以用起来了。(。◕‿◕。)ノ゜.:。撒fafa～

然后注意，模型文件、pbtxt文件以及类别数设置好了就别点第二次了 ……然后每次更换不同的图片识别就ok，否则会报错……_(:::з」∠)_嘤嘤嘤。 ...

程序完整代码如下，有的时候存在复制代码没有换行及缩进，同时上传了CSDN下载一份……但是介于要收费还是别下载了。

最后，再次恳请大佬指点。这个程序V3.0（发现压根没提到1.0）版本估计就要上摄像头实时识别了……那么，等我做出来再见了，此次一别不知何时再来发帖……

最后的最后，完成代码的关键是1.global将局部变量变成成员变量；2.将numpy.array格式的文件转换为pixmap。

# coding:utf-8
import numpy as np
import os
import os.path
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile
import pylab
from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

# Python 2.x与Python 3.x在定义类时的区别，Python2.x中（）内有object
# 形如 class UiForm(object):
class UiForm():
    
    openfile_name_pb = ''
    openfile_name_pbtxt = ''
    openpic_name = ''
    num_class = 0
    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.resize(600, 420)
        Form.setMinimumSize(QtCore.QSize(600, 420))
        Form.setMaximumSize(QtCore.QSize(600, 420))
        self.frame = QtWidgets.QFrame(Form)
        self.frame.setGeometry(QtCore.QRect(20, 20, 550, 100))
        self.frame.setFrameShape(QtWidgets.QFrame.StyledPanel)
        self.frame.setFrameShadow(QtWidgets.QFrame.Raised)
        self.frame.setObjectName("frame")
        self.horizontalLayout_2 = QtWidgets.QHBoxLayout(self.frame)
        self.horizontalLayout_2.setObjectName("horizontalLayout_2")
        # 加载模型文件按钮
        self.btn_add_file = QtWidgets.QPushButton(self.frame)
        self.btn_add_file.setObjectName("btn_add_file")
        self.horizontalLayout_2.addWidget(self.btn_add_file)
        # 加载pbtxt文件按钮
        self.btn_add_pbtxt = QtWidgets.QPushButton(self.frame)
        self.btn_add_pbtxt.setObjectName("btn_add_pbtxt")
        self.horizontalLayout_2.addWidget(self.btn_add_pbtxt)
        # 打开图片按钮
        self.btn_open_pic = QtWidgets.QPushButton(self.frame)
        self.btn_open_pic.setObjectName("btn_open_pic")
        self.horizontalLayout_2.addWidget(self.btn_open_pic)
        # 输入检测类别数目按钮
        self.btn_enter = QtWidgets.QPushButton(self.frame)
        self.btn_enter.setObjectName("btn_enter")
        self.horizontalLayout_2.addWidget(self.btn_enter)
        # 开始识别按钮
        self.btn_objdec = QtWidgets.QPushButton(self.frame)
        self.btn_objdec.setObjectName("btn_objdec")
        self.horizontalLayout_2.addWidget(self.btn_objdec)
        # 退出按钮
        self.btn_exit = QtWidgets.QPushButton(self.frame)
        self.btn_exit.setObjectName("btn_exit")
        self.horizontalLayout_2.addWidget(self.btn_exit)
        # 显示原图窗口
        self.lab_rawimg_show = QtWidgets.QLabel(Form)
        self.lab_rawimg_show.setGeometry(QtCore.QRect(20, 140, 270, 270))
        self.lab_rawimg_show.setMinimumSize(QtCore.QSize(270, 270))
        self.lab_rawimg_show.setMaximumSize(QtCore.QSize(270, 270))
        self.lab_rawimg_show.setObjectName("lab_rawimg_show")
        self.lab_rawimg_show.setStyleSheet(("border:2px solid red"))
        # 显示识别后的窗口
        self.lab_decimg_show = QtWidgets.QLabel(Form)
        self.lab_decimg_show.setGeometry(QtCore.QRect(310, 140, 270, 270))
        self.lab_decimg_show.setMinimumSize(QtCore.QSize(270, 270))
        self.lab_decimg_show.setMaximumSize(QtCore.QSize(270, 270))
        self.lab_decimg_show.setObjectName("lab_decimg_show")
        self.lab_decimg_show.setStyleSheet(("border:2px solid blue"))
        
        self.retranslateUi(Form)
        # 这里将按钮和定义的动作相连，通过click信号连接openfile槽？
        self.btn_add_file.clicked.connect(self.openpb)
        # 用于打开pbtxt文件
        self.btn_add_pbtxt.clicked.connect(self.openpbtxt)
        # 用于打开图片
        self.btn_open_pic.clicked.connect(self.openpic)
        # 用于用户输入类别数
        self.btn_enter.clicked.connect(self.enter_num_cls)
        # 开始识别
        self.btn_objdec.clicked.connect(self.object_detection)
        # 这里是将btn_exit按钮和Form窗口相连，点击按钮发送关闭窗口命令
        self.btn_exit.clicked.connect(Form.close)
        QtCore.QMetaObject.connectSlotsByName(Form)

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "目标检测"))
        self.btn_add_file.setText(_translate("Form", "加载模型文件"))
        self.btn_open_pic.setText(_translate("Form", "打开图片"))
        self.btn_add_pbtxt.setText(_translate("Form", "加载pbtxt文件"))
        self.btn_enter.setText(_translate("From", "指定识别类别数"))
        self.btn_objdec.setText(_translate("From", "开始识别"))
        self.btn_exit.setText(_translate("Form", "退出"))
        self.lab_rawimg_show.setText(_translate("Form", "原始图片"))
        self.lab_decimg_show.setText(_translate("Form", "识别后图片"))
        
    def openpb(self):
        global openfile_name_pb
        openfile_name_pb, _ = QFileDialog.getOpenFileName(self.btn_add_file,'选择pb文件','/home/kanghao/','pb_files(*.pb)')
        print('加载模型文件地址为：' + str(openfile_name_pb))
        
    def openpic(self):
        global openpic_name
        openpic_name, _  = QFileDialog.getOpenFileName(self.btn_open_pic,'选择识别图片','/home/kanghao/','image files(*.jpg)')
        self.lab_rawimg_show.setPixmap(QPixmap(openpic_name))
        self.lab_rawimg_show.setScaledContents(True)
        print(type(QPixmap(openpic_name)))
        print('加载图片文件地址为：' + str(openpic_name))
        
    def openpbtxt(self):
        global openfile_name_pbtxt
        openfile_name_pbtxt, _ = QFileDialog.getOpenFileName(self.btn_add_pbtxt,'选择pbtxt文件','/home/kanghao/','pbtxt_files(*.pbtxt)')
        print('加载标签文件地址为：' + str(openfile_name_pbtxt))
        
    def enter_num_cls(self):
        global num_class
        num_class, okPressed = QInputDialog.getInt(self.btn_enter,'指定训练类别数','你的目标有多少类？',1,1,28,1)
        if okPressed:
            print('识别目标总类为：' + str(num_class))

    def img2pixmap(self, image):
        Y, X = image.shape[:2]
        self._bgra = np.zeros((Y, X, 4), dtype=np.uint8, order='C')
        self._bgra[..., 0] = image[..., 2]
        self._bgra[..., 1] = image[..., 1]
        self._bgra[..., 2] = image[..., 0]
        qimage = QtGui.QImage(self._bgra.data, X, Y, QtGui.QImage.Format_RGB32)
        pixmap = QtGui.QPixmap.fromImage(qimage)
        return pixmap

    def object_detection(self):
        sys.path.append("..")
        from object_detection.utils import ops as utils_ops

        if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):
            raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')
        
        from utils import label_map_util

        from utils import visualization_utils as vis_util

        # Path to frozen detection graph. This is the actual model that is used for the object detection.
        PATH_TO_FROZEN_GRAPH = openfile_name_pb#ui.openpb()

        # List of the strings that is used to add correct label for each box.
        PATH_TO_LABELS = openfile_name_pbtxt

        NUM_CLASSES = num_class

        detection_graph = tf.Graph()
        with detection_graph.as_default():
          od_graph_def = tf.GraphDef()
          with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
            
        category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)

        def load_image_into_numpy_array(image):
          (im_width, im_height) = image.size
          return np.array(image.getdata()).reshape(
              (im_height, im_width, 3)).astype(np.uint8)
            
        # For the sake of simplicity we will use only 2 images:
        # image1.jpg
        # image2.jpg
        # If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS.
        TEST_IMAGE_PATHS = openpic_name
        print(TEST_IMAGE_PATHS)
        # Size, in inches, of the output images.
        IMAGE_SIZE = (12, 8)

        def run_inference_for_single_image(image, graph):
          with graph.as_default():
            with tf.Session() as sess:
              # Get handles to input and output tensors
              ops = tf.get_default_graph().get_operations()
              all_tensor_names = {output.name for op in ops for output in op.outputs}
              tensor_dict = {}
              for key in [
                  'num_detections', 'detection_boxes', 'detection_scores',
                  'detection_classes', 'detection_masks'
              ]:
                tensor_name = key + ':0'
                if tensor_name in all_tensor_names:
                  tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
                      tensor_name)
              if 'detection_masks' in tensor_dict:
                # The following processing is only for single image
                detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
                detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
                # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
                real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
                detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
                detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
                    detection_masks, detection_boxes, image.shape[0], image.shape[1])
                detection_masks_reframed = tf.cast(
                    tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                # Follow the convention by adding back the batch dimension
                tensor_dict['detection_masks'] = tf.expand_dims(
                    detection_masks_reframed, 0)
              image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

              # Run inference
              output_dict = sess.run(tensor_dict,
                                     feed_dict={image_tensor: np.expand_dims(image, 0)})

              # all outputs are float32 numpy arrays, so convert types as appropriate
              output_dict['num_detections'] = int(output_dict['num_detections'][0])
              output_dict['detection_classes'] = output_dict[
                  'detection_classes'][0].astype(np.uint8)
              output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
              output_dict['detection_scores'] = output_dict['detection_scores'][0]
              if 'detection_masks' in output_dict:
                output_dict['detection_masks'] = output_dict['detection_masks'][0]
          return output_dict
          

        image = Image.open(TEST_IMAGE_PATHS)
        # the array based representation of the image will be used later in order to prepare the
        # result image with boxes and labels on it.
        image_np = load_image_into_numpy_array(image)
        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
        image_np_expanded = np.expand_dims(image_np, axis=0)
        # Actual detection.
        output_dict = run_inference_for_single_image(image_np, detection_graph)
        # Visualization of the results of a detection.
        vis_util.visualize_boxes_and_labels_on_image_array(
            image_np,
            output_dict['detection_boxes'],
            output_dict['detection_classes'],
            output_dict['detection_scores'],
            category_index,
            instance_masks=output_dict.get('detection_masks'),
            use_normalized_coordinates=True,
            line_thickness=8)
        plt.figure(figsize=IMAGE_SIZE)
        plt.imshow(image_np)
        self.lab_decimg_show.setPixmap(self.img2pixmap(image_np))
        self.lab_decimg_show.setScaledContents(True)
        plt.savefig(str(TEST_IMAGE_PATHS)+".jpg")


## 用于显示ui界面的命令
if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    Window = QtWidgets.QWidget()
    # ui为根据类Ui_From()创建的实例
    ui = UiForm()
    ui.setupUi(Window)
    Window.show()
    sys.exit(app.exec_())  

最终运行效果展示。