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)