From f0a3bf89b0181f591feb7f6a401899ff4250b178 Mon Sep 17 00:00:00 2001 From: Brian Bjarke Jensen Date: Thu, 18 Sep 2025 23:48:07 +0200 Subject: [PATCH] added script to sync annotations from label studio --- scripts/sync_label_studio_annotations.py | 127 +++++++++++++++++++++++ 1 file changed, 127 insertions(+) create mode 100644 scripts/sync_label_studio_annotations.py diff --git a/scripts/sync_label_studio_annotations.py b/scripts/sync_label_studio_annotations.py new file mode 100644 index 0000000..74e2dd6 --- /dev/null +++ b/scripts/sync_label_studio_annotations.py @@ -0,0 +1,127 @@ +"""Script to convert Label Studio output JSON files to an Annotation.""" + +import json +from uuid import UUID +from dotenv import load_dotenv +from python_repositories.adapters import MinioAdapter +from data_store.dto import Annotation +from data_store.repositories import AnnotationRepository + + +def list_label_studio_output_files(minio_adapter: MinioAdapter) -> list[str]: + """List the Label Studio output files in MinIO.""" + # Check input + if not isinstance(minio_adapter, MinioAdapter): + raise ValueError("minio_adapter must be an instance of MinioAdapter") + if not minio_adapter.is_connected: + raise ConnectionError("minio_adapter must be connected to MinIO") + # List objects with the Label Studio output prefix + prefix = "label-studio-output/" + return minio_adapter._list_objects(prefix) + + +def get_label_studio_output(minio_adapter: MinioAdapter, object_name: str) -> dict: + """Get a Label Studio output file from MinIO.""" + # Check input + if not isinstance(minio_adapter, MinioAdapter): + raise ValueError("minio_adapter must be an instance of MinioAdapter") + if not minio_adapter.is_connected: + raise ConnectionError("minio_adapter must be connected to MinIO") + if not isinstance(object_name, str) or len(object_name) == 0: + raise ValueError("object_name must be a non-empty string") + # Get object from MinIO + buffer = minio_adapter._get(object_name) + if buffer is None: + raise FileNotFoundError(f"File not found in MinIO: {object_name}") + # Convert buffer to JSON + data: dict = json.load(buffer) + return data + + +def convert_label_studio_output(label_studio_output: dict) -> tuple[Annotation, UUID]: + """Convert Label Studio output to Annotation and corresponding job ID.""" + # Check input + if not isinstance(label_studio_output, dict): + raise ValueError("label_studio_output must be a dictionary") + # Extract job id + image_path = label_studio_output["task"]["data"]["image"] + job_id_str = image_path.rsplit("/")[-1].split(".")[0] + job_id = UUID(job_id_str) + # Extract relevant information from the Label Studio output + annotation_data: dict[str, str | None | bool] = {} + for elem in label_studio_output["result"]: + # Extract key and value + key = elem["from_name"] + choices = elem["value"]["choices"] + # Check for presence of uncertainty flag + uncertainty_key = f"{key}_uncertainty" + uncertainty_value = bool("uncertain" in choices) + if uncertainty_value: + choices.remove("uncertain") + # If no choices left, set to None + value: str | None = choices[0] if choices else None + # Store in annotation data + annotation_data[key] = value + annotation_data[uncertainty_key] = uncertainty_value + # Create the Annotation object + annotation = Annotation(**annotation_data) + return annotation, job_id + + +def sync_label_studio_annotations( + minio_adapter: MinioAdapter, + annotation_repository: AnnotationRepository, +) -> None: + """Sync Label Studio output files to AnnotationRepository.""" + # Check inputs + if not isinstance(minio_adapter, MinioAdapter): + raise ValueError("minio_adapter must be an instance of MinioAdapter") + if not minio_adapter.is_connected: + raise ConnectionError("minio_adapter must be connected to MinIO") + if not isinstance(annotation_repository, AnnotationRepository): + raise ValueError("annotation_repository must be an instance of AnnotationRepository") + if not annotation_repository.is_connected: + raise ConnectionError("annotation_repository must be connected to MinIO") + # List annotations in repository + annotation_ids = annotation_repository.list_all() + print(f"Found {len(annotation_ids)} annotations in repository.") + # Count Label Studio output files + label_studio_files = list_label_studio_output_files(minio_adapter) + print(f"Found {len(label_studio_files)} Label Studio output files.") + # Stop early if equal number of annotations and Label Studio files + if len(annotation_ids) == len(label_studio_files): + print("No new Label Studio output files to process. Exiting.") + return + # Process each Label Studio output file + for file_name in label_studio_files: + print(f"Processing file: {file_name}") + # Get Label Studio output + try: + label_studio_output = get_label_studio_output( + minio_adapter, + file_name, + ) + except Exception as exc: + print(f"Failed to get Label Studio output for {file_name}: {exc}") + continue + # Convert to Annotation + try: + annotation, job_id = convert_label_studio_output(label_studio_output) + except Exception as exc: + print(f"Failed to convert Label Studio output for {file_name}: {exc}") + continue + # Store Annotation in repository + try: + annotation_repository.put(annotation, job_id) + except Exception as exc: + print(f"Failed to store annotation for {file_name}: {exc}") + continue + print(f"Successfully processed and stored annotation for {file_name}.") + + +if __name__ == "__main__": + # Load environment variables from .env file + load_dotenv() + # Connect to data sources + with MinioAdapter() as minio_adapter, AnnotationRepository() as annotation_repo: + sync_label_studio_annotations(minio_adapter, annotation_repo)