解决问题四部曲：挑战、任务、对象、方法

正框光学目标检测（任务）

使用Python-Yolov5（方法）

一、观察标注格式（对象）

在这里插入图片描述

示例

0.txt

2 0.513333 0.493333 0.240000 0.653333

二、对数据及其标注进行可视化，便于观察分析（对象）

可视化时，需要对不同类别的目标用不同的颜色的框描绘，因此先要知晓总共都有几类，各类的序号分别是几。

统计类别序号的代码：

see_clsnum.py

import os.path as osp
import os

BASEDIR = osp.dirname(osp.abspath(__file__))

LABELDIR = osp.join(BASEDIR, 'labels')

clsnums = set()

for name in os.listdir(LABELDIR):
    path_txt = osp.join(LABELDIR, name)
    with open(path_txt, 'r') as fp:
        lines = fp.readlines()
    for line in lines:
        clsnum = line.split()[0]
        if clsnum not in clsnums:
            clsnums.add(clsnum)

print(clsnums)

输出示例：

{'5', '9', '0', '2', '3', '6', '4', '8', '1', '7'}

可视化的代码：

下面这段脚本可以直接运行。注意先按需求调整好类别号-颜色对应字典，并指定好源图片文件夹、标注文件夹、可视化输出文件夹的路径，再运行。

import cv2
import numpy as np
import os
import os.path as osp


color_dict = {
    0: (255, 000, 000),
    1: (255, 128, 000),
    2: (255, 255, 000),
    3: (000, 255, 000),
    4: (000, 255, 255),
    5: (000, 000, 255),
    6: (128, 000, 255),
    7: (255, 000, 255),
    8: (128, 000, 000),
    9: (000, 128, 000),
}


def run_watch_txts(img_dir, txt_dir, out_dir, color):
    IMAGEDIR = osp.abspath(img_dir)
    LABELDIR = osp.abspath(txt_dir)
    OUTDIR = osp.abspath(out_dir)

    if not os.path.exists(OUTDIR):
        os.makedirs(OUTDIR)

    png_names = set([name[:-4]
                    for name in os.listdir(IMAGEDIR) if name.endswith('.jpg')])
    txt_names = set([name[:-4]
                    for name in os.listdir(LABELDIR) if name.endswith('.txt')])
    intersection_names = png_names & txt_names

    for name in png_names:
        image = cv2.imread(osp.join(IMAGEDIR, name+".jpg"))

        if name not in intersection_names:
            cv2.imwrite(osp.join(OUTDIR, name+".jpg"), image)
            continue

        IMG_H, IMG_W, IMG_C = image.shape

        with open(osp.join(LABELDIR, name+".txt"), 'r') as fp:
            lines = fp.readlines()

        boxes = []
        for line in lines:
            line = line.strip().split()
            # print(line)
            label = int(line[0])
            cx = IMG_W * float(line[1])
            cy = IMG_H * float(line[2])
            w = IMG_W * float(line[3])
            h = IMG_H * float(line[4])
            # angle = int(line[5])
            rect = ((cx, cy), (w, h), 0)
            box = np.int0(cv2.boxPoints(rect))
            boxes.append(box)

        cv2.drawContours(image, boxes, -1, color_dict[label][::-1], 2)
        cv2.imwrite(osp.join(OUTDIR, name+".jpg"), image)
        print(name + ", done.")


if __name__ == '__main__':
    BASEDIR = osp.dirname(osp.abspath(__file__))

    img_dir = osp.join(BASEDIR, 'images')
    txt_dir = osp.join(BASEDIR, 'labels')
    out_dir = osp.join(BASEDIR, 'watch_gtboxes')

    run_watch_txts(img_dir, txt_dir, out_dir, color=(250, 200, 250))