数据集YOLO格式转换COCO格式

🧨只需要替换相应路径即可进行格式转换！！！

import os
import json
import shutil
from tqdm import tqdm
from PIL import Image# 类别映射：YOLO 中类别 ID -> 类别名
id2name = {0: 'dog', 1: 'cat', 2: 'person', 3: 'car'}
all_category_ids = sorted(id2name.keys())def convert_txt_to_coco(txt_path, image_dir, label_dir, save_image_dir, save_ann_path):image_id = 0annotation_id = 0coco_dict = {"images": [],"annotations": [],"categories": [{"id": cid, "name": id2name[cid]} for cid in all_category_ids]}with open(txt_path, 'r') as f:image_paths = [line.strip() for line in f if line.strip()]for img_path in tqdm(image_paths, desc=f'Processing {os.path.basename(txt_path)}'):file_name = os.path.basename(img_path)image_full_path = os.path.join(image_dir, file_name)label_name = os.path.splitext(file_name)[0] + ".txt"label_path = os.path.join(label_dir, label_name)# 获取图像尺寸try:with Image.open(image_full_path) as img:width, height = img.sizeexcept Exception as e:print(f"[WARN] Can't open {image_full_path}: {e}")continue# 复制图像到目标目录os.makedirs(save_image_dir, exist_ok=True)shutil.copy(image_full_path, os.path.join(save_image_dir, file_name))# 加入图像信息coco_dict["images"].append({"id": image_id,"file_name": file_name,"width": width,"height": height})# 处理标签if os.path.exists(label_path):with open(label_path, 'r') as lf:for line in lf:parts = line.strip().split()if len(parts) != 5:continueclass_id, x, y, w, h = map(float, parts)class_id = int(class_id)if class_id not in id2name:continuex_center = x * widthy_center = y * heightbox_w = w * widthbox_h = h * heightx_min = x_center - box_w / 2y_min = y_center - box_h / 2coco_dict["annotations"].append({"id": annotation_id,"image_id": image_id,"category_id": class_id,"bbox": [x_min, y_min, box_w, box_h],"area": box_w * box_h,"iscrowd": 0})annotation_id += 1image_id += 1os.makedirs(os.path.dirname(save_ann_path), exist_ok=True)with open(save_ann_path, 'w') as f:json.dump(coco_dict, f, indent=2)print(f"[INFO] Saved COCO JSON to: {save_ann_path}")if __name__ == "__main__":origin_root = "/data/origin-dataset-yolo"target_root = "/data/target-dataset-coco"image_dir = os.path.join(origin_root, "images")label_dir = os.path.join(origin_root, "labels")# 训练集convert_txt_to_coco(txt_path=os.path.join(origin_root, "train.txt"),image_dir=image_dir,label_dir=label_dir,save_image_dir=os.path.join(target_root, "images/train"),save_ann_path=os.path.join(target_root, "annotations/train.json"))# 验证集convert_txt_to_coco(txt_path=os.path.join(origin_root, "val.txt"),image_dir=image_dir,label_dir=label_dir,save_image_dir=os.path.join(target_root, "images/val"),save_ann_path=os.path.join(target_root, "annotations/val.json"))