CV-ToolBox系列之物体检测实现自动标注

CV-ToolBox系列之物体检测实现自动标注

CV-ToolBox:为解放双手而生

---

---

前言

在训练图像分类的网络之前，首要准备的就是标注数据，这是一个令人头疼的问题，流水线重复的操作让人抓狂。但是懒惰是驱使人类进步的根本动力，本着解放自己双手的目的，花一点时间写了自动标注的方法。

---

一、WHERE？

在哪能拿到这个工具：Github

二、HOW?

PS：在使用之前，首先你要有自己预训练的物体检测模型，即使用少量数据训练出来的模型。

1.引入库

代码如下（以yolo为例）：

from PIL import Image
import cv2
from yolo import YOLO
import xml.dom.minidom as DOC
import os


yolo = YOLO()  # 实例化模型

2.读入并生成数据

代码如下：

# ---------------------------------------------------------------------------- #
#  定义需要标注图片目录所在位置
#  遍历轮询推理，得到预标注的bbox，并生成对应的标注文件
# ---------------------------------------------------------------------------- #

images_dir = './mark_img'
images = sorted(os.listdir(images_dir))
for img in images:
    try:
        image = Image.open(os.path.join(images_dir, img))
        img_cv = cv2.imread(os.path.join(images_dir, img))
        img_size = img_cv.shape
        # ---------------------------------------------------------------------------- #
        #  注意此处返回的bbox_list:[[left, top, right, bottom, label],] label为分类的名称
        # ---------------------------------------------------------------------------- #
        r_image, bbox_list = yolo.detect_image(image)
    except Exception as e:
        print("该图片没有识别到物体，无法生成xml文件")
        continue
    else:
        auto_generate_xml(img, bbox_list, img_size, "./data_croped_xml")

该方法有详细的注释，一看就懂。

---

RESULT

talk is cheap, show me the img.这里打开标注工具labelimg，查看生成的标注文件是达到效果。

---

后记

CV-ToolBox：会持续更新更多更好用的图像处理的工具，能帮到你的话，欢迎start！