File size: 3,633 Bytes
fd27f5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import argparse
parser = argparse.ArgumentParser(description='Coco Json to Pascal VOC XML Converter.')
parser.add_argument('--coco_json', required=True, help='Target json file to convert to csv.')
parser.add_argument('--coco_folder', required=False, default='', help='Target folder to find images.')
parser.add_argument('--save_xml', required=True, help='The folder to save annotations xmls.')
parser.add_argument('--database_name', required=False, default='', help='The name of database.')
parser.add_argument('--no_skip_background', dest='skip_background', action='store_false', help='Do not skip \'background\' category.')
parser.set_defaults(skip_background=True)
args = parser.parse_args()
from pycocotools.coco import COCO
from PIL import Image
from pathlib import Path
import os
def write_to_xml(image_name, bboxes, image_folder_name, data_folder, save_folder, database_name):
with Image.open(os.path.join(data_folder, image_name)) as img:
width, height = img.size
if img.mode == 'YCbCr':
depth = 3
else:
depth = len(img.mode)
objects = ''
for bbox in bboxes:
objects = objects + '''
<object>
<name>{category_name}</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>{xmin}</xmin>
<ymin>{ymin}</ymin>
<xmax>{xmax}</xmax>
<ymax>{ymax}</ymax>
</bndbox>
</object>'''.format(
category_name = bbox[0],
xmin = bbox[1],
ymin = bbox[2],
xmax = bbox[3],
ymax = bbox[4]
)
xml = '''<annotation>
<folder>{image_folder_name}</folder>
<filename>{image_name}</filename>
<source>
<database>{database_name}</database>
</source>
<size>
<width>{width}</width>
<height>{height}</height>
<depth>{depth}</depth>
</size>
<segmented>0</segmented>{objects}
</annotation>'''.format(
image_folder_name = image_folder_name,
image_name = image_name,
database_name = database_name,
width = width,
height = height,
depth = depth,
objects = objects
)
anno_path = os.path.join(save_folder, os.path.splitext(image_name)[0] + '.xml')
with open(anno_path, 'w') as file:
file.write(xml)
Path(args.save_xml).mkdir(parents=True, exist_ok=True)
image_folder_name = os.path.basename(os.path.abspath(args.coco_folder))
coco = COCO(args.coco_json)
imgIds = coco.getImgIds()
if not args.database_name:
args.database_name = image_folder_name
print('Write annotations file...')
total = len(imgIds)
now = 1
for imgId in imgIds:
img = coco.loadImgs(imgId)[0]
if not os.path.isfile(os.path.join(args.coco_folder, img['file_name'])):
print('The image {} is not exist.'.format(img['file_name']))
exit()
anno_list = []
annIds = coco.getAnnIds(imgIds=imgId)
for annId in annIds:
ann = coco.loadAnns(annId)[0]
cat = coco.loadCats(ann['category_id'])[0]
if cat['name'] == 'background' and args.skip_background == True:
continue
box = ann['bbox']
anno_list.append([cat['name'], int(box[0]), int(box[1]), int(box[0]+box[2]), int(box[1]+box[3])])
write_to_xml(img['file_name'], anno_list, image_folder_name, args.coco_folder, args.save_xml, args.database_name)
print('Write xml files ({} / {})'.format(now, total))
now = now + 1
print('Annotations file was written!') |