近期使用mmyolo过程中发现工具自带的yolo2coco.py在转换完数据集格式后,可视化标签的时候会有标签错乱情况,具体原因也没找到,肯定是转换过程代码有问题,于是重新做一份代码直接从yolo数据集转化为coco的json格式。
代码如下:
只需要维护
category_map
data[“categories”]
这两个标签id和标签名即可。
我的txt标签格式如下
标签类别为:
{“id”: 0, “name”: “person”},
{“id”: 1, “name”: “robotic-arm”}
import os
import json
from PIL import Image
# 定义类别映射
category_map = {
"0": 0,
"1": 1
}
def txt_to_json(input_dir, output_json):
# 初始化输出数据结构
data = {
"images": [],
"annotations": [],
"categories": [
{"id": 0, "name": "person"},
{"id": 1, "name": "robotic-arm"}
]
}
# 用于生成唯一的 image_id 和 annotation_id
image_id = 1
annotation_id = 1
# 动态读取图像尺寸
def get_image_size(image_path):
with Image.open(image_path) as img:
return img.width, img.height
# 遍历输入目录中的所有文件
for filename in os.listdir(input_dir):
if filename.endswith('.txt'):
# 提取图像文件名
image_name = filename.replace('.txt', '.jpg')
image_path = os.path.join(input_dir, image_name)
# 动态读取图像尺寸
if os.path.exists(image_path):
image_width, image_height = get_image_size(image_path)
else:
print(f"警告: 图像文件 {image_name} 不存在,跳过文件 {filename}")
continue
image_info = {
"id": image_id,
"width": image_width,
"height": image_height,
"file_name": image_name
}
data['images'].append(image_info)
# 读取文本文件
with open(os.path.join(input_dir, filename), 'r') as file:
lines = file.readlines()
# 遍历文件中的每一行
for line in lines:
parts = line.strip().split()
if len(parts) != 5:
print(f"警告: 文件 {filename} 中的行格式错误: {line.strip()}")
continue # 跳过格式错误的行
category_id = int(parts[0])
x_center = float(parts[1]) * image_width
y_center = float(parts[2]) * image_height
bbox_width = float(parts[3]) * image_width
bbox_height = float(parts[4]) * image_height
x_min = int(x_center - bbox_width / 2)
y_min = int(y_center - bbox_height / 2)
bbox = [x_min, y_min, bbox_width, bbox_height]
area = bbox_width * bbox_height
annotation_info = {
"id": annotation_id,
"image_id": image_id,
"category_id": category_id,
"bbox": bbox,
"area": area,
"iscrowd": 0
}
data['annotations'].append(annotation_info)
annotation_id += 1
image_id += 1
# 保存到 JSON 文件
with open(output_json, 'w') as json_file:
json.dump(data, json_file, indent=4)
print(f"转换完成,输出文件保存为 {output_json}")
# 调用函数
input_dir = './personDataset/train'
output_json = './trainval.json'
txt_to_json(input_dir, output_json)
input_dir = './personDataset/test'
output_json = './test.json'
txt_to_json(input_dir, output_json)
结果:
![](https://i-blog.csdnimg.cn/direct/7e18c88f355340d5aaae4d466bda827e.png