目标检测之牛仔行头检测（上）—— 读取coco数据集，可视化处理，转换为yolo数据格式

目录

目标检测之牛仔行头检测（上）—— 读取coco数据集，可视化处理，转换为yolo数据格式

前言

一、数据集格式

二、利用torchvision.datasets读取数据集

三、将coco数据集转换为YOLO格式

四、以YOLOV5为baseline进行训练

五、生成val测试集并生成比赛提交文件

总结

---

前言

本系列是记录参加沐神举办的牛仔行头检测比赛，具体规则和内容看下面视频

50 课程竞赛：牛仔行头检测【动手学深度学习v2】

下载地址也是网友提供的，下载完后把后缀改成zip即可

数据集下载地址

这篇文章主要是提供一个引导参加这个比赛的思路，没有过多技巧，先参与动手为主，也欢迎大家和我一起学习讨论。

链接：https://pan.baidu.com/s/1BoO9t4sZFupXZRtYtswq9g 
提取码：e8eh

---

一、数据集格式

 下载完后，数据集的大体目录长这个样子，

其中images是其中所有的图片。

test.csv和valid.csv则为划分好的私榜和公榜数据集，里面都是id和file_name的格式，用于比赛提交使用，现在比赛已经结束了，但我们还是可以通过提交私榜的数据来测试。

 train.json则为coco数据集中目标检测类型标注的信息，我们重点读取其中的信息

其中每个类别信息如下，引至COCO数据集的使用笔记

annotation{
	"id"			: int,	# annotation的id，每个对象对应一个annotation
	"image_id"		: int, 	# 该annotation的对象所在图片的id
	"category_id"	: int, 	# 类别id，每个对象对应一个类别
	"segmentation"	: RLE or [polygon], 
	"area"			: float, 	# 面积
	"bbox"			: [x,y,width,height], 	# x,y为左上角坐标
	"iscrowd"		: 0 or 1,	# 0时segmentation为REL，1为polygon
}

categories[{
	"id"			: int,	# 类别id 
	"name"			: str, 	# 类别名称
	"supercategory"	: str,	# 类别的父类，例如：bicycle的父类是vehicle
}]

二、利用torchvision.datasets读取数据集

我这里没有使用pycocotools来读取coco数据集，而是采用了pytorch官方的torchvision.datasets来读取数据集，方法上其实是差不多的。

from torchvision.datasets.coco import CocoDetection
# coco数据集的目录
coco_image_path = 'D:\CMP\\2021Innovation\dataset\Cowboy\cowboyoutfits\images'
coco_json_path = 'D:\CMP\\2021Innovation\dataset\Cowboy\cowboyoutfits\\train.json'
# 利用官方的CocoDetection读取coco数据集
train_data = CocoDetection(coco_image_path, coco_json_path)
# 读取其中的类别信息
# categories = {cat_info['name']: cat_info['id'] for cat_info in train_data.coco.loadCats(train_data.coco.getCatIds())}
# 这里我提前读取了里面的内容，且对调键值对，方便后续处理
categories = {87: 'belt', 1034: 'sunglasses', 131: 'boot', 318: 'cowboy_hat', 588: 'jacket'}

for image, label in tqdm(train_data):
    # 将PIL格式转为cv
    image = cv.cvtColor(np.array(image), cv.COLOR_RGB2BGR)


    for i in range(len(label)):
        # 读取bbox，其中内容为x1，y1，w，h
        bbox = label[i]['bbox']
        x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])
        class_name = categories[label[i]['category_id']]

        cv.rectangle(image, (x1, y1), (x2, y2), color=(255, 0, 0), thickness=2)
        cv.putText(image, class_name, (x1, y1), cv.FONT_HERSHEY_SIMPLEX, 0.5, color=(0, 0, 255), thickness=2)
    cv.imshow('image', image)
    cv.waitKey()

三、将coco数据集转换为YOLO格式

这里我自己用yolo比较多，所有baseline用的是yolo，所以先要将其数据集转换为yolo数据集。

不过在转换的过程中比较重要的就是如何划分训练集和测试集，因为我们可以通过统计其数据发现其样本十分不均衡，这里就涉及到了均衡样本的问题。不过作为baseline，我这里采用了比较暴力简单的方法，就是直接每个类别都取出来了10个作为val。

 转换代码我写的注释比较多，且也没什么技术含量和繁琐，修改其中的数据集路径即可，就不过多讲解了。

import os
import cv2 as cv
import shutil
import numpy as np
from tqdm import tqdm
from torchvision.datasets.coco import CocoDetection


# 检测并创建需要的文件夹
def check_path(save_path):
    if not os.path.exists(save_path):
        os.mkdir(save_path)
        print("making {} file".format(save_path))


# convert the bounding box from COCO to YOLO format.
def cc2yolo_bbox(img_width, img_height, bbox):
    dw = 1. / img_width
    dh = 1. / img_height
    x = bbox[0] + bbox[2] / 2.0
    y = bbox[1] + bbox[3] / 2.0
    w = bbox[2]
    h = bbox[3]

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh

    return (x, y, w, h)


