目标检测之牛仔行头检测（上）—— 读取coco数据集，可视化处理，转换为yolo数据格式
目录
目标检测之牛仔行头检测（上）—— 读取coco数据集，可视化处理，转换为yolo数据格式
前言
本系列是记录参加沐神举办的牛仔行头检测比赛，具体规则和内容看下面视频
下载地址也是网友提供的，下载完后把后缀改成zip即可
这篇文章主要是提供一个引导参加这个比赛的思路，没有过多技巧，先参与动手为主，也欢迎大家和我一起学习讨论。
一、数据集格式
下载完后，数据集的大体目录长这个样子，
其中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 [email protected] [email protected]:.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分，还有很多地方需要改进！
总结
欢迎讨论！