Data Integrations

import json
import os
from typing import List, Any

from encord import EncordUserClient, Project
import dotenv
from encord.orm.dataset import LongPollingStatus
from encord.orm.storage import CustomerProvidedDicomSeriesDicomFileMetadata, DataUploadDicomSeriesDicomFile, \
    DataUploadItems, DataUploadDicomSeries

import json
from typing import List, Dict, Optional


def clean_dicom_metadata(metadata: Dict[str, Any]) -> Dict[str, Optional[Dict]]:
    """
    Clean DICOM metadata by converting string 'None' to actual None values
    and ensuring all values are either dicts or None.
    """
    cleaned = {}

    for tag_key, tag_value in metadata.items():
        if tag_value == 'None' or tag_value is None:
            cleaned[tag_key] = None
        elif isinstance(tag_value, dict):
            cleaned[tag_key] = tag_value
        else:
            # For any other type, set to None
            cleaned[tag_key] = None

    return cleaned


def load_dicom_series_from_json(json_path: str) -> List[DataUploadDicomSeries]:
    """
    Load DICOM series data from JSON and convert to DataUploadDicomSeries objects.

    Args:
        json_path: Path to the JSON file containing DICOM series data

    Returns:
        List of DataUploadDicomSeries objects
    """
    with open(json_path, 'r') as f:
        data = json.load(f)

    series_list = []

    for series in data.get('dicom_series', []):
        # Extract series title
        series_title = series.get('title')

        # Process metadata for this series
        metadata = None
        if 'dicom_metadata' in series and series['dicom_metadata']:
            cleaned_tags = clean_dicom_metadata(series['dicom_metadata'])
            metadata = CustomerProvidedDicomSeriesDicomFileMetadata(
                tags=cleaned_tags
            )

        # Collect all DICOM files for this series
        dicom_files = []
        for key, value in series.items():
            if key.startswith('objectUrl_'):
                # Extract file index from key
                file_index = key.replace('objectUrl_', '')

                # Create individual file title
                file_title = f"{series_title}_file_{file_index}" if series_title else f"file_{file_index}"

                # Create the DICOM file object
                dicom_file = DataUploadDicomSeriesDicomFile(
                    url=value,
                    title=file_title,
                    dicom_metadata=metadata
                )

                dicom_files.append(dicom_file)

        # Create the series object with all its files
        if dicom_files:  # Only create series if it has files
            dicom_series = DataUploadDicomSeries(
                dicom_files=dicom_files,
                title=series_title,
                client_metadata={},  # Add any custom metadata here if needed
                external_file_type="DICOM"
            )
            series_list.append(dicom_series)

    return series_list

dotenv.load_dotenv()

SSH_PATH = "/Users/chris-encord/ssh-private-key.txt" # Replace with the file path to your SSH private key
FOLDER_HASH = "00000000-0000-0000-0000-000000000000" # Replace with the unique hash for the destination Folder
INTEGRATION_HASH = "00000000-0000-0000-0000-000000000000" # Replace with the unique hash for the integration
JSON_FILE_PATH = "/Users/chris-encord/dicom_to_upload.json" # Replace with file path to the JSON file

#Authenticate using the path to your private key
client = EncordUserClient.create_with_ssh_private_key(ssh_private_key_path=SSH_PATH)

storage_folder = client.get_storage_folder(FOLDER_HASH)

# Specify the integration you want to use by replacing <integration_title> with the integration title
integrations = client.get_cloud_integrations(filter_integration_uuids=[INTEGRATION_HASH])
integration = integrations[0].id

dicom_series_for_upload = load_dicom_series_from_json(JSON_FILE_PATH)

# Initiate cloud data registration to the storage folder
upload_job_id = storage_folder.add_private_data_to_folder_start(
    integration_id=integration, private_files=DataUploadItems(dicom_series=dicom_series_for_upload), ignore_errors=True
)

# timeout_seconds determines how long the code waits after initiating upload before checking the upload status
res = storage_folder.add_private_data_to_folder_get_result(upload_job_id, timeout_seconds=30)
print(f"Execution result: {res}")

if res.status == LongPollingStatus.PENDING:
    print("Upload is still in progress, try again later!")
elif res.status == LongPollingStatus.DONE:
    print("Upload completed")
    if res.unit_errors:
        print("The following URLs failed to upload:")
        for e in res.unit_errors:
            print(e.object_urls)
else:
    print(f"Upload failed: {res.errors}")