# save image and label
def save_image_label(image_path_copy_to, label_path_copy_to, label_text):
    # 保存地址和复制图片
    shutil.copyfile(img_path, image_path_copy_to)
    # 保存val label
    with open(label_path_copy_to, "w") as f:
        index = 0
        for write in label_text:
            label_info = [str(write[0]), str(write[1][0]), str(write[1][1]), str(write[1][2]), str(write[1][3])]
            if index == 0:
                f.write(" ".join(label_info))
                index += 1
            else:
                f.write("\n" + " ".join(label_info))

# 以下两个是coco数据集的路径，按需修改！！！！！
coco_image_path = 'D:\CMP\\2021Innovation\dataset\Cowboy\cowboyoutfits\images'
coco_json_path = 'D:\CMP\\2021Innovation\dataset\Cowboy\cowboyoutfits\\train.json'
# 拼接出保存的目录结构
save_path_root = './yolo_datasets'
save_path_train = os.path.join(save_path_root, 'train')
save_path_train_images = os.path.join(save_path_train, 'images')
save_path_train_labels = os.path.join(save_path_train, 'labels')
save_path_val = os.path.join(save_path_root, 'val')
save_path_val_images = os.path.join(save_path_val, 'images')
save_path_val_labels = os.path.join(save_path_val, 'labels')
save_path = [save_path_root, save_path_train, save_path_train_images, save_path_train_labels, save_path_val,
             save_path_val_images, save_path_val_labels]
# 生成保存目录
for path in save_path:
    check_path(path)
# 用于后续统计及划分数据集
categories = {87: 'belt', 1034: 'sunglasses', 131: 'boot', 318: 'cowboy_hat', 588: 'jacket'}
categories_name = ['belt', 'sunglasses', 'boot', 'cowboy_hat', 'jacket']
categories_map = {'belt': 0, 'sunglasses': 1, 'boot': 2, 'cowboy_hat': 3, 'jacket': 4}  # 类别映射
categories_static = {'belt': 0, 'sunglasses': 0, 'boot': 0, 'cowboy_hat': 0,
                     'jacket': 0}  # 统计验证集
# 载入数据
train_data = CocoDetection(coco_image_path, coco_json_path)
# 获取id和filename的对应字典
id_filenames = {}
for i in range(len(train_data.coco.dataset['images'])):
    id = train_data.coco.dataset['images'][i]['id']
    file_name = train_data.coco.dataset['images'][i]['file_name']
    id_filenames.update({id: file_name})

# 给每个种类划分10个给验证集
sample_num = 10
# 统计数据

for image, label in tqdm(train_data):
    # 将PIL图片转为cv
    image = cv.cvtColor(np.array(image), cv.COLOR_RGB2BGR)
    # 通过id获取图片名称
    image_name = id_filenames[label[0]['image_id']]
    # 去除.jpg后最
    image_name = image_name[:-4]
    # 获取图片路径
    img_path = os.path.join(coco_image_path, image_name + '.jpg')
    # 获取图片长宽
    img_width, img_height = image.shape[1], image.shape[0]
    # 保存的信息
    info = []
    for i in range(len(label)):
        # 读取bbox并转换为yolo格式
        bbox = label[i]['bbox']
        # 用于display用
        # x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])
        bbox = cc2yolo_bbox(img_width, img_height, bbox)
        # 获取对应label的class名字
        class_name = categories[label[i]['category_id']]
        info.append([categories_map[class_name], bbox])

    #     cv.rectangle(image, (x1, y1), (x2, y2), color=(255, 0, 0), thickness=2)
    #     cv.putText(image, class_name, (x1, y1), cv.FONT_HERSHEY_SIMPLEX, 0.5, color=(0, 0, 255), thickness=2)
    # cv.imshow('image', image)
    # cv.waitKey()

    class_counter = [0, 0, 0, 0, 0]
    label_text = []
    flag = 1
    for text in info:
        class_counter[text[0]] += 1
        label_text.append(text)
        # 是腰带类
        if text[0] == 0:
            class_index = 0
            flag = 0
    # print(label_text)
    if flag:
        class_index = np.argmax(class_counter)
    # 如果少于10个，就加入测试集
    if categories_static[categories_name[class_index]] < 10:
        # 对应类别数加一
        categories_static[categories_name[class_index]] += 1
        # 保存到val
        image_path_copy_to = os.path.join(save_path_val_images, image_name + '.jpg')
        label_path_copy_to = os.path.join(save_path_val_labels, image_name + '.txt')
        save_image_label(image_path_copy_to, label_path_copy_to, label_text)
    # 否则加入训练集
    else:
        image_path_copy_to = os.path.join(save_path_train_images, image_name + '.jpg')
        label_path_copy_to = os.path.join(save_path_train_labels, image_name + '.txt')
        save_image_label(image_path_copy_to, label_path_copy_to, label_text)

运行后，可以看到在当前目录生成了yolo格式的数据集

四、以YOLOV5为baseline进行训练

这里不知道怎么训练的可以看看我这篇文章，或者CSDN上也有很多人讲怎么训练的

