### 将 YOLO 格式数据集转换为 COCO 格式
要实现从 YOLO 格式的标注文件到 COCO 格式的转换,需要完成以下几个核心任务:
1. **解析 YOLO 的 `.txt` 文件**:YOLO 标签是以相对坐标的形式存储的,每行表示一个目标框的信息。其格式如下:
```
<class_id> <x_center> <y_center> <width> <height>
```
这些值均为相对于图像尺寸的比例。
2. **构建 COCO 数据结构**:COCO 格式的标注是一个 JSON 文件,其中包含了 `images`, `categories`, 和 `annotations` 三个主要部分。具体字段定义可以参考官方文档[^1]。
3. **编写脚本进行转换**:通过 Python 脚本读取 YOLO 数据集中所有的图片和标签文件,并将其重新组织成符合 COCO 格式的 JSON 结构。
以下是具体的代码示例以及说明:
#### 转换逻辑与代码实现
```python
import os
import json
from PIL import Image
def yolo_to_coco(yolo_dir, output_json_path):
images_dir = os.path.join(yolo_dir, 'images')
labels_dir = os.path.join(yolo_dir, 'labels')
categories = [
{"id": 1, "name": "object_1"},
{"id": 2, "name": "object_2"}
] # 替换为目标类别列表
coco_data = {
"info": {},
"licenses": [],
"categories": categories,
"images": [],
"annotations": []
}
annotation_id = 1
image_id = 1
for subset in ['train', 'val', 'test']:
subset_images_dir = os.path.join(images_dir, subset)
subset_labels_dir = os.path.join(labels_dir, subset)
if not os.path.exists(subset_images_dir) or not os.path.exists(subset_labels_dir):
continue
for label_file in os.listdir(subset_labels_dir):
if not label_file.endswith('.txt'):
continue
img_name = label_file.replace('.txt', '.jpg') # 假设图片扩展名为 .jpg
img_path = os.path.join(subset_images_dir, img_name)
if not os.path.isfile(img_path):
print(f"Image {img_name} missing.")
continue
with open(os.path.join(subset_labels_dir, label_file), 'r') as f:
lines = f.readlines()
img = Image.open(img_path)
width, height = img.size
image_info = {
"file_name": img_name,
"height": height,
"width": width,
"id": image_id
}
coco_data["images"].append(image_info)
for line in lines:
parts = list(map(float, line.strip().split()))
class_id = int(parts[0])
x_center_rel, y_center_rel, w_rel, h_rel = parts[1:]
x_center_abs = x_center_rel * width
y_center_abs = y_center_rel * height
box_width = w_rel * width
box_height = h_rel * height
x_min = round(x_center_abs - (box_width / 2))
y_min = round(y_center_abs - (box_height / 2))
bbox = [x_min, y_min, box_width, box_height]
annotation = {
"segmentation": [[]],
"area": box_width * box_height,
"iscrowd": 0,
"image_id": image_id,
"bbox": bbox,
"category_id": class_id + 1, # COCO 类别 ID 应该从 1 开始
"id": annotation_id
}
coco_data["annotations"].append(annotation)
annotation_id += 1
image_id += 1
with open(output_json_path, 'w') as f:
json.dump(coco_data, f)
if __name__ == "__main__":
yolo_dataset_directory = './dataset-yolo' # 输入 YOLO 数据集路径
output_coco_json = './output-coco.json' # 输出 COCO JSON 文件路径
yolo_to_coco(yolo_dataset_directory, output_coco_json)
```
上述代码实现了以下功能:
- 遍历 YOLO 数据集中的所有子目录(如 train、val、test),逐一处理每个图片及其对应的标签文件。
- 对于每一个标签文件,提取边界框信息并计算绝对像素位置。
- 构建 COCO 所需的 JSON 数据结构,并最终保存至指定路径。
#### 关键点解释
- 图片大小获取:使用 Pillow (`PIL.Image`) 来打开图片文件以获得宽度和高度。
- 边界框转换:YOLO 使用的是中心点加宽高的比例形式,而 COCO 使用左上角顶点加上宽高作为矩形描述方式。因此需要将两者之间的坐标系相互映射[^2]。
- 类别匹配:确保 YOLO 中使用的类索引能够正确对应到 COCO 的分类体系中去。
---
###