发现网上的转换效率太低了,所以提供高效版本的转换代码
import os
import json
from tqdm import tqdm
import argparse
def parse_args():
parser = argparse.ArgumentParser(description="Convert COCO JSON to YOLO format")
parser.add_argument('--json_path', default='./train.json', type=str, help="Input: COCO format (JSON)")
parser.add_argument('--save_path', default='./train2017', type=str, help="Output directory for labels")
return parser.parse_args()
def convert(size, box):
dw = 1. / size[0]
dh = 1. / size[1]
x = box[0] + box[2] / 2.0
y = box[1] + box[3] / 2.0
w = box[2]
h = box[3]
return round(x * dw, 6), round(y * dh, 6), round(w * dw, 6), round(h * dh, 6)
def main():
args = parse_args()
json_file = args.json_path
ana_txt_save_path = args.save_path
with open(json_file, 'r') as f:
data = json.load(f)
os.makedirs(ana_txt_save_path, exist_ok=True)
id_map = {category['id']: i for i, category in enumerate(data['categories'])}
with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
for category in data['categories']:
f.write(f"{category['name']}\n")
annotations = {img_id: [] for img_id in [img['id'] for img in data['images']]}
for ann in data['annotations']:
annotations[ann['image_id']].append(ann)
list_file_path = os.path.join(ana_txt_save_path, 'train2017.txt')
with open(list_file_path, 'w') as list_file:
for img in tqdm(data['images']):
filename = img["file_name"]
img_width, img_height = img["width"], img["height"]
img_id = img["id"]
head, _ = os.path.splitext(filename)
ana_txt_name = f"{head}.txt"
with open(os.path.join(ana_txt_save_path, ana_txt_name), 'w') as f_txt:
for ann in annotations[img_id]:
box = convert((img_width, img_height), ann["bbox"])
f_txt.write(f"{id_map[ann['category_id']]} {box[0]} {box[1]} {box[2]} {box[3]}\n")
list_file.write(f'./images/train2017/{filename}\n')
if __name__ == '__main__':
main()