130 lines
5.4 KiB
Python
130 lines
5.4 KiB
Python
"""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) # type: ignore[arg-type]
|
|
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)
|