【深度学习】【工具向】做一个直观展示和检查COCO数据集的工具

1. 背景

最近在研究深度学习领域的目标检测和实例分割，发现很多框架默认的数据集格式是COCO格式，如果刚好手头上有其他格式的数据集，那么就需要借助官方的工具或自己编写工具进行转换，转换完成后最初可能会有很多格式问题（新手），而且碰到一个新的别人提供的数据集，也不能直观的浏览图片上的表述区域和类别，于是，就想要编写一个工具，可以打开annotation文件，分析和展示里面的标注信息，并且可以一键诊断问题

2. 效果图

2.1 浏览标注区域和类别 

2.2 一键诊断问题

 成功

失败

  

3. 技术领域

基本没啥技术难点，主要就是PyQt，json解析和数据集格式的掌握，下面面向小白对几个名词做出解释：

3.1 PyQt

PyQt 是一个 Python 库，用于创建图形用户界面（GUI）。它是由 Riverbank Computing 开发，是一个跨平台的应用程序开发框架。PyQt 允许开发者使用 Python 语言创建具有各种控件的图形界面，如窗口、按钮、文本框等。

PyQt 支持 Qt 的所有模块，并且具有高度的可定制性。你可以使用它的信号和槽机制来处理事件和交互。此外，PyQt 还提供了大量的工具和库，可以帮助你更轻松地创建复杂的 GUI 应用程序。

总的来说，PyQt 是一个强大而灵活的库，可以让你用 Python 语言快速创建高质量的 GUI 应用程序。无论你是需要创建简单的界面，还是需要构建复杂的应用程序，PyQt 都能够满足你的需求。

3.2 JSON

JSON，全称JavaScript Object Notation，是一种轻量级的数据交换格式。它基于JavaScript的子集，易于人阅读和编写，同时也易于机器解析和生成。JSON 是基于文本的，并且独立于语言。这些特性使得 JSON 成为理想的数据交换语言。

3.3 COCO数据集

COCO数据集的格式主要包含以下部分：

info：表示测试/验证/训练所有实例的整体信息，如创建年份、版本等。
licenses：包含数据的许可证信息。
images：包含多个对应实例的数组，存储图像的相关信息。
annotations：包含多个对应实例的数组，存储标注信息。

COCO数据集使用json文件存储，每个json文件包括上述四个基本类型。

4. 上代码

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QVBoxLayout, QPushButton, QLabel, QWidget, QHBoxLayout, QTextEdit, QProgressBar
from PyQt5.QtGui import QPixmap, QImage, QPainter, QColor, QFont, QPen, QPalette
from PyQt5.QtCore import Qt
import json
import os
import random

