import tensorflow as tf
from PyQt5.QtGui import QFont,QIcon,QPixmap
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QTabWidget, QPushButton,QHBoxLayout,QFileDialog
import cv2
import sys
from PIL import Image
import numpy as np
import shutil


class MainWindow(QTabWidget):
    # 初始化
    def __init__(self):             #__init__() 是python中的构造函数，在创建对象的时"自动调用"。
        super().__init__()          #调用父类__init__方法
        self.setWindowIcon(QIcon('./images/logo.png'))#设置窗体图标
        self.setWindowTitle('垃圾分类系统')  # 修改系统名称
        # 模型初始化
        self.model = tf.keras.models.load_model("./model/garbage_classification_model.h5")  # 加载模型
        self.class_names = ['纸板','玻璃','金属','纸张','塑料','厨余垃圾'] # 类名
        self.resize(1000, 800)#设置窗体大小
        self.initUI()

    # 界面初始化，设置界面布局
    def initUI(self):
        main_widget = QWidget()#QWidget直接提供的成员函数
        main_layout = QHBoxLayout()#QBOXLayout类可以在水平和垂直方向上排列控件
        font = QFont('宋体', 15)#设置字体

        # 主页面，设置组件并在组件放在布局上
        left_widget = QWidget()#QWidget直接提供的成员函数
        left_layout = QVBoxLayout()#QBOXLayout类可以在水平和垂直方向上排列控件
        img_title = QLabel("图片显示")#QLabel 类代表标签，它是一个用于显示文本或图像的窗口部件
        img_title.setFont(font)#设置字体
        img_title.setAlignment(Qt.AlignCenter)#文字垂直水平居中
        self.img_label = QLabel()#QLabel 类代表标签，它是一个用于显示文本或图像的窗口部件
        self.img_label.setPixmap(QPixmap("./images/background.jpg"))#设置背景图片
        left_layout.addWidget(img_title)#将按钮添加到布局中
        left_layout.addWidget(self.img_label, 1, Qt.AlignCenter)
        left_widget.setLayout(left_layout)#小部件的布局管理器设置为layout。
        right_widget = QWidget()
        right_layout = QVBoxLayout()
        btn_change = QPushButton("选择图片")#创建按钮
        btn_change.clicked.connect(self.change_img)#点击按钮，执行change_img函数
        btn_change.setFont(font)
        btn_predict = QPushButton("开始识别")
        btn_predict.setFont(font)
        btn_predict.clicked.connect(self.predict_img)
        label_result = QLabel('分类结果')
        self.result = QLabel("等待识别")
        label_result.setFont(QFont('宋体', 18))
        self.result.setFont(QFont('宋体', 18))
        right_layout.addStretch()#在布局器中增加一个伸缩量
        right_layout.addWidget(label_result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addWidget(self.result, 0, Qt.AlignCenter)
        right_layout.addStretch()
        right_layout.addStretch()
        right_layout.addWidget(btn_change)
        right_layout.addWidget(btn_predict)
        right_layout.addStretch()
        right_widget.setLayout(right_layout)
        main_layout.addWidget(left_widget)
        main_layout.addWidget(right_widget)
        main_widget.setLayout(main_layout)

        # 模型页面，设置组件并把组件放在布局上
        model_widget = QWidget()
        model_layout = QVBoxLayout()
        bnt_model = QPushButton("模型选择")
        bnt_leibie = QPushButton("类别选择")
        bnt_model.setFont(QFont('宋体', 24))
        bnt_leibie.setFont(QFont('宋体', 24))
        bnt_model.clicked.connect(self.change_model)
        bnt_leibie.clicked.connect(self.change_leibie)
        model_layout.addWidget(bnt_model)
        model_layout.addWidget(bnt_leibie)
        model_layout.addStretch()
        model_widget.setLayout(model_layout)
        
        # 将控件添加到选项卡
        self.addTab(main_widget, '主页')
        self.addTab(model_widget, '设置')
    #选择模型
    def change_model(self):
        fileName,_=QFileDialog.getOpenFileName(self, 'chose files', '','Image files(*.h5)')
        self.model = tf.keras.models.load_model(fileName)  # 加载模型
        print("选择的模型:", fileName)
    #选择类别
    def change_leibie(self):
        fileName,_=QFileDialog.getOpenFileName(self, 'chose files', '','Image files(*.txt)')
        with open(fileName, 'r',encoding='utf-8') as file:  
             # 使用列表推导式去除每一行末尾的换行符  
            lines = [line.strip() for line in file.readlines()] 
        self.class_names=lines 
        print("选择的类别:",lines)
        
    # 上传并显示图片
    def change_img(self):
        openfile_name = QFileDialog.getOpenFileName(self, 'chose files', '','Image files(*.jpg *.png *jpeg)')  # 打开文件选择框选择文件
        img_name = openfile_name[0]  # 获取图片名称
        if img_name == '':
            pass
        else:
            target_image_name = "./images/tmp_up." + img_name.split(".")[-1]  # 将图片移动到当前目录
            shutil.copy(img_name, target_image_name)
            self.to_predict_name = target_image_name
            img_init = cv2.imread(self.to_predict_name)  # 打开图片
            h, w, c = img_init.shape
            scale = 500 / h
            img_show = cv2.resize(img_init, (0, 0), fx=scale, fy=scale)  # 将图片的大小统一调整到500的高，方便界面显示
            cv2.imwrite("./images/show.png", img_show)
            img_init = cv2.resize(img_init, (224, 224))  # 将图片大小调整到224*224用于模型推理
            cv2.imwrite('./images/target.png', img_init)
            self.img_label.setPixmap(QPixmap("./images/show.png"))
            self.result.setText("等待识别")

    # 预测图片
    def predict_img(self):
        img = Image.open('./images/target.png')  # 读取图片
        img = np.asarray(img)  # 将图片转化为numpy的数组
        outputs = self.model.predict(img.reshape(1, 224, 224, 3))  # 将图片输入模型得到结果
        result_index = int(np.argmax(outputs))
        result = self.class_names[result_index]  # 获得对应的名称
        self.result.setText(result)  # 在界面上做显示

if __name__ == "__main__":
    app = QApplication(sys.argv)
    x = MainWindow()
    x.show()
    sys.exit(app.exec_())