光学目标检测yolov5笔记

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

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

使用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'}

可视化的代码：

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

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

    errorimgs = []
    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

        print('processing:', name)
        try:
            IMG_H, IMG_W, IMG_C = image.shape
        except:
            errorimgs.append(name)
            continue

        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.")

    print(
        f"There are {len(errorimgs)} images failed to visualize, they are:\n{errorimgs}")


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))

等待片刻，从输出文件夹中收割可视化结果。

可视化效果图示例：

三、数据集 train val 划分

为数据集划分train，val。注意在yolov5架构下，希望按照如下目录进行组织：

objdet_data
|——images
|——|——train
|——|——val
|——labels
|————train
|————val

划分train-val的代码：

import os
import os.path as osp
import shutil
import random
import cv2

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

train_ratio = 0.9
random.seed(0)

IMGDIR = osp.join(BASEDIR, 'images')
LABDIR = osp.join(BASEDIR, 'labels')

imgnames = [name for name in os.listdir(IMGDIR) if name.endswith('.jpg')]
labnames = [name for name in os.listdir(LABDIR) if name.endswith('.txt')]
print(len(imgnames))
print(len(labnames))


# create output dir
def mkdir_if_missing(OUTDIR):
    if not osp.exists(OUTDIR):
        os.makedirs(OUTDIR)


OUTDIR = osp.join(BASEDIR, 'trainval')
OUT_IMG_TRAIN = osp.join(OUTDIR, 'images', 'train')
OUT_IMG_VAL = osp.join(OUTDIR, 'images', 'val')
OUT_LAB_TRAIN = osp.join(OUTDIR, 'labels', 'train')
OUT_LAB_VAL = osp.join(OUTDIR, 'labels', 'val')
for DIR in [OUT_IMG_TRAIN, OUT_IMG_VAL, OUT_LAB_TRAIN, OUT_LAB_VAL]:
    mkdir_if_missing(DIR)

# filter available images
available_imgnames = []
for imgname in imgnames:
    path_img = osp.join(IMGDIR, imgname)
    img = cv2.imread(path_img)
    if img is None:
        continue
    else:
        available_imgnames.append(imgname)

print(len(available_imgnames))

random.shuffle(available_imgnames)

available_txtnames = [
    name.split('.')[0] + '.txt' for name in available_imgnames
]

# available_imgnames.sort(key=lambda x: int(x.split('.')[0]))
# available_txtnames.sort(key=lambda x: int(x.split('.')[0]))

thresh = int(len(available_imgnames) * train_ratio)

for i, (imgname,
        txtname) in enumerate(zip(available_imgnames, available_txtnames)):
    if imgname.split('.')[0] != txtname.split('.')[0]:
        print('mismatch! ', imgname, 'VS', txtname)
        continue
    else:
        src_img_path = osp.join(IMGDIR, imgname)
        src_lab_path = osp.join(LABDIR, txtname)
        if i < thresh:  # for train
            dst_img_path = osp.join(OUT_IMG_TRAIN, imgname)
            dst_lab_path = osp.join(OUT_LAB_TRAIN, txtname)
        else:  # for val
            dst_img_path = osp.join(OUT_IMG_VAL, imgname)
            dst_lab_path = osp.join(OUT_LAB_VAL, txtname)
        shutil.copy(src_img_path, dst_img_path)
        shutil.copy(src_lab_path, dst_lab_path)
    print(i, '/', thresh)

OK。。。