import os
import json
import xml.etree.ElementTree as ET
def xml_to_coco(xml_folder, output_json):
coco_data = {
"images": [],
"annotations": [],
"categories": [],
}
# Mapping between class names and category ids
class_to_category_id = {}
image_id = 1
annotation_id = 1
for xml_file in os.listdir(xml_folder):
if xml_file.endswith(".xml"):
xml_path = os.path.join(xml_folder, xml_file)
tree = ET.parse(xml_path)
root = tree.getroot()
# Extract image information
image_info = {
"id": image_id,
"file_name": root.find("filename").text,
"width": int(root.find("size/width").text),
"height": int(root.find("size/height").text),
}
coco_data["images"].append(image_info)
# Extract annotations
for obj in root.findall("object"):
category_name = obj.find("name").text
if category_name not in class_to_category_id:
class_to_category_id[category_name] = len(class_to_category_id)
category_info = {
"id": class_to_category_id[category_name],
"name": category_name,
"supercategory": "object",
}
coco_data["categories"].append(category_info)
category_id = class_to_category_id[category_name]
bbox = [
float(obj.find("bndbox/xmin").text),
float(obj.find("bndbox/ymin").text),
float(obj.find("bndbox/xmax").text) - float(obj.find("bndbox/xmin").text),
float(obj.find("bndbox/ymax").text) - float(obj.find("bndbox/ymin").text),
]
area = bbox[2] * bbox[3]
annotation_info = {
"id": annotation_id,
"image_id": image_id,
"category_id": category_id,
"bbox": bbox,
"area": area,
"iscrowd": 0,
}
coco_data["annotations"].append(annotation_info)
annotation_id += 1
image_id += 1
print('coco_data')
# with open(output_json, "w") as json_file:
# json.dump(coco_data, json_file)
# Example usage
xml_folder_path = ###your xml folder
output_json_path = ##your destination file to save json
xml_to_coco(xml_folder_path, output_json_path)
目标检测数据集转换--labelme格式转coco格式通用代码
最新推荐文章于 2024-05-15 16:38:28 发布