workspace datapoint management

docs/sdk.rst

@@ -21,7 +21,7 @@ Organization
 Workspace
 ----------------------
 .. autoclass:: redbrick.workspace.RBWorkspace
-   :members: name, org_id, workspace_id, metadata_schema, classification_schema, cohorts, update_schema, update_cohorts, get_datapoints, archive_datapoints, unarchive_datapoints, add_datapoints_to_cohort, remove_datapoints_from_cohort, update_datapoint_attributes
+   :members: name, org_id, workspace_id, metadata_schema, classification_schema, cohorts, update_schema, update_cohorts, get_datapoints, create_datapoints, archive_datapoints, unarchive_datapoints, delete_datapoints, add_datapoints_to_cohort, add_datapoints_to_projects, remove_datapoints_from_cohort, update_datapoint_attributes
    :show-inheritance:
 
 .. _project:

redbrick/cli/command/upload.py

@@ -11,6 +11,7 @@
 from redbrick.cli.project import CLIProject
 from redbrick.cli.cli_base import CLIUploadInterface
 from redbrick.common.enums import StorageMethod, ImportTypes
+from redbrick.upload.interact import create_tasks, prepare_json_files
 from redbrick.utils.logging import assert_validation, logger
 from redbrick.utils.files import find_files_recursive
 from redbrick.types.task import InputTask
@@ -214,17 +215,20 @@ def handle_upload(self) -> None:  # noqa: ignore=C901
                 with open(item_group[0], "r", encoding="utf-8") as file_:
                     files_data.append(json.load(file_))
             logger.debug("Preparing json files for upload")
