yolo数据集格式

dataset_yolo
    images
        |--train
        |--test
        |--val
    labels
        |--train
        |--test
        |--val

yolo2coco.py

from genericpath import exists
import os
import json
from PIL import Image# 设置数据集路径
dataset_path = "./dataset_1223_yolo"
images_path = os.path.join(dataset_path, "images")
labels_path = os.path.join(dataset_path, "labels")save_dir="output"
os.makedirs(save_dir,exist_ok=True)# 类别映射
categories = [{"id": 1, "name": "yb_text"},{"id": 2, "name": "kk_text"},{"id": 3, "name": "zsd_text"},{"id": 4, "name": "xn_text"},{"id": 5, "name": "controls_text"},{"id": 6, "name": "water_mark"},# 添加更多类别
]# YOLO格式转COCO格式的函数
def convert_yolo_to_coco(x_center, y_center, width, height, img_width, img_height):x_min = (x_center - width / 2) * img_widthy_min = (y_center - height / 2) * img_heightwidth = width * img_widthheight = height * img_heightreturn [x_min, y_min, width, height]# 初始化COCO数据结构
def init_coco_format():return {"images": [],"annotations": [],"categories": categories}suffixes=('.png', '.jpg')# 处理每个数据集分区
for split in ['train', 'test', 'val']:coco_format = init_coco_format()annotation_id = 1folder_path=os.path.join(images_path, split)if not os.path.exists(folder_path):continue for img_name in os.listdir(folder_path):if img_name.lower().endswith(suffixes):img_path = os.path.join(images_path, split, img_name)label_path = os.path.join(labels_path, split, os.path.splitext(img_name)[0]+".txt")img = Image.open(img_path)img_width, img_height = img.sizeimage_info = {"file_name": img_name,"id": len(coco_format["images"]) + 1,"width": img_width,"height": img_height}coco_format["images"].append(image_info)if os.path.exists(label_path):with open(label_path, "r") as file:for line in file:category_id, x_center, y_center, width, height = map(float, line.split())bbox = convert_yolo_to_coco(x_center, y_center, width, height, img_width, img_height)annotation = {"id": annotation_id,"image_id": image_info["id"],"category_id": int(category_id) + 1,"bbox": bbox,"area": bbox[2] * bbox[3],"iscrowd": 0}coco_format["annotations"].append(annotation)annotation_id += 1# 为每个分区保存JSON文件with open(f"{save_dir}/{split}_coco_format.json", "w") as json_file:json.dump(coco_format, json_file, indent=4)