Pytorch机器学习（四）——YOLOV5训练自己的VOC数据集

我这里以YOLOV5S为网络结构去训练的50个epoch，得到结果如下。

其实这个结果也不意外，腰带的检测效果十分的差，主要就是样本不均导致的结果，这个是后期重点需要解决的，也是这个比赛的难点。

     Epoch   gpu_mem       box       obj       cls    labels  img_size
     49/49     3.88G    0.0274   0.01624 0.0008392        19       640: 100%|██████████| 189/189 [06:54<00:00,  2.19s/it]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 2/2 [00:03<00:00,  1.87s/it]
                 all         50         85      0.825      0.512      0.566       0.36
                belt         50         12          1          0      0.167     0.0417
          sunglasses         50         14      0.667      0.857      0.757      0.514
                boot         50         21      0.801      0.381      0.515      0.259
          cowboy_hat         50         13      0.883      0.769       0.81      0.595
              jacket         50         25      0.776      0.554      0.582       0.39

五、生成val测试集并生成比赛提交文件

训练好后，我们需要使用比赛提供的测试集，去测试效果。

生成val文件，我也为大家写好了，直接改一下目录路径运行即可

import csv
import shutil
import os
# 修改为自己test.csv的路径和图片路径
test_path = './cowboyoutfits/test.csv'
root = './cowboyoutfits/images'
# 后面不需要修改
val_root = 'val'
if not os.path.exists(val_root):
    os.mkdir(val_root)

index = 0
with open(test_path, 'r') as f:
    reader = csv.reader(f)
    for row in reader:
        if index == 0:
            index = 1
            continue

        filename = row[1]
        image_path = os.path.join(root, filename)
        image_copy_to = os.path.join(val_root, filename)
        print(image_path, image_copy_to)
        shutil.copyfile(image_path, image_copy_to)

生成val文件夹后，将其中的图片复制到yolov5的data下的images文件夹下，将其中的权重参数改成训练好的，进行预测即可。

在detect下，注意要把这个save-txt的参数改为默认为True

预测完成后，我们需要将yolo生成的txt文件，转换为体积的coco格式，这里我直接是用的官方提供的一种方法。

import os
import pandas as pd
from PIL import Image
import zipfile
# 以下是需要修改的路径
test_df = pd.read_csv('D:\CMP\\2021Innovation\dataset\Cowboy\cowboyoutfits\\test.csv')  # test.csv的路径
PRED_PATH = "D:\CMP\\2021Innovation\yolov5_cowboy\\runs\detect\exp2\labels"             # 跑出结果的labels路径
IMAGE_PATH = "D:\CMP\\2021Innovation\yolov5_cowboy\data\images"                         # images路径
test_df.head()
# id 映射
cate_id_map = {87: 0, 1034: 1, 131: 2, 318: 3, 588: 4}

# list our prediction files path
prediction_files = os.listdir(PRED_PATH)
print('Number of test images with detections: ', len(prediction_files))


# convert yolo to coco annotation format
def yolo2cc_bbox(img_width, img_height, bbox):
    x = (bbox[0] - bbox[2] * 0.5) * img_width
    y = (bbox[1] - bbox[3] * 0.5) * img_height
    w = bbox[2] * img_width
    h = bbox[3] * img_height
    return (x, y, w, h)


# reverse the categories numer to the origin id
re_cate_id_map = dict(zip(cate_id_map.values(), cate_id_map.keys()))


def make_submission(df, PRED_PATH, IMAGE_PATH):
    output = []
    # for i in tqdm(range(len(df))):
    for i in range(len(df)):
        print(i)
        row = df.loc[i]
        image_id = row['id']
        file_name = row['file_name'].split('.')[0]
        if f'{file_name}.txt' in prediction_files:
            img = Image.open(f'{IMAGE_PATH}/{file_name}.jpg')
            width, height = img.size
            with open(f'{PRED_PATH}/{file_name}.txt', 'r') as file:
                for line in file:
                    # print(line)
                    preds = line.strip('\n').split(' ')

                    preds = list(map(float, preds))  # conver string to float
                    cc_bbox = yolo2cc_bbox(width, height, preds[1: ])
                    result = {
                        'image_id': image_id,
                        'category_id': re_cate_id_map[preds[0]],
                        'bbox': cc_bbox,
                        'score': preds[-1]
                    }

                    output.append(result)
    return output


sub_data = make_submission(test_df, PRED_PATH, IMAGE_PATH)
op_pd = pd.DataFrame(sub_data)
op_pd.sample(10)
if not os.path.exists('cow_answer'):
    os.mkdir('cow_answer')
op_pd.to_json('cow_answer/answer.json', orient='records')
zf = zipfile.ZipFile('cow_answer/sample_answer.zip', 'w')
zf.write('cow_answer/answer.json', 'answer.json')
zf.close()

运行脚本后，可以看到有个新的文件夹cow_answer下有

 我们只需要在以下网站提交zip文件即可。

比赛网站

最后这个baseline是33.9分，还有很多地方需要改进！

总结

欢迎讨论！