-            points = self.project.project.upload.prepare_json_files(
-                files_data,
-                storage_id,
-                label_storage_id,
-                segmentation_mapping,
-                task_dirs,
-                upload_cache,
-                self.args.rt_struct,
-                self.args.mhd,
-                self.args.label_validate,
-                self.args.concurrency,
+            points = prepare_json_files(
+                context=self.project.context,
+                org_id=self.project.org_id,
+                taxonomy=self.project.project.taxonomy,
+                files_data=files_data,
+                storage_id=storage_id,
+                label_storage_id=label_storage_id,
+                task_segment_map=segmentation_mapping,
+                task_dirs=task_dirs,
+                uploaded=upload_cache,
+                rt_struct=self.args.rt_struct,
+                mhd_mask=self.args.mhd,
+                label_validate=self.args.label_validate,
+                concurrency=self.args.concurrency,
             )
             segmentation_mapping = {}
         else:
@@ -368,16 +372,19 @@ def handle_upload(self) -> None:  # noqa: ignore=C901
             logger.info(f"Found {len(points)} items")
 
             uploads = asyncio.run(
-                project.upload._create_tasks(
-                    points,
-                    segmentation_mapping,
-                    self.args.ground_truth,
-                    storage_id,
-                    label_storage_id,
-                    self.args.label_validate,
-                    self.args.prune_segmentations,
-                    self.args.concurrency,
-                    False,
+                create_tasks(
+                    context=project.context,
+                    org_id=project.org_id,
+                    workspace_id=None,
+                    project_id=project.project_id,
+                    points=points,
+                    segmentation_mapping=segmentation_mapping,
+                    is_ground_truth=self.args.ground_truth,
+                    storage_id=storage_id,
+                    label_storage_id=label_storage_id,
+                    label_validate=self.args.label_validate,
+                    prune_segmentations=self.args.prune_segmentations,
+                    concurrency=self.args.concurrency,
                 )
             )
 

redbrick/common/upload.py

@@ -16,7 +16,8 @@ async def create_datapoint_async(  # pylint: disable=too-many-locals
         self,
         aio_client: aiohttp.ClientSession,
         org_id: str,
-        project_id: str,
+        workspace_id: Optional[str],
+        project_id: Optional[str],
         storage_id: str,
         name: str,
         items: List[str],
@@ -43,9 +44,10 @@ async def update_items_async(
         self,
         aio_client: aiohttp.ClientSession,
         org_id: str,
-        project_id: str,
         storage_id: str,
-        task_id: str,
+        dp_id: Optional[str] = None,
+        project_id: Optional[str] = None,
+        task_id: Optional[str] = None,
         items: Optional[List[str]] = None,
         series_info: Optional[List[Dict]] = None,
         heat_maps: Optional[List[Dict]] = None,
@@ -61,6 +63,12 @@ def items_upload_presign(
     ) -> List[Dict[Any, Any]]:
         """Get a presigned url for uploading items."""
 
+    @abstractmethod
+    async def delete_datapoints(
+        self, aio_client: aiohttp.ClientSession, org_id: str, dp_ids: List[str]
+    ) -> bool:
+        """Delete datapoints in a workspace."""
+
     @abstractmethod
     async def delete_tasks(
         self,

redbrick/common/workspace.py

@@ -77,3 +77,14 @@ def update_datapoint_attributes(
         self, org_id: str, dp_id: str, attributes: Dict
     ) -> None:
         """Update datapoint attributes."""
+
+    @abstractmethod
+    def add_datapoints_to_projects(
+        self,
+        org_id: str,
+        workspace_id: str,
+        project_ids: List[str],
+        dp_ids: List[str],
+        ground_truth: bool,
+    ) -> None:
+        """Add datapoints to project."""

redbrick/labeling/public.py

@@ -14,11 +14,11 @@
 from redbrick.common.enums import StorageMethod
 from redbrick.stage import Stage
 from redbrick.types.taxonomy import Taxonomy
+from redbrick.upload.interact import validate_json
 from redbrick.utils.upload import (
     convert_mhd_to_nii_labels,
     convert_rt_struct_to_nii_labels,
     process_segmentation_upload,
-    validate_json,
 )
 from redbrick.utils.logging import log_error, logger
 from redbrick.utils.async_utils import gather_with_concurrency, get_session

redbrick/repo/shards.py

@@ -284,6 +284,10 @@ def datapoint_shard(raw_items: bool, presigned_items: bool) -> str:
             storageId
         }}
         attributes
+        archived
+        cohorts {{
+            name
+        }}
     """
 
 

redbrick/repo/upload.py

@@ -21,7 +21,8 @@ async def create_datapoint_async(
         self,
         aio_client: aiohttp.ClientSession,
         org_id: str,
-        project_id: str,
+        workspace_id: Optional[str],
+        project_id: Optional[str],
         storage_id: str,
         name: str,
         items: List[str],
@@ -46,7 +47,8 @@ async def create_datapoint_async(
         query_string = """
             mutation createDatapointSDK(
                 $orgId: UUID!
-                $projectId: UUID!
+                $workspaceId: UUID
+                $projectId: UUID
                 $items: [String!]!
                 $heatMaps: [HeatMapInput!]
                 $transforms: [TransformInput!]
@@ -64,6 +66,7 @@ async def create_datapoint_async(
             ) {
                 createDatapoint(
                     orgId: $orgId
+                    workspaceId: $workspaceId
                     projectId: $projectId
                     items: $items
                     heatMaps: $heatMaps
@@ -80,6 +83,7 @@ async def create_datapoint_async(
                     priority: $priority
                     attributes: $attributes
                 ) {
+                    dpId
                     taskId
                     taskIds
                 }
@@ -88,6 +92,7 @@ async def create_datapoint_async(
 
         query_variables = {
             "orgId": org_id,
+            "workspaceId": workspace_id,
             "projectId": project_id,
             "items": items,
             "heatMaps": heat_maps,
@@ -124,9 +129,10 @@ async def update_items_async(
         self,
         aio_client: aiohttp.ClientSession,
         org_id: str,
-        project_id: str,
         storage_id: str,
-        task_id: str,
+        dp_id: Optional[str] = None,
+        project_id: Optional[str] = None,
+        task_id: Optional[str] = None,
         items: Optional[List[str]] = None,
         series_info: Optional[List[Dict]] = None,
         heat_maps: Optional[List[Dict]] = None,
@@ -135,12 +141,14 @@ async def update_items_async(
         meta_data: Optional[Dict] = None,
     ) -> Dict:
         """Update items in a datapoint."""
+        # pylint: disable=too-many-locals
         query_string = """
             mutation updateTaskItemsSDK(
                 $orgId: UUID!
-                $projectId: UUID!
+                $dpId: UUID
+                $projectId: UUID
+                $taskId: UUID
                 $storageId: UUID!
-                $taskId: UUID!
                 $items: [String!]
                 $seriesInfo: [SeriesInfoInput!]
                 $heatMaps: [HeatMapInput!]
@@ -150,9 +158,10 @@ async def update_items_async(
             ) {
                 updateTaskItems(
                     orgId: $orgId
+                    dpId: $dpId
                     projectId: $projectId
-                    storageId: $storageId
                     taskId: $taskId
+                    storageId: $storageId
                     items: $items
                     seriesInfo: $seriesInfo
                     heatMaps: $heatMaps
@@ -168,9 +177,10 @@ async def update_items_async(
 
         query_variables = {
             "orgId": org_id,
+            "dpId": dp_id,
             "projectId": project_id,
-            "storageId": storage_id,
             "taskId": task_id,
+            "storageId": storage_id,
             "items": items,
             "seriesInfo": series_info,
             "heatMaps": heat_maps,
@@ -228,6 +238,32 @@ def items_upload_presign(
         presigned: List[Dict] = result["itemsUploadPresign"]["items"]
         return presigned
 
+    async def delete_datapoints(
+        self, aio_client: aiohttp.ClientSession, org_id: str, dp_ids: List[str]
+    ) -> bool:
+        """Delete datapoints in a workspace."""
+        query_string = """
+        mutation deleteDatapointsSDK($orgId: UUID!, $dpIds: [UUID!]!) {
+            deleteDatapoints(
+                orgId: $orgId
+                dpIds: $dpIds
+            ) {
+                ok
+            }
+        }
+        """
+        # EXECUTE THE QUERY
+        query_variables = {
+            "orgId": org_id,
+            "dpIds": dp_ids,
+        }
+
+        result: Dict[str, Dict] = await self.client.execute_query_async(
+            aio_client, query_string, query_variables
+        )
+
+        return (result.get("deleteDatapoints", {}) or {}).get("ok", False)
+
     async def delete_tasks(
         self,
         aio_client: aiohttp.ClientSession,

redbrick/repo/workspace.py

@@ -240,3 +240,41 @@ def update_datapoint_attributes(
             "attributes": json.dumps(attributes),
         }
         self.client.execute_query(query, variables)
+
+    def add_datapoints_to_projects(
+        self,
+        org_id: str,
+        workspace_id: str,
+        project_ids: List[str],
+        dp_ids: List[str],
+        ground_truth: bool,
+    ) -> None:
+        """Add datapoints to project."""
+        query = """
+        mutation importDatapointsFromWorkspaceSDK(
+            $orgId: UUID!
+            $workspaceId: UUID!
+            $projectIds: [UUID!]!
+            $dpIds: [UUID!]!
+            $isGroundTruth: Boolean
+        ) {
+            importDatapointsFromWorkspace(
+                orgId: $orgId
+                workspaceId: $workspaceId
+                projectIds: $projectIds
+                dpIds: $dpIds
+                isGroundTruth: $isGroundTruth
+            ) {
+                ok
+                message
+            }
+        }
+        """
+        variables = {
+            "orgId": org_id,
+            "workspaceId": workspace_id,
+            "projectIds": project_ids,
+            "dpIds": dp_ids,
+            "isGroundTruth": ground_truth,
+        }
+        self.client.execute_query(query, variables)