You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
262 lines
9.4 KiB
Python
262 lines
9.4 KiB
Python
#!/usr/bin/env python
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
"""
|
|
Given a CVAT XML, this script reads the CVAT XML and writes the
|
|
annotations in YOLO format into a given directory.
|
|
|
|
This implementation supports both interpolation tracks from video and
|
|
annotated images.
|
|
"""
|
|
|
|
import os
|
|
import argparse
|
|
import glog as log
|
|
from lxml import etree
|
|
import requests
|
|
|
|
|
|
def parse_args():
|
|
"""Parse arguments of command line"""
|
|
parser = argparse.ArgumentParser(
|
|
description='Convert CVAT XML annotations to YOLO format'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--cvat-xml', metavar='FILE', required=True,
|
|
help='input file with CVAT annotation in xml format'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--image-dir', metavar='DIRECTORY', required=False,
|
|
help='directory which contains original images'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--output-dir', metavar='DIRECTORY', required=True,
|
|
help='directory for output annotations in YOLO format'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--username', metavar='USERNAME', required=False,
|
|
help='Username from CVAT Login page, required to download images'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--password', metavar='PASSWORD', required=False,
|
|
help='Password from CVAT Login page, required to download images'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--labels', metavar='ILABELS', required=False,
|
|
help='Labels (separated by comma) to extract. Example: car,truck,motorcycle'
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def process_cvat_xml(xml_file, image_dir, output_dir,username,password,ilabels):
|
|
"""
|
|
Transforms a single XML in CVAT format to YOLO TXT files and download images when not in IMAGE_DIR
|
|
|
|
:param xml_file: CVAT format XML
|
|
:param image_dir: image directory of the dataset
|
|
:param output_dir: directory of annotations with YOLO format
|
|
:param username: Username used to login CVAT. Required to download images
|
|
:param password: Password used to login CVAT. Required to download images
|
|
:param ilabels: Comma separated ordered labels
|
|
:return:
|
|
"""
|
|
KNOWN_TAGS = {'box', 'image', 'attribute'}
|
|
|
|
if (image_dir is None):
|
|
image_dir=os.path.join(output_dir,"data/obj")
|
|
os.makedirs(image_dir, exist_ok=True)
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
cvat_xml = etree.parse(xml_file)
|
|
basename = os.path.splitext( os.path.basename( xml_file ) )[0]
|
|
current_labels = {}
|
|
traintxt = ""
|
|
|
|
if (ilabels is not None):
|
|
vlabels=ilabels.split(',')
|
|
current_labels = {label: idx for idx, label in enumerate(vlabels)}
|
|
else:
|
|
current_labels = {label.text: idx for idx, label in enumerate(cvat_xml.findall('meta/task/labels/label/name'))}
|
|
|
|
tracks= cvat_xml.findall( './/track' )
|
|
|
|
if (tracks is not None) and (len(tracks) > 0):
|
|
frames = {}
|
|
|
|
for track in tracks:
|
|
trackid = int(track.get("id"))
|
|
label = track.get("label")
|
|
boxes = track.findall( './box' )
|
|
for box in boxes:
|
|
frameid = int(box.get('frame'))
|
|
outside = int(box.get('outside'))
|
|
#occluded = int(box.get('occluded')) #currently unused
|
|
#keyframe = int(box.get('keyframe')) #currently unused
|
|
xtl = float(box.get('xtl'))
|
|
ytl = float(box.get('ytl'))
|
|
xbr = float(box.get('xbr'))
|
|
ybr = float(box.get('ybr'))
|
|
|
|
frame = frames.get( frameid, {} )
|
|
|
|
if outside == 0:
|
|
frame[ trackid ] = { 'xtl': xtl, 'ytl': ytl, 'xbr': xbr, 'ybr': ybr, 'label': label }
|
|
|
|
frames[ frameid ] = frame
|
|
|
|
width = int(cvat_xml.find('.//original_size/width').text)
|
|
height = int(cvat_xml.find('.//original_size/height').text)
|
|
|
|
taskid = int(cvat_xml.find('.//task/id').text)
|
|
|
|
urlsegment = cvat_xml.find(".//segments/segment/url").text
|
|
urlbase = urlsegment.split("?")[0]
|
|
|
|
httpclient = requests.session()
|
|
httpclient.get(urlbase)
|
|
|
|
csrftoken = "none"
|
|
sessionid = "none"
|
|
|
|
# Spit out a list of each object for each frame
|
|
for frameid in sorted(frames.keys()):
|
|
image_name = "%s_%08d.jpg" % (basename, frameid)
|
|
image_path = os.path.join(image_dir, image_name)
|
|
if not os.path.exists(image_path):
|
|
if username is None:
|
|
log.warn('{} image cannot be found. Is `{}` image directory correct?\n'.format(image_path, image_dir))
|
|
else:
|
|
log.info('{} image cannot be found. Downloading from task ID {}\n'.format(image_path, taskid))
|
|
|
|
if sessionid == "none":
|
|
if "csrftoken" in httpclient.cookies:
|
|
csrftoken = httpclient.cookies["csrftoken"]
|
|
elif "csrf" in httpclient.cookies:
|
|
csrftoken = httpclient.cookies["csrf"]
|
|
|
|
login_data = dict(username=username, password=password,
|
|
csrfmiddlewaretoken=csrftoken, next='/dashboard')
|
|
|
|
urllogin = urlbase+"/auth/login"
|
|
httpclient.post(urllogin, data=login_data,
|
|
headers=dict(Referer=urllogin))
|
|
|
|
if ("sessionid" in httpclient.cookies):
|
|
sessionid = httpclient.cookies["sessionid"]
|
|
|
|
url = urlbase+"/api/v1/tasks/"+str(taskid)+"/frames/"+ str(frameid)
|
|
|
|
req = httpclient.get(url, headers=dict(
|
|
csrftoken=csrftoken, sessionid=sessionid))
|
|
|
|
with open(image_path, 'wb') as fo:
|
|
fo.write(req.content)
|
|
print('Url saved as %s\n' % image_path)
|
|
|
|
|
|
frame = frames[frameid]
|
|
|
|
_yoloAnnotationContent=""
|
|
|
|
objids = sorted(frame.keys())
|
|
|
|
for objid in objids:
|
|
|
|
box = frame[objid]
|
|
|
|
label = box.get('label')
|
|
xmin = float(box.get('xtl'))
|
|
ymin = float(box.get('ytl'))
|
|
xmax = float(box.get('xbr'))
|
|
ymax = float(box.get('ybr'))
|
|
|
|
if not label in current_labels:
|
|
raise Exception('Unexpected label name {}'.format(label))
|
|
|
|
labelid=current_labels[label]
|
|
yolo_x= (xmin + ((xmax-xmin)/2))/width
|
|
yolo_y= (ymin + ((ymax-ymin)/2))/height
|
|
yolo_w = (xmax - xmin) / width
|
|
yolo_h = (ymax - ymin) / height
|
|
|
|
if len(_yoloAnnotationContent) != 0:
|
|
_yoloAnnotationContent += "\n"
|
|
|
|
_yoloAnnotationContent+=str(labelid)+" "+"{:.6f}".format(yolo_x) +" "+"{:.6f}".format(yolo_y) +" "+"{:.6f}".format(yolo_w) +" "+"{:.6f}".format(yolo_h)
|
|
anno_name = os.path.basename(os.path.splitext(image_name)[0] + '.txt')
|
|
anno_path = os.path.join(image_dir, anno_name)
|
|
|
|
_yoloFile = open(anno_path, "w", newline="\n")
|
|
_yoloFile.write(_yoloAnnotationContent)
|
|
_yoloFile.close()
|
|
|
|
if len(traintxt)!=0:
|
|
traintxt+="\n"
|
|
|
|
traintxt+=image_path
|
|
|
|
else:
|
|
for img_tag in cvat_xml.findall('image'):
|
|
image_name = img_tag.get('name')
|
|
width = int(img_tag.get('width'))
|
|
height = int(img_tag.get('height'))
|
|
image_path = os.path.join(image_dir, image_name)
|
|
if not os.path.exists(image_path):
|
|
log.warn('{} image cannot be found. Is `{}` image directory correct?'.
|
|
format(image_path, image_dir))
|
|
|
|
unknown_tags = {x.tag for x in img_tag.iter()}.difference(KNOWN_TAGS)
|
|
if unknown_tags:
|
|
log.warn('Ignoring tags for image {}: {}'.format(image_path, unknown_tags))
|
|
|
|
_yoloAnnotationContent = ""
|
|
|
|
for box in img_tag.findall('box'):
|
|
label = box.get('label')
|
|
xmin = float(box.get('xtl'))
|
|
ymin = float(box.get('ytl'))
|
|
xmax = float(box.get('xbr'))
|
|
ymax = float(box.get('ybr'))
|
|
|
|
if not label in current_labels:
|
|
raise Exception('Unexpected label name {}'.format(label))
|
|
|
|
labelid = current_labels[label]
|
|
yolo_x = (xmin + ((xmax-xmin)/2))/width
|
|
yolo_y = (ymin + ((ymax-ymin)/2))/height
|
|
yolo_w = (xmax - xmin) / width
|
|
yolo_h = (ymax - ymin) / height
|
|
|
|
if len(_yoloAnnotationContent) != 0:
|
|
_yoloAnnotationContent += "\n"
|
|
|
|
_yoloAnnotationContent += str(labelid)+" "+"{:.6f}".format(yolo_x) + " "+"{:.6f}".format(
|
|
yolo_y) + " "+"{:.6f}".format(yolo_w) + " "+"{:.6f}".format(yolo_h)
|
|
|
|
anno_name = os.path.basename(os.path.splitext(image_name)[0] + '.txt')
|
|
anno_path = os.path.join(image_dir, anno_name)
|
|
|
|
_yoloFile = open(anno_path, "w", newline="\n")
|
|
_yoloFile.write(_yoloAnnotationContent)
|
|
_yoloFile.close()
|
|
|
|
traintxt_file=open(output_dir+"/train.txt","w",newline="\n")
|
|
traintxt_file.write(traintxt)
|
|
traintxt_file.close()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
process_cvat_xml(args.cvat_xml, args.image_dir, args.output_dir, args.username,args.password,args.labels)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|