From 4dad78a3528ccc2282e63f714868fa37f8622a97 Mon Sep 17 00:00:00 2001 From: Samir Mohan Date: Wed, 1 May 2024 19:55:57 -0400 Subject: [PATCH] Increase column width allowance for YOLO label outputs to accomodate prediction confidence scores --- label_studio_converter/imports/yolo.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/label_studio_converter/imports/yolo.py b/label_studio_converter/imports/yolo.py index 66d3e06..eb66ae0 100644 --- a/label_studio_converter/imports/yolo.py +++ b/label_studio_converter/imports/yolo.py @@ -111,7 +111,9 @@ def convert_yolo_to_ls( # convert all bounding boxes to Label Studio Results lines = file.readlines() for line in lines: - label_id, x, y, width, height = line.split() + values = line.split() + label_id, x, y, width, height = values[0:5] + score = float(values[5]) if len(values) >= 6 else None x, y, width, height = ( float(x), float(y), @@ -135,6 +137,8 @@ def convert_yolo_to_ls( "original_width": image_width, "original_height": image_height, } + if score: + item["score"] = score task[out_type][0]['result'].append(item) tasks.append(task)