CUB_200_2011数据集转Yolo格式
CUB200数据集是细粒度图像识别领域的基准数据集，该数据集共有11788张鸟类图像，包含200类鸟类子类，其中训练数据集有5994张图像，测试集有5794张图像。每张图像均提供了图像类标记信息，图像中鸟的bounding box，鸟的关键part信息，以及鸟类的属性信息。
所以CUB数据集也能用于鸟类目标检测模型训练，可以使用如下代码将原始CUB数据集格式转换为Yolo格式：
注：
转成的yolo格式数据集用于粗粒度bird类目标检测，仅有1种label：0.
若要实现细粒度的200种label，请在__getitem__( )返回值中添加label，并修改64行format的第一个变量为具体的细粒度label.
代码块
Python
自动换行
复制代码
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
@Time: 2023/2/19 10:50
@Author: rumi_summer
@Description: transfer dataset from CUB form to Yolo form
'''
import os
import cv2
class Manager:
def __init__(self, path, target, train=True):
self.root = path
self.target = target
self.is_train = train
self.images_path = {}
self.phase_name = 'train' if train else 'val'
with open(os.path.join(self.root, 'images.txt')) as f:
for line in f:
image_id, path = line.split()
self.images_path[image_id] = path
# 获取类别标签dict
self.class_ids = {}
with open(os.path.join(self.root, 'image_class_labels.txt')) as f:
for line in f:
image_id, class_id = line.split()
self.class_ids[image_id] = class_id
# 获取标注框
self.bondingbox = {}
with open(os.path.join(self.root, 'bounding_boxes.txt')) as f:
for line in f:
image_id, x, y, w, h = line.split()
x, y, w, h = float(x), float(y), float(w), float(h)
self.bondingbox[image_id] = (x, y, w, h)
# 获取train/test数据id列表
self.data_id = []
if self.is_train:
with open(os.path.join(self.root, 'train_test_split.txt')) as f:
for line in f:
image_id, is_train = line.split()
if int(is_train):
self.data_id.append(image_id)
if not self.is_train:
with open(os.path.join(self.root, 'train_test_split.txt')) as f:
for line in f:
image_id, is_train = line.split()
if not int(is_train):
self.data_id.append(image_id)
def transfer(self):
# 清空 train.txt / val.txt
path_file_clear = open(os.path.join(self.target, (self.phase_name + '.txt')), 'w')
path_file_clear.close()
for i in range(self.__len__()):
image, file_name, x, y, w, h = self.__getitem__(i)
file_path = os.path.join(self.target, 'images', self.phase_name, file_name)
cv2.imwrite(file_path, image)
label_file = open(os.path.join(self.target, 'labels', self.phase_name, (os.path.splitext(file_name)[0] +
'.txt')), 'w')
label_file.write('{} {} {} {} {}'.format(0, x, y, w, h))
label_file.close()
path_file = open(os.path.join(self.target, (self.phase_name + '.txt')), 'a')
path_file.write('{}{}'.format(('\n' if i != 0 else ''), file_path))
path_file.close()
def __len__(self):
return len(self.data_id)
def __getitem__(self, index):
image_id = self.data_id[index]
label = int(self._get_class_by_id(image_id)) - 1
path = self._get_path_by_id(image_id)
file_name = os.path.basename(path)
image = cv2.imread(os.path.join(self.root, 'images', path))
width = image.shape[1]
height = image.shape[0]
x, y, w, h = self.bondingbox[image_id]
x = (x + (w / 2)) / width
w /= width
y = (y + (h / 2)) / height
h /= height
return image, file_name, x, y, w, h
def _get_path_by_id(self, image_id):
return self.images_path[image_id]
def _get_class_by_id(self, image_id):
return self.class_ids[image_id]
def _get_bbox_by_id(self, image_id):
return self.bondingbox[image_id]
def mkdirs(path):
if not os.path.exists(path):
os.makedirs(path)
for phase in ['train', 'val']:
images_path = os.path.join(path, 'images', phase)
labels_path = os.path.join(path, 'labels', phase)
if not os.path.exists(images_path):
os.makedirs(images_path)
if not os.path.exists(labels_path):
os.makedirs(labels_path)
if __name__ == "__main__":
# 原数据集根目录
root = r'../xxx/xxx/CUB_200_2011'
# 目标yolo格式数据集根目录 - 须是完整绝对路径
des = r'D:\xxx\xxx\data_yolo\CUB_200_2011'
# 创建目标路径下的文件夹
mkdirs(des)
manager_train = Manager(root, des, train=True)
manager_train.transfer()
manager_val = Manager(root, des, train=False)
manager_val.transfer()
复制成功
需要修改的是root和des的路径，其中des为绝对路径。
参考资料：
opencv读取CUB数据集：https://blog.csdn.net/weixin_41735859/article/details/106937174
处理CUB数据集标注框：https://blog.csdn.net/rocketeerLi/article/details/104931869
cv21940047
分享至
投诉或建议
