小工具——筛选图像小工具

news2024/11/25 7:06:06

最近在公司手动筛图片，需要将某些含有检测目标的图像手动筛选出来用于做新模型的测试。我最开始是两个文件夹，来回复制粘贴，后来感觉这种效率太低了，就随手写了一个图像筛查小工具。代码如下：

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QAction, QFileDialog
from PyQt5.QtGui import QPixmap, QIcon, QImage
from PyQt5.QtCore import Qt, QDir


class ImageViewer(QMainWindow):
    def __init__(self):
        super().__init__()

        self.current_index = 0
        self.image_files = []
        self.input_folder_path = ""
        self.output_folder_path = ""

        self.init_ui()

    def init_ui(self):
        # 创建布局
        layout = QVBoxLayout()

        # 创建标签用于显示图像
        self.image_label = QLabel(self)
        layout.addWidget(self.image_label)

        # 创建主窗口
        main_widget = QWidget()
        main_widget.setLayout(layout)
        self.setCentralWidget(main_widget)

        # 创建菜单栏
        menu_bar = self.menuBar()
        file_menu = menu_bar.addMenu("文件")
        self.open_input_folder_action = QAction("打开导入图像文件夹", self)
        self.open_output_folder_action = QAction("打开保存图像文件夹", self)
        file_menu.addAction(self.open_input_folder_action)
        file_menu.addAction(self.open_output_folder_action)

        # 绑定事件处理函数
        self.open_input_folder_action.triggered.connect(self.open_input_folder)
        self.open_output_folder_action.triggered.connect(self.open_output_folder)

        # 设置窗口属性
        self.setWindowTitle("图像筛查工具")
        self.setGeometry(100, 100, 800, 600)

    def open_input_folder(self):
        # 打开输入文件夹
        folder_path = QFileDialog.getExistingDirectory(self, "选择输入文件夹", QDir.homePath())
        if folder_path:
            self.input_folder_path = folder_path
            self.load_images()

    def open_output_folder(self):
        # 打开输出文件夹
        folder_path = QFileDialog.getExistingDirectory(self, "选择输出文件夹", QDir.homePath())
        if folder_path:
            self.output_folder_path = folder_path

    def load_images(self):
        # 加载文件夹中的图像文件
        self.image_files = QDir(self.input_folder_path).entryList(["*.png", "*.jpg", "*.jpeg"], QDir.Files)
        if self.image_files:
            self.current_index = 0
            self.show_current_image()

    def show_current_image(self):
        # 显示当前图像
        image_path = QDir(self.input_folder_path).filePath(self.image_files[self.current_index])
        image = QImage(image_path)
        if not image.isNull():
            pixmap = QPixmap.fromImage(image)
            self.image_label.setPixmap(pixmap.scaled(self.image_label.size(), Qt.KeepAspectRatio,
                                                     Qt.SmoothTransformation))
            self.setWindowTitle(f"图像筛查工具 - {self.current_index + 1}/{len(self.image_files)}")

    def keyPressEvent(self, event):
        # 键盘事件处理
        if event.key() == Qt.Key_Right:
            # 切换到下一张图像
            if self.current_index < len(self.image_files) - 1:
                self.current_index += 1
                self.show_current_image()
        elif event.key() == Qt.Key_Left:
            # 切换到上一张图像
            if self.current_index > 0:
                self.current_index -= 1
                self.show_current_image()
        elif event.key() == Qt.Key_Space:
            # 保存当前图像到输出文件夹
            if self.current_index < len(self.image_files):
                image_path = QDir(self.input_folder_path).filePath(self.image_files[self.current_index])
                output_path = QDir(self.output_folder_path).filePath(self.image_files[self.current_index])
                QDir().rename(image_path, output_path)
                self.image_files.pop(self.current_index)
                self.show_current_image()


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

首先要选择两个文件夹，一个用来导入数据，另外一个用来保存数据。导入文件夹，文件夹里面有大量图像，点击键盘右键可以切换到下一张，左键上一张，点击空格可以把这个图像保存到另外一个文件夹中。不过到最后一张的时候可能回异常退出，不过不影响使用所以我就没进行改进了。

下面是该工具的效果：