轻量级检测框标注工具详解

提示：该项目参考于超级轻量标注工具

---

前言

最近实验室在做一个关于气体检测的项目，网上download的数据集很少，而自己合成的数据集又缺少真实性，所以需要自己拍摄视频，进行数据的手动标注。然而需要标注的数据拍摄好后大概有近一万张，如果使用labelme这样的标注软件标注估计要耗费大量的时间。labelme的标注功能很强大，而且程序运行逻辑也很完善，操作起来也很简单，但是在大量的数据集面前再简单的操作也会变得很复杂，所以我在想能不能针对我们拍摄的数据集特点做一个简单的标注工具。

---

一、数据集特点

拍摄的数据集是在室内空旷的地方拍摄的，雾化器喷出来的烟雾大小，浓度，方向在30-40S的视频中变化很小，只是通过移动摄像头改变烟雾在视频中的位置。这就意味着我可以通过手动标几张基准图片，确定烟雾和源点的大小和相对位置，后面的图像只要确定源点的中心点进行平移即可，然后再对其的大小进行微调就行了。

二、标注工具介绍

这里对轻量级的标注工具做代码解析记录，方便后续再次更改时，能够直接查阅。

1.整体框架

如图所示，整个界面可以分为两个部分，上面的菜单部分和下面的图片标注部分。

所以程序把整体结构分为了main_widget和image_widget两个部分。main_widget包括菜单中的按键组件和4个image_widget子窗口。通过菜单的操作，将图片的信息和标注类别信息传递给4个image_widget，然后image_widget读取图片信息，根据用户鼠标的操作记录和绘制标注框，并提供接口向main_widget传递标注框的信息。简单来说image_widget负责标注，main_widget负责参数设置和数据保存。

2.功能详解

这里对一些主要的功能代码做一下记录。

1.框架构造

这是在main_widget.py的初始化函数中的一段添加image_widget子窗口的代码。这里side是代表每行每列放多少个子窗口。这里可以看到main_widget把image_widget组件存放到image_widget这个列表中，并将其排列在主窗口中。

       vbox = QVBoxLayout()
        for i in range(self.side):
            hbox = QHBoxLayout()
            for j in range(self.side):
                self.image_widgets.append(CImageWidget())
                hbox.addWidget(self.image_widgets[-1])
            vbox.addLayout(hbox)

2.文件读取

文件读取的流程大致为，进入到用户选取的文件夹下，遍历该文件夹下所有的文件，将所有的图像文件信息存放到label_info这个字典中，并初始化标签为none（这里值得一提的是直接读出的文件列表是无序的，如果你的文件是按顺序命名的需要手动排序），然后默认该文件夹的标签保存文件为label.txt。下一步就是读取标签保存文件，对label_info中的图像赋值。

    def slot_btn_open(self):
        self.image_dir = QFileDialog.getExistingDirectory(self, u"选择标定图片文件夹", r'D:\YJS\h2o\h2o\3')
        if os.path.exists(self.image_dir):
            self.btn_back.show()
            self.btn_next.show()
            files = os.listdir(self.image_dir)
            for img_name in files:
                if str(img_name).endswith('txt'):
                    continue
                self.label_info[str(img_name)] = None

            self.label_file = self.image_dir + "/label.txt"
            # 按照名称排序
            tem = {}
            # def cmpare(x, y):
            #     x = int(x.split('.')[0])
            #     y = int(y.split('.')[0])
            #     if x>y:
            #         return 1
            #     elif x==y:
            #         return 0
            #     else:
            #         return -1
            for i in sorted(self.label_info, key=lambda x: int(x.split('.')[0])):
                tem[i] = self.label_info[i]
            self.label_info = tem
                
            print("slot_btn_open:", "total images:", len(self.label_info))

            self.read_label_file()
            self.slot_btn_next()

标签读取函数 read_label_file()中首先判断标签文件是否存在，不存在则直接返回，label_info的值仍然全部为none。这里标签的保存格式为
imageName location0 type0 location1 type1 …
所以根据这种格式对txt文件中的信息进行字符串解析，将得到的边框信息存放到label_info的对应value中。

    def read_label_file(self):
        if os.path.exists(self.label_file):
            with open(self.label_file, 'r') as f:
                for line in f:
                    info = line.strip('\r\n')
                    if len(info) == 0:
                        continue
                    domains = info.split(' ')
                    name = domains[0]
                    boxes = []
                    img_path = self.image_dir + '/' + name
                    if not os.path.exists(img_path):
                        print('error:', img_path, 'not exist, ignore it')
                        continue

                    if len(domains) > 1:
                        boxes_str = domains[1:]
                        assert (len(boxes_str) % 5 == 0)
                        box_count = int(len(boxes_str) / 5)
                        for i in range(box_count):
                            box_str = boxes_str[i*5:(i+1)*5]
                            box = [float(x) for x in box_str]
                            boxes.append(box)
                    self.label_info[name] = boxes

然后就是对图片信息的展示，算了感觉代码好简单，没啥讲的必要了，就是将先获取每个image_widget当前的图片和标注信息，然后保存到label_info中，然后再获得要展示的图片信息和标签信息传递给imagge_widget进行展示。

    def slot_btn_next(self):
        self.save_box_info()
        image_names = self.update_batch_index(next=True, pre=False)
        if image_names is not None:
            for i in range(self.total):
                if i + 1 <= len(image_names):
                    img_path = self.image_dir + '/' + image_names[i]
                    self.image_widgets[i].set_info(img_path, self.label_info[image_names[i]])
                else:
                    self.image_widgets[i].set_info(None, None)

3.文件标注

文件标注image_widget组件的代码其实也挺简单的，梳理一下大致流程。根据得到的图片信息读取图片并将已存在的标注信息保存，然后绘制图片和标注框。标注的时候检测鼠标点击操作，鼠标点击左键，记录初始位置，鼠标按住移动时时刻记录鼠标位置，然后根据当前位置和初始位置绘制标注框形成动态的效果。鼠标松开时记录结束位置，将结束位置和初始位置和类别信息保存为一个标注框记录。点击鼠标右键时删除标注记录中最晚的一个框。代码就不贴了.

---

总结

这里简单的对标注工具的代码逻辑做了解析，虽然代码很简单，但是还是佩服原作者的编程能力，能够自己写出来还是很了不起的，这里贴上我根据自己数据集的特点对标注工具的进一步扩展https://github.com/TNTSAYou/labeltool，里面的功能实现了不少，但代码逻辑还存在一些问题，等有空再完善一下。