class COCOAnnotationChecker(QMainWindow):
    def __init__(self):
        super().__init__()
        self.error_count = 0
        self.initUI()

    def initUI(self):
        self.setWindowTitle('COCO Annotation Checker')

        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.setMinimumSize(700, 500)
        
        self.layout = QVBoxLayout()
        self.image_layout = QHBoxLayout()
        self.buttons_layout = QHBoxLayout()
        
        self.text_show = QTextEdit()
        self.text_show.setMaximumWidth(300)

        self.image_label = QLabel()

        palette = QPalette()
        color = QColor(78, 78, 78)  # 创建一个颜色对象
        self.image_label.setStyleSheet(f"background-color: {color.name()};")  # 将颜色设置为样式表的一部分
        self.image_label.setMinimumWidth(700)
        self.image_label.setMinimumHeight(500)
        
        self.open_button = QPushButton('Open')
        self.open_button.clicked.connect(self.load_annotations)
        self.open_button.setMinimumHeight(30)
        
        self.previous_button = QPushButton('Previous')
        self.previous_button.clicked.connect(self.show_previous_image)
        self.previous_button.setMinimumHeight(30)
        
        self.next_button = QPushButton('Next')
        self.next_button.clicked.connect(self.show_next_image)
        self.next_button.setMinimumHeight(30)
        
        self.check_button = QPushButton('Check all')
        self.check_button.clicked.connect(self.check_all_annotations)
        self.check_button.setMinimumHeight(30)
        
        self.progress_bar = QProgressBar()
        self.progress_bar.setMinimumWidth(50)
        self.progress_bar.setVisible(False)
        
        self.central_widget.setLayout(self.layout)
        self.image_layout.addWidget(self.image_label)
        self.image_layout.addWidget(self.text_show)
        self.layout.addLayout(self.image_layout)
        self.buttons_layout.addWidget(self.open_button)
        self.buttons_layout.addWidget(self.previous_button)
        self.buttons_layout.addWidget(self.next_button)
        self.buttons_layout.addWidget(self.check_button)
        self.buttons_layout.addWidget(self.progress_bar)
        self.layout.addLayout(self.buttons_layout)

        self.annotations_file = None
        self.images_dir = None
        self.current_image_index = 0
        
    def get_resize_image(self, image : QImage):
        zoom = 1.0
        if (image.height() / image.width() > self.image_label.height() / self.image_label.width()):
            zoom = image.height() / self.image_label.height()
        else:
            zoom = image.width() / self.image_label.width()
        new_image = image.scaled(int(image.width()/zoom), int(image.height()/zoom))
        return new_image
    
    def record_error(self, message):
        self.error_count += 1
        self.text_show.append(str(self.error_count) + ". " + message + "\n")
        
    def check_all_annotations(self):
        try:
            self.progress_bar.setVisible(True)
            percent = 0
            self.progress_bar.setValue(percent)
            self.text_show.clear()
            self.error_count = 0
            self.show_error_on_pic("Checking all annotation elements...")
            if self.annotations_file is None:
                self.record_error("标注文件不存在")
                return
            if self.images_dir is None or self.current_image_index >= len(os.listdir(self.images_dir)) - 1:
                self.record_error("图片文件夹 {} 不存在".format(self.images_dir))
                return
            self.current_image_index += 1
            with open(self.annotations_file, 'r') as f:
                annotations = json.load(f)
            if (annotations.get("images") is None):
                self.record_error("images 栏目不存在")
                return
            if (annotations.get("annotations") is None):
                self.record_error("annotations 栏目不存在")
                return
            if (annotations.get("categories") is None):
                self.record_error("categories 栏目不存在")
                return
                
            images_list = annotations['images']
            image_item_index = 0
            item_index = 0
            annotations_check = False
            step_per_percent = int(len(images_list) * len(annotations['annotations']) / 100)
            for image_item in images_list:
                image_item_index += 1
                if (len(image_item) == 0):
                    self.record_error("标注文件中 images 栏目为空")
                    return
                
                column_missing = False
                if (image_item.get('id') is None):
                    self.record_error("在images栏目中第{}个图片信息中的 id 栏目不存在".format(image_item_index))
                    column_missing = True
                if (image_item.get('file_name') is None):
                    self.record_error("在images栏目中第{}个图片信息中的 file_name 栏目不存在".format(image_item_index))
                    column_missing = True
                if (image_item.get('height') is None):
                    self.record_error("在images栏目中第{}个图片信息中的 height 栏目不存在".format(image_item_index))
                    column_missing = True
                if (image_item.get('width') is None):
                    self.record_error("在images栏目中第{}个图片信息中的 width 栏目不存在".format(image_item_index))
                    column_missing = True
                if (column_missing):
                    continue
                
                image_id = image_item['id']
                image_name = image_item['file_name']

                image_path = os.path.join(self.images_dir, image_name)
                if not (os.path.isfile(image_path)):
                    self.record_error("图片文件不存在\n{}".format(image_path))
                
                annotation_list = annotations['annotations']
                index = 0
                image_exist = False
                id_list = set()
                
                for annotation in annotation_list:
                    index += 1
                    item_index += 1
                    if (item_index % step_per_percent == 0):
                        percent += 1
                        self.progress_bar.setValue(percent)
                        QApplication.processEvents()
                        
                    if (len(annotation) == 0):
                        self.record_error('在annotations栏目中第 {} 个标注信息为空'.format(index))
                    if not (annotations_check):
                        if (annotation.get("image_id") is None):
                            self.record_error('在annotations栏目中第 {} 个标注信息image_id 部分不存在'.format(index))
                        if (annotation.get("category_id") is None):
                            self.record_error('在annotations栏目中第 {} 个标注信息category_id 部分不存在'.format(index))
                        if (annotation.get("bbox") is None):
                            self.record_error('在annotations栏目中第 {} 个标注信息bbox 部分不存在'.format(index))
                        if (annotation.get("area") is None):
                            self.record_error('在annotations栏目中第 {} 个标注信息area 部分不存在'.format(index))
                    
                        if (annotation.get("id") is not None):
                            if not (annotation["id"] in id_list):
                                id_list.add(annotation["id"])
                            else:
                                self.record_error('在annotations栏目中第 {} 个标注信息 id 部分是重复的'.format(index))
                    if (annotation.get("image_id") is not None):  
                        if (annotation["image_id"] == image_id):
                            image_exist = True
                        
                annotations_check = True
                    
                if not (image_exist):
                    self.record_error("图像文件 {} image_id {} 在annotations栏目中没有标注信息".format(self.images_dir, image_id))
        except Exception as e:
            self.record_error("检查异常: {}".format( e))
        finally:
            if (self.error_count == 0):
                self.show_error_on_pic("标注文件验证成功", QColor(0, 255, 0))
            else:
                self.show_error_on_pic("标注文件验证失败", QColor(255, 0, 0))
                
    def show_annotations(self, par_image_id, par_image_file):
        if self.annotations_file is None:
            self.show_error_on_pic("annotation file not exist")
            return

        with open(self.annotations_file, 'r') as f:
            annotations = json.load(f)
        annotation_list = annotations['annotations']
        annotation = [item for item in annotation_list if item["image_id"] == par_image_id]
        if (len(annotation) == 0):
            self.show_error_on_pic('Image id {} not exist'.format(par_image_id))
            return
        
        image_path = os.path.join(self.images_dir, par_image_file)
        image = QImage(image_path)
        pen_color_table = {}
        for annotation_item in annotation:
            category_id = annotation_item['category_id']
            bbox = annotation_item['bbox']
            painter = QPainter(image)
            
            pen_color = QColor()
            if (pen_color_table.get(category_id) is not None):
                pen_color = pen_color_table[category_id]
            else:
                pen_color = self.get_random_color(category_id)
                pen_color_table[category_id] = pen_color
                
            pen = QPen(pen_color, 5)
            painter.setPen(pen)
            painter.drawRect(int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))
            # 设置画笔颜色和宽度
            painter.setPen(QColor(255, 255, 255))
            # 设置字体大小和样式
            font = QFont('Arial', 12, QFont.Bold)
            painter.setFont(font)
            painter.drawText(int(bbox[0]), int(bbox[1]), 50, 20, Qt.AlignmentFlag.AlignLeft, str(category_id))
            painter.end()
        image = self.get_resize_image(image)
        pixmap = QPixmap.fromImage(image)
        self.image_label.setPixmap(pixmap)
        self.image_label.setAlignment(Qt.AlignCenter)
        
    def get_random_color(self, seed):
        offset = seed % 10
        r = random.randint(155, 255) - offset * 10
        g = random.randint(155, 255) - offset * 10
        b = random.randint(155, 255) - offset * 10
        return QColor(int(r), int(g), int(b))

    def show_previous_image(self):
        self.progress_bar.setVisible(False)
        if self.images_dir is None or self.current_image_index >= len(os.listdir(self.images_dir)) - 1:
            self.show_error_on_pic("Image not exist")
            return
        self.current_image_index -= 1
        with open(self.annotations_file, 'r') as f:
            annotations = json.load(f)
        images_list = annotations['images']
        image_item = images_list[self.current_image_index]
        
        if (len(image_item) == 0):
            self.show_error_on_pic("Annotation file error")
            return
        image_id = image_item['id']
        image_name = image_item['file_name']

        image_path = os.path.join(self.images_dir, image_name)
        if not (os.path.isfile(image_path)):
            self.show_error_on_pic("Image file not exist\n{}".format(image_path))
            return

        self.show_annotations(image_id, image_name)

    def get_image_id(self, file_name):
        with open(self.annotations_file, 'r') as f:
            annotations = json.load(f)
        images_list = annotations['images']
        image_item = {}
        for item in images_list:
            if (item["file_name"] == file_name):
                image_item = item  
        #image = [item for item in images_list if (item["file_name"] == file_name)]
        if (len(image_item) == 0):
            return
        image_id = image_item['id']
        return image_id
    
    def show_error_on_pic(self, message, color = QColor(255, 255, 255)):
        image = QImage(700, 500, QImage.Format_RGB32)

        # 创建一个QColor对象，表示黑色
        black_color = QColor(0, 0, 0)

        # 使用QImage的fill方法填充背景色
        image.fill(black_color)

        painter = QPainter(image)
        painter.setPen(color)
        # 设置字体大小和样式
        font = QFont('Arial', 15, QFont.Bold)
        painter.setFont(font)
        painter.drawText(int(150), int(150), 500, 200, Qt.AlignmentFlag.AlignLeft, message)
        painter.end()
        pixmap = QPixmap.fromImage(image)
        self.image_label.setPixmap(pixmap)
        self.image_label.setAlignment(Qt.AlignCenter)
        
    def show_next_image(self):
        self.progress_bar.setVisible(False)
        if self.images_dir is None or self.current_image_index >= len(os.listdir(self.images_dir)) - 1:
            self.show_error_on_pic("Image not exist")
            return
        self.current_image_index += 1
        with open(self.annotations_file, 'r') as f:
            annotations = json.load(f)
        images_list = annotations['images']
        image_item = images_list[self.current_image_index]
        
        if (len(image_item) == 0):
            self.show_error_on_pic("Annotation file error")
            return
        image_id = image_item['id']
        image_name = image_item['file_name']

        image_path = os.path.join(self.images_dir, image_name)
        if not (os.path.isfile(image_path)):
            self.show_error_on_pic("Image file not exist\n{}".format(image_path))
            return

        self.show_annotations(image_id, image_name)

    def load_annotations(self):
        self.progress_bar.setVisible(False)
        options = QFileDialog.Options()
        options |= QFileDialog.ReadOnly
        file_name, _ = QFileDialog.getOpenFileName(self, "Load Annotations", "", "JSON Files (*.json);;All Files (*)", options=options)
        if file_name:
            self.error_count = 0
            self.annotations_file = file_name
            self.images_dir = os.path.dirname(os.path.dirname(file_name)) + "/images"
            self.current_image_index = -1
            self.show_next_image()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    checker = COCOAnnotationChecker()
    checker.show()
    sys.exit(app.exec_())