1.YOLO格式文件组织
这是要转化的YOLO格式文件组织,分别有images和labels文件,里面划分了train和val。
2.代码实现
将其转化为coco格式。
import os
import cv2
import json
import argparse
from tqdm import tqdmCOCO_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, Nonereturn 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 loadimages_dict.append({'file_name': image_name,'height': height,'width': width,'id': start_image_id + i})return images_dictdef 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_idprint("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 filewith open(label_file, 'r') as f:lines = f.readlines()for line in lines:parts = line.strip().split()if len(parts) != 5:continue # Skip invalid annotationcategory_id = int(parts[0]) + 1 # YOLO starts from 0, COCO from 1x, y, w, h = map(float, parts[1:])x_min = (x - w / 2) * widthy_min = (y - h / 2) * heightbbox_width = w * widthbbox_height = h * heightarea = bbox_width * bbox_heightannotations_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 += 1return {'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.")continueprint(f"Processing {subset} dataset...")coco_data = convert_yolo_to_coco(image_dir, label_dir)save_json(coco_data, save_path)
执行完代码后分别有train和val的json文件,如下。
将其组织成coco数据集格式,即可。