|
|
import json |
|
|
import os |
|
|
from uuid import uuid4 |
|
|
|
|
|
|
|
|
def export_one_prediction_to_label_studio(prediction, label_studio_relative_path: str, output_folder: str, |
|
|
model_name: str, from_name: str, |
|
|
threshold_to_export: float = 0.1): |
|
|
image_name = label_studio_relative_path.split('/')[-1] |
|
|
output_file = image_name.replace('.jpg', ".json").replace('.jpeg', ".json") |
|
|
output_path = os.path.join(output_folder, output_file) |
|
|
|
|
|
original_width, original_height = prediction.orig_shape |
|
|
boxes_filtered = filter(lambda x: float(x.conf[0]) > threshold_to_export, prediction.boxes) |
|
|
json_content = { |
|
|
"data": { |
|
|
"img": label_studio_relative_path, |
|
|
}, |
|
|
"predictions": [ |
|
|
{ |
|
|
"model_version": model_name, |
|
|
"result": [ |
|
|
{ |
|
|
"original_width": original_width, |
|
|
"original_height": original_height, |
|
|
"score": float(box.conf[0]), |
|
|
"value": { |
|
|
|
|
|
"x": float(box.xywhn[0][0] - box.xywhn[0][2] / 2) * 100, |
|
|
"y": float(box.xywhn[0][1] - box.xywhn[0][3] / 2) * 100, |
|
|
"width": float(box.xywhn[0][2]) * 100, |
|
|
"height": float(box.xywhn[0][3]) * 100, |
|
|
"rotation": 0 |
|
|
}, |
|
|
"id": str(uuid4()), |
|
|
"from_name": from_name, |
|
|
"to_name": "image", |
|
|
"type": "rectangle" |
|
|
} for box in boxes_filtered], |
|
|
}] |
|
|
} |
|
|
|
|
|
with open(output_path, 'w') as fp: |
|
|
json.dump(json_content, fp) |
|
|
|
|
|
return json_content |
|
|
|
|
|
|
|
|
def export_one_yolo_output_to_label_studio_rectangles(yolo_prediction, |
|
|
model_version: str, |
|
|
label_studio_relative_path: str, |
|
|
output_folder: str, |
|
|
label_studio_to_name: str = 'image', |
|
|
label_studio_from_name: str = 'label', |
|
|
confidence_threshold: float = 0.1, |
|
|
overwrite_existing_files: bool = False): |
|
|
os.makedirs(output_folder, exist_ok=True) |
|
|
|
|
|
image_name = label_studio_relative_path.split('/')[-1] |
|
|
output_file = image_name.replace('.jpg', ".json").replace('.jpeg', ".json") |
|
|
output_path = os.path.join(output_folder, output_file) |
|
|
|
|
|
original_width, original_height = yolo_prediction.orig_shape |
|
|
class_names = yolo_prediction.names |
|
|
boxes_filtered = filter(lambda x: float(x.conf[0]) >= confidence_threshold |
|
|
, yolo_prediction.boxes) |
|
|
|
|
|
if not os.path.exists(output_path) or overwrite_existing_files: |
|
|
json_content = { |
|
|
"data": { |
|
|
"image": label_studio_relative_path, |
|
|
}, |
|
|
"predictions": [ |
|
|
{ |
|
|
"model_version": model_version, |
|
|
"result": [ |
|
|
{ |
|
|
"original_width": original_width, |
|
|
"original_height": original_height, |
|
|
"score": float(box.conf[0]), |
|
|
"value": { |
|
|
|
|
|
|
|
|
"x": float(box.xywhn[0][0] - box.xywhn[0][2] / 2) * 100, |
|
|
"y": float(box.xywhn[0][1] - box.xywhn[0][3] / 2) * 100, |
|
|
"width": float(box.xywhn[0][2]) * 100, |
|
|
"height": float(box.xywhn[0][3]) * 100, |
|
|
"rotation": 0, |
|
|
"rectanglelabels": [class_names[int(box.cls)]] |
|
|
}, |
|
|
"id": str(uuid4()), |
|
|
"from_name": label_studio_from_name, |
|
|
"to_name": label_studio_to_name, |
|
|
"type": "rectanglelabels", |
|
|
} for box in boxes_filtered], |
|
|
}] |
|
|
} |
|
|
|
|
|
with open(output_path, 'w') as fp: |
|
|
json.dump(json_content, fp) |
|
|
|
|
|
return json_content |
|
|
|
