From 553731ed70a0f0a9164d7450f9b25344133d71b1 Mon Sep 17 00:00:00 2001 From: Andrey Zhavoronkov <41117609+azhavoro@users.noreply.github.com> Date: Tue, 27 Aug 2019 15:11:23 +0300 Subject: [PATCH] fixed label order (#669) --- utils/yolo/converter.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/utils/yolo/converter.py b/utils/yolo/converter.py index 6539eb15..4bbcde29 100644 --- a/utils/yolo/converter.py +++ b/utils/yolo/converter.py @@ -78,13 +78,12 @@ def process_cvat_xml(xml_file, image_dir, output_dir,username,password,ilabels): basename = os.path.splitext( os.path.basename( xml_file ) )[0] current_labels = {} traintxt = "" - auto_lbl_count = 0 if (ilabels is not None): vlabels=ilabels.split(',') - for _label in vlabels: - current_labels[_label]=auto_lbl_count - auto_lbl_count+=1 + 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' ) @@ -179,8 +178,7 @@ def process_cvat_xml(xml_file, image_dir, output_dir,username,password,ilabels): ymax = float(box.get('ybr')) if not label in current_labels: - current_labels[label] = auto_lbl_count - auto_lbl_count+=1 + raise Exception('Unexpected label name {}'.format(label)) labelid=current_labels[label] yolo_x= (xmin + ((xmax-xmin)/2))/width @@ -228,8 +226,7 @@ def process_cvat_xml(xml_file, image_dir, output_dir,username,password,ilabels): ymax = float(box.get('ybr')) if not label in current_labels: - current_labels[label] = auto_lbl_count - auto_lbl_count += 1 + raise Exception('Unexpected label name {}'.format(label)) labelid = current_labels[label] yolo_x = (xmin + ((xmax-xmin)/2))/width