281 lines
8.7 KiB
Python
281 lines
8.7 KiB
Python
"""
|
||
COCO到YOLO数据集转换工具
|
||
|
||
功能说明:
|
||
将COCO格式的数字标注数据集转换为YOLO训练所需的格式。
|
||
COCO格式使用JSON存储标注，YOLO格式使用文本文件存储。
|
||
|
||
主要功能:
|
||
- 解析COCO格式的标注文件（coco.json）
|
||
- 提取图片和边界框信息
|
||
- 转换边界框格式：COCO [x,y,w,h] → YOLO [x_center,y_center,w,h] (归一化)
|
||
- 自动划分训练集和验证集
|
||
- 创建YOLO标准目录结构
|
||
- 生成dataset.yaml配置文件
|
||
|
||
格式转换详解:
|
||
COCO格式:
|
||
- bbox: [x, y, width, height] (像素坐标，左上角)
|
||
- 绝对坐标，单位为像素
|
||
|
||
YOLO格式:
|
||
- bbox: [x_center, y_center, width, height] (归一化坐标)
|
||
- 相对坐标，值在0-1之间
|
||
- x_center = (x + width/2) / img_width
|
||
- y_center = (y + height/2) / img_height
|
||
|
||
目录结构:
|
||
输入（COCO格式）:
|
||
digit-validation/
|
||
├── coco.json # 标注文件
|
||
└── images/ # 图片文件
|
||
|
||
输出（YOLO格式）:
|
||
yolo_dataset/
|
||
├── dataset.yaml # 配置文件
|
||
├── images/
|
||
│ ├── train/ # 训练集图片
|
||
│ └── val/ # 验证集图片
|
||
└── labels/
|
||
├── train/ # 训练集标注(.txt)
|
||
└── val/ # 验证集标注(.txt)
|
||
|
||
数据划分:
|
||
- 使用固定随机种子保证可重复性
|
||
- 默认20%作为验证集
|
||
- 保持图片和标注的对应关系
|
||
|
||
使用示例:
|
||
# 基础使用（默认参数）
|
||
python scripts/prepare_yolo_dataset.py
|
||
|
||
# 自定义参数
|
||
python scripts/prepare_yolo_dataset.py \
|
||
--root digit-validation \
|
||
--out yolo_dataset \
|
||
--val-ratio 0.2 \
|
||
--seed 42
|
||
|
||
# 只使用训练集（不划分验证集）
|
||
python scripts/prepare_yolo_dataset.py --val-ratio 0.0
|
||
|
||
注意事项:
|
||
- 确保coco.json中的file_name与实际图片文件匹配
|
||
- 类别ID必须是0-9（数字）
|
||
- 边界框坐标不能超出图片范围
|
||
- 输出目录会被覆盖，注意备份
|
||
|
||
作者: Gavin Chan
|
||
版本: 1.0
|
||
日期: 2025-10-30
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import random
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Dict, List
|
||
|
||
|
||
@dataclass
|
||
class CocoImage:
|
||
id: int
|
||
file_name: str
|
||
width: int
|
||
height: int
|
||
|
||
|
||
@dataclass
|
||
class CocoAnnotation:
|
||
id: int
|
||
image_id: int
|
||
bbox: List[float] # x, y, width, height
|
||
property_info: str
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
"""
|
||
解析命令行参数
|
||
|
||
Returns:
|
||
argparse.Namespace: 参数对象
|
||
- root: COCO数据集根目录
|
||
- out: YOLO数据集输出目录
|
||
- val_ratio: 验证集比例（0-1）
|
||
- seed: 随机种子（保证可重复性）
|
||
"""
|
||
parser = argparse.ArgumentParser(description="Prepare YOLO dataset from digit-validation COCO json")
|
||
parser.add_argument("--root", type=Path, default=Path("digit-validation"), help="digit-validation directory")
|
||
parser.add_argument("--out", type=Path, default=Path("yolo_dataset"), help="output dataset directory")
|
||
parser.add_argument("--val-ratio", type=float, default=0.2, help="validation split ratio")
|
||
parser.add_argument("--seed", type=int, default=20240305, help="random seed")
|
||
return parser.parse_args()
|
||
|
||
|
||
def load_coco(root: Path) -> tuple[List[CocoImage], List[CocoAnnotation]]:
|
||
"""
|
||
加载并解析COCO格式的标注文件
|
||
|
||
Args:
|
||
root (Path): COCO数据集根目录，应包含coco.json文件
|
||
|
||
Returns:
|
||
tuple: 二元组 (images, annotations)
|
||
- images: 图片信息列表（CocoImage对象）
|
||
- annotations: 标注信息列表（CocoAnnotation对象）
|
||
|
||
Raises:
|
||
FileNotFoundError: 如果coco.json不存在
|
||
JSONDecodeError: 如果JSON格式错误
|
||
"""
|
||
coco_path = root / "coco.json"
|
||
if not coco_path.exists():
|
||
raise FileNotFoundError(f"COCO file not found at {coco_path}")
|
||
|
||
with coco_path.open("r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
images = [
|
||
CocoImage(id=img["id"], file_name=img["file_name"], width=img["width"], height=img["height"])
|
||
for img in data["images"]
|
||
]
|
||
annotations = [
|
||
CocoAnnotation(
|
||
id=ann["id"],
|
||
image_id=ann["image_id"],
|
||
bbox=ann["bbox"],
|
||
property_info=ann.get("property_info", "").strip(),
|
||
)
|
||
for ann in data["annotations"]
|
||
]
|
||
return images, annotations
|
||
|
||
|
||
def ensure_dirs(out_root: Path) -> Dict[str, Path]:
|
||
"""
|
||
创建YOLO数据集所需的目录结构
|
||
|
||
创建的目录:
|
||
- images/train/ 训练集图片
|
||
- images/val/ 验证集图片
|
||
- labels/train/ 训练集标注
|
||
- labels/val/ 验证集标注
|
||
|
||
Args:
|
||
out_root (Path): 输出根目录
|
||
|
||
Returns:
|
||
Dict[str, Path]: 目录路径字典
|
||
- images_train: 训练集图片目录
|
||
- images_val: 验证集图片目录
|
||
- labels_train: 训练集标注目录
|
||
- labels_val: 验证集标注目录
|
||
"""
|
||
dirs = {
|
||
"images_train": out_root / "images" / "train",
|
||
"images_val": out_root / "images" / "val",
|
||
"labels_train": out_root / "labels" / "train",
|
||
"labels_val": out_root / "labels" / "val",
|
||
}
|
||
for directory in dirs.values():
|
||
directory.mkdir(parents=True, exist_ok=True)
|
||
return dirs
|
||
|
||
|
||
def coco_to_yolo(
|
||
images: List[CocoImage],
|
||
annotations: List[CocoAnnotation],
|
||
image_dir: Path,
|
||
out_root: Path,
|
||
val_ratio: float,
|
||
seed: int,
|
||
) -> Path:
|
||
id_to_image = {image.id: image for image in images}
|
||
image_to_annotations: Dict[int, List[CocoAnnotation]] = {}
|
||
for ann in annotations:
|
||
image_to_annotations.setdefault(ann.image_id, []).append(ann)
|
||
|
||
valid_images = [img for img in images if (image_dir / img.file_name).exists()]
|
||
random.Random(seed).shuffle(valid_images)
|
||
|
||
split_idx = int(len(valid_images) * (1 - val_ratio))
|
||
train_imgs = valid_images[:split_idx]
|
||
val_imgs = valid_images[split_idx:]
|
||
|
||
dirs = ensure_dirs(out_root)
|
||
|
||
def process(image: CocoImage, split: str) -> None:
|
||
src_path = image_dir / image.file_name
|
||
dst_img_dir = dirs["images_train"] if split == "train" else dirs["images_val"]
|
||
dst_lbl_dir = dirs["labels_train"] if split == "train" else dirs["labels_val"]
|
||
|
||
dst_img_path = dst_img_dir / image.file_name
|
||
dst_img_path.write_bytes(src_path.read_bytes())
|
||
|
||
anns = image_to_annotations.get(image.id, [])
|
||
lines: List[str] = []
|
||
|
||
for ann in anns:
|
||
digit_str = ann.property_info.strip()
|
||
if not digit_str.isdigit():
|
||
continue
|
||
digit = int(digit_str)
|
||
if digit < 0 or digit > 9:
|
||
continue
|
||
|
||
x, y, w, h = ann.bbox
|
||
x_center = (x + w / 2) / image.width
|
||
y_center = (y + h / 2) / image.height
|
||
w_norm = w / image.width
|
||
h_norm = h / image.height
|
||
|
||
lines.append(f"{digit} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
|
||
|
||
dst_lbl_path = dst_lbl_dir / (image.file_name.rsplit(".", 1)[0] + ".txt")
|
||
dst_lbl_path.write_text("\n".join(lines), encoding="utf-8")
|
||
|
||
for img in train_imgs:
|
||
process(img, "train")
|
||
for img in val_imgs:
|
||
process(img, "val")
|
||
|
||
data_yaml = {
|
||
"path": str(out_root.resolve()),
|
||
"train": "images/train",
|
||
"val": "images/val",
|
||
"names": {i: str(i) for i in range(10)},
|
||
}
|
||
|
||
data_yaml_path = out_root / "dataset.yaml"
|
||
with data_yaml_path.open("w", encoding="utf-8") as f:
|
||
f.write("# auto-generated by scripts/prepare_yolo_dataset.py\n")
|
||
for key, value in data_yaml.items():
|
||
if isinstance(value, dict):
|
||
f.write(f"{key}:\n")
|
||
for k, v in value.items():
|
||
f.write(f" {k}: {v}\n")
|
||
else:
|
||
f.write(f"{key}: {value}\n")
|
||
|
||
print(f"YOLO dataset prepared at {out_root}")
|
||
print(f"Train images: {len(train_imgs)}, Val images: {len(val_imgs)}")
|
||
print(f"Data config written to: {data_yaml_path}")
|
||
return data_yaml_path
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
image_dir = args.root / "images"
|
||
if not image_dir.exists():
|
||
raise FileNotFoundError(f"Images directory not found at {image_dir}")
|
||
|
||
images, annotations = load_coco(args.root)
|
||
coco_to_yolo(images, annotations, image_dir, args.out, args.val_ratio, args.seed)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|