图片查看器demo

main_opencv.py

# main_window.py
import cv2
import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QFileDialog, QVBoxLayout, QPushButton,
                             QLabel, QWidget, QHBoxLayout, QGraphicsView, QGraphicsScene,
                             QStatusBar, QAction, QMenuBar, QMenu, QTextEdit, QLineEdit, QMessageBox, QInputDialog, QComboBox, QDesktopWidget, QGraphicsRectItem)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QMouseEvent, QWheelEvent
from PyQt5.QtCore import Qt, QPointF, QRectF
from background_color_dialog import BackgroundColorDialog, ImageViewer
import xml.etree.ElementTree as ET
from pathlib import Path
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.defects = []
        self.HaveLoadImage = False
        self.HaveLoadBiaozhu = False
    def init_ui(self):
        self.setWindowTitle('Image Viewer')
        self.setGeometry(100, 100, 800, 600)

        self.create_central_widget()
        self.create_status_bar()
        self.create_menu_bar()

    # 中心控件
    def create_central_widget(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout()

        # Image display area
        self.image_viewer = ImageViewer(self)  # Assuming ImageViewer is defined elsewhere
        self.image_viewer.update_dropdown_signal.connect(self.update_dropdown)
        main_layout.addWidget(self.image_viewer, 3)
        right_layout = QVBoxLayout()

        self.add_image_load_section(right_layout)
        self.add_algorithm_load_section(right_layout)
        self.add_image_navigation_section(right_layout)
        self.add_annotation_section(right_layout)
        self.deleteimg(right_layout)
        self.add_annotation_dropdown(right_layout)
        self.add_status_bar_section(right_layout)

        main_layout.addLayout(right_layout)
        central_widget.setLayout(main_layout)

    # 加载图片
    def add_image_load_section(self, layout):
        imageload_layout = QHBoxLayout()
        self.load_button = QPushButton('加载图片')
        self.load_button.clicked.connect(self.load_image)
        imageload_layout.addWidget(self.load_button)

        self.status_input = QLineEdit()
        imageload_layout.addWidget(self.status_input)
        layout.addLayout(imageload_layout)

    # 加载标注文件夹
    def add_algorithm_load_section(self, layout):
        algload_layout = QHBoxLayout()
        self.load_alg = QPushButton('加载标注文件夹')
        self.load_alg.clicked.connect(self.load_biaozhu_dir)  # Corrected method name
        algload_layout.addWidget(self.load_alg)

        self.alg_input = QLineEdit()
        algload_layout.addWidget(self.alg_input)
        layout.addLayout(algload_layout)

    def add_image_navigation_section(self, layout):
        select_image_layout = QHBoxLayout()
        select_image_layout.setSpacing(20)

        self.previousImage_button = QPushButton('上一张')
        self.previousImage_button.setFixedSize(70, 20)
        self.previousImage_button.setStyleSheet(self.get_button_style())
        self.previousImage_button.clicked.connect(self.clickpreviousImage)
        select_image_layout.addWidget(self.previousImage_button)

        self.nextImage_button = QPushButton('下一张')
        self.nextImage_button.setFixedSize(70, 20)
        self.nextImage_button.setStyleSheet(self.get_button_style())
        self.nextImage_button.clicked.connect(self.clicknextImage)
        select_image_layout.addWidget(self.nextImage_button)

        layout.addLayout(select_image_layout)

    def add_annotation_section(self, layout):
        jiancebiaozhu_layout = QHBoxLayout()
        self.jiancepencile_button = QPushButton('目标检测标注')
        self.jiancepencile_button.clicked.connect(self.drawjiance)
        jiancebiaozhu_layout.addWidget(self.jiancepencile_button)

        self.quxiao_button = QPushButton('取消标注')
        self.quxiao_button.clicked.connect(self.drawjiance)  # Corrected method name
        jiancebiaozhu_layout.addWidget(self.quxiao_button)
        layout.addLayout(jiancebiaozhu_layout)

    def deleteimg(self, layout):
        self.deletimg_button = QPushButton("删除该张图片")
        self.deletimg_button.clicked.connect(self.deletimg_hanshu)
        layout.addWidget(self.deletimg_button)

    def add_annotation_dropdown(self, layout):
        defect_dropdown_layout = QHBoxLayout()
        self.dropdown = QComboBox()
        defect_dropdown_layout.addWidget(self.dropdown)
        layout.addLayout(defect_dropdown_layout)

    def add_status_bar_section(self, layout):
        # 文本状态栏
        status_bar_layout = QVBoxLayout()
        self.text_label = QLabel('<b>文本 <font color="red">状态栏</font></b>')
        status_bar_layout.addWidget(self.text_label)

        self.textbox = QTextEdit(self)
        self.textbox.setReadOnly(True)
        status_bar_layout.addWidget(self.textbox)

        layout.addLayout(status_bar_layout)
        layout.addStretch(1)  # Ensure status bar area is at the bottom

    def create_status_bar(self):
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)

    # 创建tab页
    def create_menu_bar(self):
        menu_bar = self.menuBar()

        file_menu = menu_bar.addMenu('文件')
        open_action = QAction('打开图片', self)
        open_action.triggered.connect(self.load_image)
        file_menu.addAction(open_action)

        open_project = QAction('打开文件夹', self)
        open_project.triggered.connect(self.load_dir)
        file_menu.addAction(open_project)


        detection_menu = menu_bar.addMenu('检测')
        select_detection = QAction('选择检测方法（传统/深度）', self)
        select_detection.triggered.connect(self.load_image)
        detection_menu.addAction(select_detection)

        model_select = QAction('选择模型', self)
        model_select.triggered.connect(self.load_image)
        detection_menu.addAction(model_select)

        annotation_menu = menu_bar.addMenu('标注')
        select_annotation = QAction('选择标注目的（目标分割/目标检测）', self)
        select_annotation.triggered.connect(self.select_biaozhu)
        annotation_menu.addAction(select_annotation)

        settings_menu = menu_bar.addMenu('设置')
        background_color_action = QAction('背景颜色设置', self)
        background_color_action.triggered.connect(self.open_background_color_dialog)
        settings_menu.addAction(background_color_action)

        language_set = QAction('语言设置', self)
        language_set.triggered.connect(self.load_image)
        settings_menu.addAction(language_set)

    def get_button_style(self):
        return (
            "QPushButton {"
            "border-radius: 10px;"
            "background-color: #3418db;"
            "color: white;"
            "border: none;"
            "font-size: 10px;"
            "}"
        )
    def load_image(self):
        file_path, _ = QFileDialog.getOpenFileName(self, 'Open Image File', '', 'Images (*.png *.xpm *.jpg *.bmp)')
        if file_path:
            file_name = Path(file_path).stem + '.xml'
            xml_path = os.path.join(self.biaozhu_dir_path, file_name)
            annotation = self.anlyVOC(xml_path)
            self.image_viewer.load_image(file_path, annotation, xml_path)
            self.status_input.setText(file_path)
            self.textbox.append("加载状态：image\n图片路径：{}".format(file_path))
    def load_dir(self):
        # 打开文件夹选择对话框
        self.dir_path = QFileDialog.getExistingDirectory(self, 'Select Directory')
        if self.dir_path:
            # 获取目录中的所有文件
            self.image_files = [f for f in os.listdir(self.dir_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.xpm'))]
            if not self.image_files:
                QMessageBox.information(self, 'No Images Found', 'No image files found in the selected directory.')
                return
            self.textbox.append("加载状态：dir")
            self.len_imageFiles = len(self.image_files)
            self.add = 0
            self.file_path = os.path.join(self.dir_path, self.image_files[self.add])
            file_name = Path(self.file_path).stem + '.xml'
            if self.HaveLoadBiaozhu == True:
                xml_path = os.path.join(self.biaozhu_dir_path, file_name)
                annotation = self.anlyVOC(xml_path)
                self.image_viewer.load_image(self.file_path, annotation, xml_path)
                self.status_input.setText(f'Loaded: {self.file_path}')
                self.textbox.append("图片路径：{}".format(self.file_path))
            else:
                self.image_viewer.load_image(self.file_path, None, None)
                QMessageBox.information(self, '加载顺序错误', '应当先加载标注，再加载图片！')

    def deletimg_hanshu(self):
        # 弹窗确认删除
        reply = QMessageBox.question(self, '确认删除',
                                     f"确认删除文件 {self.file_path} 吗？",
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.No)

        if reply == QMessageBox.Yes:
            try:
                os.remove(self.file_path)
                self.textbox.append(f"已删除: {self.file_path}")
            except Exception as e:
                self.textbox.append(f"无法删除 {self.file_path}: {e}")
        else:
            self.textbox.append("删除操作已取消")


    def load_biaozhu_dir(self):
        # 加载标注文件夹，当前仅支持加载xmlVOC格式标注
        self.biaozhu_dir_path = QFileDialog.getExistingDirectory(self, 'Select Directory')
        if self.biaozhu_dir_path:
            self.biaozhu_files = [f for f in os.listdir(self.biaozhu_dir_path) if f.lower().endswith(('.xml', '.txt'))]
            self.textbox.append("标注个数：{}\n标注路径：{}".format(len(self.biaozhu_files), self.biaozhu_dir_path))
            if not self.biaozhu_files:
                QMessageBox.information(self, '加载错误', '文件夹中没有有效的标注文件.')
                return

            for file_name in self.biaozhu_files:
                file_path = os.path.join(self.biaozhu_dir_path, file_name)
                objects = self.anlyVOC(file_path)
                if objects != None:
                    for object in objects:
                        defect = object["name"]
                        if defect not in self.defects:
                            self.defects.append(defect)
            self.textbox.append("当前加载的标注有：{}".format(self.defects))
            self.dropdown.clear()
            self.dropdown.addItems(self.defects)
        else:
            QMessageBox.information(self, 'No biaozhu Found', 'No biaozhu files found in the selected directory.')
    def update_dropdown(self, label):
        # 缺陷类别下拉框
        self.dropdown.clear()
        self.dropdown.setCurrentText(label)
    def anlyVOC(self, xmlfilepath):
        # 解析XML文件
        try:
            # 解析 XML 文件
            tree = ET.parse(xmlfilepath)
            root = tree.getroot()
            # 提取物体信息
            objects = []
            for obj in root.findall('object'):
                obj_info = {}
                # 提取类别信息
                obj_info['name'] = obj.find('name').text

                # 提取边界框信息
                bndbox = obj.find('bndbox')
                obj_info['xmin'] = int(bndbox.find('xmin').text)
                obj_info['ymin'] = int(bndbox.find('ymin').text)
                obj_info['xmax'] = int(bndbox.find('xmax').text)
                obj_info['ymax'] = int(bndbox.find('ymax').text)
                objects.append(obj_info)
            return objects
        except Exception as e:
            print(f"An error occurred: {e}")
            return None

    def clickpreviousImage(self):
        # 上一张
        self.add -= 1
        if self.add < 0:
            self.textbox.append("当前图片索引超规，请重新选择！")
            QMessageBox.information(self, '没有找到图片', '当前图片索引超规，请重新选择')
            return
        self.file_path = os.path.join(self.dir_path, self.image_files[self.add])
        file_name = Path(self.file_path).stem + '.xml'
        xml_path = os.path.join(self.biaozhu_dir_path, file_name)
        annotation = self.anlyVOC(xml_path)
        self.image_viewer.load_image(self.file_path, annotation, xml_path)
        self.status_input.setText(f'Loaded: {self.file_path}')
        self.textbox.append("图片路径：{}".format(self.file_path))

    def clicknextImage(self):
        # 下一张
        self.add += 1
        try:
            if self.add >= self.len_imageFiles:
                self.textbox.append("当前图片索引超规，请重新选择！")
                QMessageBox.information(self, '没有找到图片', '当前图片索引超规，请重新选择')
                return
        except:
            self.textbox.append("请先选择图片文件夹！")
            QMessageBox.information(self, '没有找到图片', '当前图片索引超规，请重新选择')
            return
        self.file_path = os.path.join(self.dir_path, self.image_files[self.add])
        file_name = Path(self.file_path).stem + '.xml'
        xml_path = os.path.join(self.biaozhu_dir_path, file_name)
        annotation = self.anlyVOC(xml_path)
        self.image_viewer.load_image(self.file_path, annotation, xml_path)
        self.status_input.setText(f'Loaded: {self.file_path}')
        self.textbox.append("图片路径：{}".format(self.file_path))
        print("图片路径：{}".format(self.file_path))

    def drawjiance(self):
        self.textbox.append("开始标注")
        # self.image_viewer.setDragMode(QGraphicsView.NoDrag)
        if self.image_viewer.drawing == 3:
            self.image_viewer.drawing = 1
        self.image_viewer.set_drawing_mode(self.image_viewer.drawing)
        self.image_viewer.drawing += 1



    def keyPressEvent(self, event):
        # 下一张
        if event.key() == Qt.Key_D:
            self.clicknextImage()
        elif event.key() == Qt.Key_A:
            self.clickpreviousImage()
        super().keyPressEvent(event)


    def select_biaozhu(self):
        # 设置标注策略 分割（点子）或者检测（画框）
        # 弹出对话框让用户选择标注目的
        options = ["目标分割", "目标检测"]
        selected, ok = QInputDialog.getItem(self, "选择标注目的", "请选择标注目的：", options, 0, False)
        if ok and selected:
            if selected == "目标分割":
                self.textbox.append("已经切换标注模式为目标分割！")
                self.set_annotation_mode("segmentation")
            elif selected == "目标检测":
                self.textbox.append("已经切换标注模式为目标检测！")
                self.set_annotation_mode("detection")
        else:
            # 用户取消选择时的处理
            QMessageBox.information(self, "取消选择", "您取消了标注目的的选择。")

    def set_annotation_mode(self, mode):
        # 根据标注模式设置当前模式
        if mode == "segmentation":
            # 设置为目标分割模式
            self.current_mode = "Segnationbiaozhu"
            QMessageBox.information(self, "模式设置", "当前标注模式：目标分割")
        elif mode == "detection":
            # 设置为目标检测模式
            self.current_mode = "Detectionbiaozhu"
            QMessageBox.information(self, "模式设置", "当前标注模式：目标检测")
        else:
            # 处理无效模式
            QMessageBox.warning(self, "无效模式", "未识别的标注模式。")


    def open_background_color_dialog(self):
        dialog = BackgroundColorDialog(self)
        dialog.exec_()



def main():
    app = QApplication(sys.argv)
    window = MainWindow()
    size = window.geometry()
    screen = QDesktopWidget().screenGeometry()
    window.move((screen.width() - size.width()) // 2, (screen.height() - size.height()) // 2 - 120)
    window.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()

background_color_dialog.py

# background_color_dialog.py
from PyQt5.QtWidgets import QDialog, QVBoxLayout, QPushButton, QColorDialog
from PyQt5.QtWidgets import (QVBoxLayout, QPushButton, QGraphicsView, QGraphicsScene,
                             QGraphicsRectItem, QGraphicsEllipseItem)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QMouseEvent, QWheelEvent, QPen, QKeyEvent, QColor
from PyQt5.QtCore import Qt, QPointF, QRectF, pyqtSignal, QObject
import traceback
class BackgroundColorDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle('Background Color Setting')
        self.setGeometry(100, 100, 300, 200)

        layout = QVBoxLayout()

        self.deep_color_button = QPushButton('深色系')
        self.light_color_button = QPushButton('浅色系')

        self.deep_color_button.clicked.connect(self.set_deep_color)
        self.light_color_button.clicked.connect(self.set_light_color)

        layout.addWidget(self.deep_color_button)
        layout.addWidget(self.light_color_button)

        self.setLayout(layout)

    def set_deep_color(self):
        if self.parent() is not None:
            self.parent().setStyleSheet("QMainWindow { background-color: white; color: black; }")

    def set_light_color(self):
        if self.parent() is not None:
            self.parent().setStyleSheet("QMainWindow { background-color: black; color: white; }")

class ResizableControlPoint(QGraphicsEllipseItem):
    def __init__(self, parent, position):
        super().__init__(-2, -2, 5, 5, parent)
        self.setBrush(QColor('blue'))
        self.setFlag(QGraphicsEllipseItem.ItemIsMovable)
        self.setFlag(QGraphicsEllipseItem.ItemSendsGeometryChanges)
        self.setZValue(1)
        self.parent_item = parent
        self.position = position

    def mouseMoveEvent(self, event):
        if self.parent_item:
            rect = self.parent_item.rect()
            pos = self.mapToParent(event.pos())
            if self.position == "topLeft":
                rect.setTopLeft(pos)
            elif self.position == "topRight":
                rect.setTopRight(pos)
            elif self.position == "bottomLeft":
                rect.setBottomLeft(pos)
            elif self.position == "bottomRight":
                rect.setBottomRight(pos)
            self.parent_item.setRect(rect)
            self.parent_item.update_control_points()
        super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event):
        super().mouseReleaseEvent(event)
        self.parent_item.update_control_points()

class EditableRectItem(QGraphicsRectItem):
    def __init__(self, rect, parent=None, label=None):
        super().__init__(rect, parent)
        self.label = label
        self.setPen(QPen(Qt.red, 2))
        self.setFlag(QGraphicsRectItem.ItemIsMovable)
        self.setFlag(QGraphicsRectItem.ItemSendsGeometryChanges)
        self.setFlag(QGraphicsRectItem.ItemIsSelectable)
        self.setFlag(QGraphicsRectItem.ItemIsFocusable)

        self.control_points = []
        self.create_control_points()

    def create_control_points(self):
        positions = [
            "topLeft", "topRight",
            "bottomLeft", "bottomRight"
        ]
        for pos in positions:
            control_point = ResizableControlPoint(self, pos)
            self.control_points.append(control_point)
        self.update_control_points()

    def update_control_points(self):
        rect = self.rect()
        if len(self.control_points) == 4:
            self.control_points[0].setPos(rect.topLeft())
            self.control_points[1].setPos(rect.topRight())
            self.control_points[2].setPos(rect.bottomLeft())
            self.control_points[3].setPos(rect.bottomRight())

    def itemChange(self, change, value):
        if change == QGraphicsRectItem.ItemPositionChange:
            self.update_control_points()
        return super().itemChange(change, value)

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            if self.scene() is not None:
                self.scene().itemClicked.emit(self)  # Notify the scene
        super().mousePressEvent(event)

class GraphicsScene(QGraphicsScene):
    itemClicked = pyqtSignal(object)  # Define a custom signal

    def __init__(self):
        super().__init__()
class ImageViewer(QGraphicsView):
    update_dropdown_signal = pyqtSignal(list)

    def __init__(self, main_window):
        super().__init__()
        self.setRenderHint(QPainter.Antialiasing)
        self.setRenderHint(QPainter.SmoothPixmapTransform)
        self.scene = GraphicsScene()
        self.setScene(self.scene)
        self.image_item = None
        self.main_window = main_window
        self.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)

        self.drawing = False
        self.rect_item = None
        self.start_pos = None
        self.scroll_hand_drag = False
        self.rect_items = []

        self.scene.itemClicked.connect(self.on_rect_item_clicked)


    def load_image(self, file_path, annotation, xmlfilepath):
        try:
            self.xmlfilepath = xmlfilepath
            pixmap = QPixmap(file_path)

            if self.scene:
                self.scene.clear()
            if hasattr(self, 'image_item') and self.image_item:
                self.image_item = None
                self.scene.removeItem(self.image_item)
            a = self.scene.addPixmap(pixmap)
            self.image_item = self.scene.addPixmap(pixmap)
            self.image_item.setTransformationMode(Qt.SmoothTransformation)
            # Set the scene rectangle and fit the view
            self.setSceneRect(QRectF(pixmap.rect()))
            self.fitInView(self.sceneRect(), Qt.KeepAspectRatio)

            # Process annotations
            if annotation:
                for obj in annotation:
                    try:
                        xmin = int(obj['xmin'])
                        ymin = int(obj['ymin'])
                        xmax = int(obj['xmax'])
                        ymax = int(obj['ymax'])
                        label = str(obj['name'])
                        if xmin >= xmax or ymin >= ymax:
                            raise ValueError("Invalid bounding box coordinates.")

                        rect_item = EditableRectItem(QRectF(xmin, ymin, xmax - xmin, ymax - ymin), label=label)
                        rect_item.setPen(QColor('blue'))
                        rect_item.setBrush(QColor(255, 0, 0, 50))
                        self.scene.addItem(rect_item)
                    except Exception as e:
                        print(f"Error processing object: {e}")
        except Exception as e:
            print("Error loading image or drawing annotations:", e)
            traceback.print_exc()

    def handle_update_dropdown(self, labels):
        pass
    def on_rect_item_clicked(self, rect_item):
        rect = rect_item.rect()
        label = rect_item.label
        print(f"Clicked rectangle with label '{label}' at position ({rect.left()}, {rect.top()}), size ({rect.width()}, {rect.height()})")
    def wheelEvent(self, event: QWheelEvent):
        factor = 1.2
        if event.angleDelta().y() < 0:
            factor = 1 / factor
        self.scale(factor, factor)

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton and event.modifiers() == Qt.ControlModifier:
            self.scroll_hand_drag = True
            self.setDragMode(QGraphicsView.ScrollHandDrag)
        elif self.drawing and event.button() == Qt.LeftButton:
            self.start_pos = self.mapToScene(event.pos())
            label = self.main_window.dropdown.currentText()
            self.rect_item = EditableRectItem(QRectF(self.start_pos, self.start_pos), label=label)
            self.scene.addItem(self.rect_item)
            self.rect_items.append(self.rect_item)
        else:
            super().mousePressEvent(event)

    def mouseMoveEvent(self, event: QMouseEvent):
        if self.drawing and self.start_pos:
            end_pos = self.mapToScene(event.pos())
            rect = QRectF(self.start_pos, end_pos).normalized()
            if self.rect_item:
                self.rect_item.setRect(rect)
        else:
            super().mouseMoveEvent(event)

    def mouseReleaseEvent(self, event: QMouseEvent):
        if self.scroll_hand_drag:
            self.scroll_hand_drag = False
            self.setDragMode(QGraphicsView.NoDrag)
        elif self.drawing:
            self.start_pos = None
        else:
            super().mouseReleaseEvent(event)

    def set_drawing_mode(self, drawing: bool):
        self.drawing = drawing
        if drawing:
            self.setDragMode(QGraphicsView.NoDrag)
            self.setCursor(Qt.CrossCursor)
        else:
            self.setDragMode(QGraphicsView.ScrollHandDrag)
            self.setCursor(Qt.ArrowCursor)

    def keyPressEvent(self, event: QKeyEvent):
        if event.modifiers() == Qt.ControlModifier and event.key() == Qt.Key_Z:
            if self.rect_items:
                last_rect = self.rect_items.pop()
                self.scene.removeItem(last_rect)
                print("撤销")
        elif event.key() == Qt.Key_E:
            self.set_drawing_mode(not self.drawing)
        elif event.key() == Qt.Key_S:
            if self.scene:
                items = self.scene.items()
                for item in items:
                    if isinstance(item, QGraphicsRectItem):
                        self.main_window.anlyVOCgenggai(self.scene, self.xmlfilepath)
        super().keyPressEvent(event)