【配置 YOLOX 用于按目录分类的图片数据集】
现在的图标点选越来越多,如何一步解决,采用 YOLOX 目标检测模式则可以轻松解决
要在 YOLOX 中使用按目录分类的图片数据集(每个目录代表一个类别,目录下是该类别的所有图片),你需要进行以下配置步骤:
将需要检测的图标截取为独立的图片,放入到对应的目录中,如 banner ,apple 等
1. 准备数据集结构
假设你的原始数据集结构如下:
dataset_root/class1/img1.jpgimg2.jpg...class2/img1.jpgimg2.jpg......
2. 转换为 YOLOX 所需的格式
YOLOX 通常需要以下格式:
图片文件对应的标注文件(每个图片一个 .txt 文件)训练/验证集的划分文件
创建 Python 转换脚本
import os
import glob
from sklearn.model_selection import train_test_splitdef convert_directory_to_yolox(dataset_root, output_dir):# 获取所有类别classes = [d for d in os.listdir(dataset_root) if os.path.isdir(os.path.join(dataset_root, d))]classes.sort()# 创建类别映射文件with open(os.path.join(output_dir, 'classes.txt'), 'w') as f:f.write('\n'.join(classes))# 收集所有图片路径和标签image_paths = []labels = []for class_id, class_name in enumerate(classes):class_dir = os.path.join(dataset_root, class_name)for img_path in glob.glob(os.path.join(class_dir, '*')):if img_path.lower().endswith(('.png', '.jpg', '.jpeg')):image_paths.append(img_path)labels.append(class_id)# 划分训练集和验证集train_img, val_img, train_lbl, val_lbl = train_test_split(image_paths, labels, test_size=0.2, random_state=42)# 创建标注文件和数据集列表def write_dataset(img_list, lbl_list, list_file):with open(list_file, 'w') as f_list:for img_path, label in zip(img_list, lbl_list):# 创建标注文件路径base_name = os.path.splitext(os.path.basename(img_path))[0]txt_path = os.path.join(output_dir, 'labels', base_name + '.txt')# 写入标注文件os.makedirs(os.path.dirname(txt_path), exist_ok=True)with open(txt_path, 'w') as f_txt:# YOLO格式: class_id x_center y_center width height# 对于全图分类,bbox是整张图片f_txt.write(f"{label} 0.5 0.5 1.0 1.0\n")# 写入数据集列表f_list.write(f"{img_path}\n")os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)write_dataset(train_img, train_lbl, os.path.join(output_dir, 'train.txt'))write_dataset(val_img, val_lbl, os.path.join(output_dir, 'val.txt'))# 使用示例
convert_directory_to_yolox('path/to/your/dataset_root', 'path/to/yolox_dataset')
3. 配置 YOLOX
修改 exps 中的配置文件
找到或创建你的实验配置文件(如 exps/example/yolox_s.py),修改以下部分:
class Exp(yolox_base.Exp):def __init__(self):super(Exp, self).__init__()self.num_classes = 你的类别数量self.data_dir = "转换后的数据集路径"self.train_ann = "train.txt"self.val_ann = "val.txt"self.input_size = (640, 640) # 根据你的需求调整self.test_size = (640, 640)
修改数据加载部分
确保你的数据加载器能够处理这种格式。YOLOX 默认使用 MosaicDetection 数据集,你可能需要自定义:
from yolox.data import get_yolox_datadir
from yolox.data.datasets import DATASET_REGISTRY@DATASET_REGISTRY.register()
class YourDataset(Dataset):def __init__(self, ...):# 实现你的数据集加载逻辑pass
4. 训练命令
使用修改后的配置运行训练:
python tools/train.py -f exps/example/yolox_s.py -d 1 -b 64 --fp16 -o -c yolox_s.pth
注意事项
这种方法将整个图像作为一个"边界框"处理,适用于图像分类任务如果你需要真正的目标检测(图像中有多个对象),需要为每个对象提供精确的边界框标注考虑调整输入尺寸以适应你的图像比例对于分类任务,可能需要修改损失函数为更适合分类的损失(如交叉熵)
如果你需要真正的目标检测而不是图像分类,则需要为每个图像提供包含对象位置和类别的详细标注文件。