保姆级教程:用Python脚本将OPIXray/HIXray安检X光数据集转成YOLO格式(附完整代码)
·
从安检X光到YOLO格式:OPIXray/HIXray数据集转换实战指南
当你第一次拿到OPIXray或HIXray这样的专业安检X光数据集时,可能会被其复杂的标注格式难住。作为计算机视觉领域的新手,我们往往更熟悉YOLO这种简洁的标注方式——每个物体只需一行五个数字就能描述清楚。本文将带你一步步完成从原始VOC格式到YOLO格式的完整转换,并提供可直接运行的Python脚本。
1. 理解数据集与格式差异
在开始转换前,我们需要清楚两种标注格式的本质区别。VOC格式采用XML文件存储标注信息,每个物体用四个坐标点表示其边界框(xmin, ymin, xmax, ymax),而YOLO格式则使用归一化后的中心点坐标和宽高(x_center, y_center, width, height)。
关键差异对比 :
| 特征 | VOC格式 | YOLO格式 |
|---|---|---|
| 坐标表示 | 绝对像素值 | 归一化相对值(0-1) |
| 边界框描述 | 左上+右下角坐标 | 中心点坐标+宽高 |
| 文件存储 | 每个图像对应XML文件 | 所有标注在一个txt文件 |
对于OPIXray和HIXray这两个安检专用数据集,还需要特别注意它们的类别定义不同:
# OPIXray类别字典
opixray_classes = {
'Straight_Knife': 0,
'Folding_Knife': 1,
'Scissor': 2,
'Utility_Knife': 3,
'Multi-tool_Knife': 4
}
# HIXray类别字典
hixray_classes = {
'Mobile_Phone': 0,
'Laptop': 1,
'Portable_Charger_2': 2,
'Portable_Charger_1': 3,
'Tablet': 4,
'Cosmetic': 5,
'Water': 6,
'Nonmetallic_Lighter': 7
}
2. 环境准备与脚本解析
转换过程需要Python环境和几个基础库。建议使用conda创建虚拟环境:
conda create -n xray_converter python=3.8
conda activate xray_converter
pip install opencv-python numpy
核心转换脚本包含几个关键函数:
- 坐标转换函数 :将VOC的(xmin,ymin,xmax,ymax)转换为YOLO的(x_center,y_center,width,height)
def voc_to_yolo(size, box):
"""将VOC格式坐标转换为YOLO格式
Args:
size: 图像宽高 (w,h)
box: VOC格式边界框 [xmin,ymin,xmax,ymax]
Returns:
list: YOLO格式坐标 [x_center,y_center,width,height]
"""
dw = 1./size[0] # 宽度归一化因子
dh = 1./size[1] # 高度归一化因子
x_center = (box[0] + box[2])/2.0 * dw
y_center = (box[1] + box[3])/2.0 * dh
width = (box[2] - box[0]) * dw
height = (box[3] - box[1]) * dh
return [x_center, y_center, width, height]
- 类别映射函数 :根据数据集类型返回对应的类别索引
def get_class_index(class_name, dataset_type='OPIXray'):
"""获取类别对应的索引编号
Args:
class_name: 类别名称字符串
dataset_type: 数据集类型('OPIXray'或'HIXray')
Returns:
int: 类别索引
"""
if dataset_type == 'OPIXray':
class_dict = opixray_classes
else:
class_dict = hixray_classes
return class_dict.get(class_name, -1) # 未找到返回-1
3. 完整转换流程实现
现在我们将各个功能模块整合成完整的转换流程。假设原始数据按如下结构组织:
dataset_root/
├── images/ # 存放所有X光图像
├── annotations/ # 存放VOC格式标注文件
└── labels/ # 输出YOLO格式标签(转换后自动创建)
完整转换脚本如下:
import os
import cv2
def convert_voc_to_yolo(dataset_root, dataset_type='OPIXray'):
"""主转换函数
Args:
dataset_root: 数据集根目录路径
dataset_type: 数据集类型('OPIXray'或'HIXray')
"""
# 路径配置
img_dir = os.path.join(dataset_root, 'images')
anno_dir = os.path.join(dataset_root, 'annotations')
output_dir = os.path.join(dataset_root, 'labels')
os.makedirs(output_dir, exist_ok=True)
# 处理每个标注文件
for anno_file in os.listdir(anno_dir):
if not anno_file.endswith('.txt'):
continue
anno_path = os.path.join(anno_dir, anno_file)
output_path = os.path.join(output_dir, anno_file)
with open(anno_path, 'r') as f_in, open(output_path, 'w') as f_out:
for line in f_in:
# 解析每行标注: img_name class x1 y1 x2 y2
parts = line.strip().split()
if len(parts) != 6:
continue
img_name, class_name = parts[0], parts[1]
box = list(map(float, parts[2:6]))
# 获取图像尺寸
img_path = os.path.join(img_dir, img_name)
img = cv2.imread(img_path)
if img is None:
continue
h, w = img.shape[:2]
# 坐标转换
yolo_box = voc_to_yolo((w, h), box)
class_idx = get_class_index(class_name, dataset_type)
if class_idx == -1:
continue
# 写入YOLO格式
f_out.write(f"{class_idx} {' '.join(map(str, yolo_box))}\n")
if __name__ == '__main__':
dataset_path = "path/to/your/dataset" # 修改为实际路径
convert_voc_to_yolo(dataset_path, dataset_type='OPIXray')
提示:运行前请确保所有路径设置正确,特别是图像和标注文件的路径。建议先在小批量数据上测试脚本是否正常工作。
4. 验证与可视化
转换完成后,我们需要验证结果是否正确。最简单的方法是可视化几个样本的标注:
import matplotlib.pyplot as plt
import matplotlib.patches as patches
def visualize_yolo_annotation(img_path, label_path):
"""可视化YOLO格式标注
Args:
img_path: 图像路径
label_path: 标签路径
"""
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
fig, ax = plt.subplots(1)
ax.imshow(img)
with open(label_path, 'r') as f:
for line in f:
parts = line.strip().split()
if len(parts) != 5:
continue
class_idx, xc, yc, bw, bh = map(float, parts)
# 转换回像素坐标
x = (xc - bw/2) * w
y = (yc - bh/2) * h
width = bw * w
height = bh * h
# 绘制边界框
rect = patches.Rectangle(
(x, y), width, height,
linewidth=2, edgecolor='r', facecolor='none'
)
ax.add_patch(rect)
# 添加类别标签
plt.text(
x, y-5, f'Class {int(class_idx)}',
color='red', fontsize=12, weight='bold'
)
plt.show()
# 示例使用
sample_img = "path/to/image.jpg"
sample_label = "path/to/label.txt"
visualize_yolo_annotation(sample_img, sample_label)
常见问题排查 :
- 标注框位置偏移 :检查坐标归一化计算是否正确,特别是宽高顺序
- 类别索引错误 :确认使用的类别字典与数据集匹配
- 图像加载失败 :检查图像路径和文件扩展名是否一致
- 内存不足 :处理大型数据集时,考虑分批处理
5. 高级技巧与优化建议
对于大规模数据集转换,可以考虑以下优化:
多进程处理 :利用Python的multiprocessing加速
from multiprocessing import Pool
def process_single_file(args):
"""包装单文件处理函数供多进程使用"""
file_path, output_dir, img_dir, dataset_type = args
# 实现单文件处理逻辑...
def batch_convert_voc_to_yolo(dataset_root, dataset_type, num_workers=4):
"""多进程批量转换"""
img_dir = os.path.join(dataset_root, 'images')
anno_dir = os.path.join(dataset_root, 'annotations')
output_dir = os.path.join(dataset_root, 'labels')
os.makedirs(output_dir, exist_ok=True)
file_args = [
(os.path.join(anno_dir, f), output_dir, img_dir, dataset_type)
for f in os.listdir(anno_dir) if f.endswith('.txt')
]
with Pool(num_workers) as p:
p.map(process_single_file, file_args)
日志记录 :添加详细日志帮助调试
import logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s',
filename='conversion.log'
)
def convert_voc_to_yolo_with_logging(...):
try:
# 转换逻辑...
logging.info(f"成功处理 {anno_file}")
except Exception as e:
logging.error(f"处理 {anno_file} 时出错: {str(e)}")
数据校验 :转换后自动检查数据完整性
def validate_conversion(dataset_root):
"""验证转换结果完整性"""
img_dir = os.path.join(dataset_root, 'images')
label_dir = os.path.join(dataset_root, 'labels')
missing_pairs = []
# 检查图像和标签是否一一对应
for img_name in os.listdir(img_dir):
base_name = os.path.splitext(img_name)[0]
label_name = f"{base_name}.txt"
if not os.path.exists(os.path.join(label_dir, label_name)):
missing_pairs.append((img_name, label_name))
if missing_pairs:
print(f"发现 {len(missing_pairs)} 个不匹配的图像-标签对")
for img, label in missing_pairs[:5]: # 只打印前5个示例
print(f"图像 {img} 缺少对应的标签 {label}")
else:
print("所有图像都有对应的标签文件")
在实际项目中,我遇到过因路径包含中文导致的文件读取失败问题,建议所有路径都使用英文命名。另一个常见陷阱是忘记归一化坐标,导致YOLO无法正确读取标注。
更多推荐
所有评论(0)