1.YOLO格式文件组织
这是要转化的YOLO格式文件组织,分别有images和labels文件,里面划分了train和val。
2.代码实现
将其转化为coco格式。
import os
import cv2
import json
import argparse
from tqdm import tqdm
COCO_DICT = ['images', 'annotations', 'categories']
IMAGES_DICT = ['file_name', 'height', 'width', 'id']
ANNOTATIONS_DICT = ['image_id', 'iscrowd', 'area', 'bbox', 'category_id', 'id']
CATEGORIES_DICT = ['id', 'name']
YOLO_CATEGORIES = [
'Double hexagonal column', 'Flange nut', 'Hexagon nut', 'Hexagon pillar', 'Hexagon screw',
'Hexagonal steel column', 'Horizontal bubble', 'Keybar', 'Plastic cushion pillar',
'Rectangular nut', 'Round head screw', 'Spring washer', 'T-shaped screw'
]#这里替换成自己的类别
parser = argparse.ArgumentParser(description='Convert YOLO to COCO format')
parser.add_argument('--image_path', type=str, default='/home/yu/Yolov8/ultralytics-main/mydata0/images', help='Image folder containing train and val subfolders')#这里换成images的路径,images路径下包含train和val文件夹
parser.add_argument('--annotation_path', type=str, default='/home/yu/Yolov8/ultralytics-main/mydata0/labels', help='Annotation folder containing train and val subfolders')#这里换成标签的路径,同样包含标签的train和val文件夹
parser.add_argument('--save_dir', type=str, default='/home/yu/Yolov8/ultralytics-main/mydata0', help='Directory to save COCO JSON files')
args = parser.parse_args()
def save_json(data, path):
print(f'Saving JSON: {path}')
with open(path, 'w') as f:
json.dump(data, f, indent=4)
print('Successfully saved JSON:', path)
def load_image(image_path):
img = cv2.imread(image_path)
if img is None:
print(f"Warning: Cannot load image {image_path}")
return None, None
return img.shape[0], img.shape[1]
def generate_categories_dict(category_list):
return [{CATEGORIES_DICT[0]: i + 1, CATEGORIES_DICT[1]: name} for i, name in enumerate(category_list)]
def generate_images_dict(image_list, image_dir, start_image_id=0):
images_dict = []
for i, image_name in enumerate(tqdm(image_list, desc="Processing images")):
image_path = os.path.join(image_dir, image_name)
height, width = load_image(image_path)
if height is None or width is None:
continue # Skip if image failed to load
images_dict.append({
'file_name': image_name,
'height': height,
'width': width,
'id': start_image_id + i
})
return images_dict
def convert_yolo_to_coco(image_dir, label_dir, start_image_id=0, start_annotation_id=0):
categories_dict = generate_categories_dict(YOLO_CATEGORIES)
image_list = os.listdir(image_dir)
images_dict = generate_images_dict(image_list, image_dir, start_image_id)
annotations_dict = []
annotation_id = start_annotation_id
print("Processing annotations...")
for image_data in images_dict:
image_id = image_data['id']
image_name = image_data['file_name']
width, height = image_data['width'], image_data['height']
label_file = os.path.join(label_dir, image_name.replace('.jpg', '.txt'))
if not os.path.exists(label_file):
continue # Skip if no label file
with open(label_file, 'r') as f:
lines = f.readlines()
for line in lines:
parts = line.strip().split()
if len(parts) != 5:
continue # Skip invalid annotation
category_id = int(parts[0]) + 1 # YOLO starts from 0, COCO from 1
x, y, w, h = map(float, parts[1:])
x_min = (x - w / 2) * width
y_min = (y - h / 2) * height
bbox_width = w * width
bbox_height = h * height
area = bbox_width * bbox_height
annotations_dict.append({
'image_id': image_id,
'iscrowd': 0,
'area': area,
'bbox': [x_min, y_min, bbox_width, bbox_height],
'category_id': category_id,
'id': annotation_id
})
annotation_id += 1
return {
'images': images_dict,
'annotations': annotations_dict,
'categories': categories_dict
}
if __name__ == '__main__':
for subset in ['train', 'val']:
image_dir = os.path.join(args.image_path, subset)
label_dir = os.path.join(args.annotation_path, subset)
save_path = os.path.join(args.save_dir, f'instances_{subset}2017.json')
if not os.path.exists(image_dir) or not os.path.exists(label_dir):
print(f"Skipping {subset} as folders do not exist.")
continue
print(f"Processing {subset} dataset...")
coco_data = convert_yolo_to_coco(image_dir, label_dir)
save_json(coco_data, save_path)
执行完代码后分别有train和val的json文件,如下。
将其组织成coco数据集格式,即可。