Users might want to save splits on the server (DAT-934)

datamint/dataset/__init__.py

@@ -19,6 +19,7 @@
 from .sliced_video_dataset import SlicedVideoDataset
 from .detection_dataset import DetectionDataset, detection_collate_fn
 from .factory import build_dataset
+from .split_result import SplitResult
 
 __all__ = [
     # Core
@@ -35,4 +36,5 @@
     'detection_collate_fn',
     # Factory
     'build_dataset',
+    'SplitResult',
 ]

datamint/dataset/base.py

@@ -20,10 +20,12 @@
 from datamint.entities.annotations.annotation_spec import AnnotationSpec, CategoryAnnotationSpec
 from datamint.entities.annotations import AnnotationType
 
+
 if TYPE_CHECKING:
     from datamint.entities import Resource, Project, Annotation
     from albumentations import BaseCompose
     from datamint.mlflow.data import DatamintMLflowDataset
+    from .split_result import SplitResult
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -1162,7 +1164,8 @@ def split(
         by_patient: bool = False,
         none_patient_id_strategy: Literal['individual', 'group', 'skip', 'error'] = 'individual',
         **splits: float,
-    ) -> dict[str, 'DatamintBaseDataset']:
+    ) -> 'SplitResult':
+
         """Split the dataset into multiple named subsets.
 
         The mode is selected automatically when no explicit split mode is
@@ -1215,23 +1218,26 @@ def split(
         Raises:
             ValueError: If ratios are invalid or arguments conflict.
         """
+
+        from .split_result import SplitResult
+
         if by_patient:
             if use_project_splits or use_server_splits:
                 raise ValueError(
                     "by_patient=True cannot be combined with use_project_splits or use_server_splits."
                 )
-        
+
             if not splits:
                 raise ValueError(
                     "by_patient=True requires ratio kwargs (e.g. train=0.8, test=0.2) to determine split sizes."
                 )
-            return self._split_locally_by_patient(dict(splits), seed, none_patient_id_strategy)
-        
+            return SplitResult(self._split_locally_by_patient(dict(splits), seed, none_patient_id_strategy))
+
         if use_project_splits is None and use_server_splits is None and not splits:
             use_project_splits = getattr(self, 'project', None) is not None or as_of_timestamp is not None
 
         if use_project_splits:
-            return self._split_by_project_api(splits, as_of_timestamp=as_of_timestamp)
+            return SplitResult(self._split_by_project_api(splits, as_of_timestamp=as_of_timestamp))
 
         if as_of_timestamp is not None:
             raise ValueError(
@@ -1246,9 +1252,9 @@ def split(
             import warnings
             warnings.warn("use_server_splits and splitting by resource tags are deprecated in favor of use_project_splits. "
                           "Please migrate to project-scoped splits for better reproducibility and management.", DeprecationWarning)
-            return self._split_by_server_tags(splits)
+            return SplitResult(self._split_by_server_tags(splits))
 
-        return self._split_locally(splits, seed)
+        return SplitResult(self._split_locally(splits, seed))
 
     def _split_by_project_api(
         self,

datamint/dataset/split_result.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .base import DatamintBaseDataset
+
+
+class SplitResult(dict):
+    """A dict of split name → dataset that can persist itself to the server.
+
+    Behaves exactly like a plain ``dict`` -- indexing, iteration, and all
+    standard dict operations work unchanged.  The extra ``.save()`` method
+    pushes the split assignments to the Datamint project so they become the
+    official server-side split.
+
+    Example::
+
+        parts = dataset.split(train=0.8, test=0.2, by_patient=True, seed=42)
+
+        parts['train']          # works as before
+        for name, ds in parts.items(): ...  # works as before
+
+        parts.save()            # persist to server
+        parts.save(force=True)  # overwrite existing assignments
+    """
+
+    def save(self, force: bool = False) -> None:
+        """Persist split assignments to the server.
+
+        Args:
+            force: If ``True``, overwrite any existing split assignments on
+                the project.  If ``False`` (default) and the project already
+                has assignments, a ``ValueError`` is raised.
+
+        Raises:
+            ValueError: If the dataset was not loaded from a project, or if
+                the project already has split assignments and ``force=False``.
+        """
+        any_ds: DatamintBaseDataset = next(iter(self.values()))
+        project = getattr(any_ds, 'project', None)
+        if project is None:
+            raise ValueError(
+                "Cannot save splits: the dataset was not loaded from a project. "
+                "Load with ImageDataset(project='...') or VolumeDataset(project='...') first."
+            )
+
+        api = any_ds._api
+        existing = api.projects.get_splits(project)
+        if existing and not force:
+            raise ValueError(
+                f"Project '{project.name}' already has split assignments. "
+                "Use save(force=True) to overwrite."
+            )
+
+        for split_name, ds in self.items():
+            api.projects.assign_splits(project, ds.resources, split_name)

notebooks/patient_wise_split.ipynb

@@ -46,10 +46,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "c3d4e5f6",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Project 'split_save_test' already exists, skipping upload.\n",
+      "\n",
+      "Loaded 20 resources\n"
+     ]
+    }
+   ],
    "source": [
     "import pydicom\n",
     "import pydicom.data\n",
@@ -58,7 +68,7 @@
     "from datamint import Api\n",
     "from datamint.dataset import ImageDataset\n",
     "\n",
-    "PROJECT_NAME = \"patient_wise_split_demo\"\n",
+    "PROJECT_NAME = \"split_save_test\"\n",
     "\n",
     "# 8 patients selected from pydicom's \n",
     "SELECTED_PATIENTS = {\n",
@@ -111,19 +121,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 8,
    "id": "e5f6a7b8",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "MR_truncated.dcm                          patient_id='4MR1'\n",
       "MR_small_RLE.dcm                          patient_id='4MR1'\n",
-      "JPEG2000.dcm                              patient_id='8NM1'\n",
       "JPEG2000-embedded-sequence-delimiter.dcm  patient_id='8NM1'\n",
-      "MR_truncated.dcm                          patient_id='4MR1'\n",
-      "MR_small_expb.dcm                         patient_id='4MR1'\n"
+      "JPEG2000.dcm                              patient_id='8NM1'\n",
+      "MR_small.dcm                              patient_id='4MR1'\n"
      ]
     }
    ],
@@ -155,7 +165,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 9,
    "id": "b8c9d0e1",
    "metadata": {},
    "outputs": [
@@ -172,9 +182,9 @@
       "  1CT1                   1 resource(s)\n",
       "  13US1                  2 resource(s)\n",
       "  CQ500-CT-310           1 resource(s)\n",
-      "  204                    1 resource(s)\n",
       "  11-05-25-142825        1 resource(s)\n",
-      "  021234567              1 resource(s)\n"
+      "  021234567              1 resource(s)\n",
+      "  204                    1 resource(s)\n"
      ]
     }
    ],
@@ -202,7 +212,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 10,
    "id": "d0e1f2a3",
    "metadata": {},
    "outputs": [
@@ -236,35 +246,26 @@
    "source": [
     "## 5. Persist Splits to the Server\n",
     "\n",
-    "The `parts` dict computed above lives only in memory. To make the split **reproducible across sessions and shareable with the team**, write the assignments back to the project using `api.projects.assign_splits()`.\n",
+    "The `parts` object returned by `split()` is a `SplitResult` — it behaves exactly like a dict, but it also knows how to save itself to the server via `.save()`.\n",
+    "\n",
+    "Once persisted, anyone can reload exactly the same split (without re-running the patient split logic) by calling `dataset.split()` with no arguments.\n",
     "\n",
-    "Once persisted, anyone can reload exactly the same split (without re-running the patient split logic) by calling `dataset.split()` with no arguments (project-backed datasets prefer the server assignments automatically)."
+    "If the project already has split assignments, `.save()` raises an error. Pass `force=True` to overwrite."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 11,
    "id": "a3b4c5d6",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Assigned    7 resources → 'train'\n",
-      "Assigned    9 resources → 'val'\n",
-      "Assigned    4 resources → 'test'\n",
-      "\n",
-      "Patient-wise split assignments saved to the project.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "for split_name, split_ds in parts.items():\n",
-    "    api.projects.assign_splits(project, split_ds.resources, split_name=split_name)\n",
-    "    print(f\"Assigned {len(split_ds):4d} resources → '{split_name}'\")\n",
+    "parts.save()\n",
     "\n",
-    "print(\"\\nPatient-wise split assignments saved to the project.\")"
+    "# If the project already has splits, this raises:\n",
+    "#   ValueError: Project '...' already has split assignments. Use save(force=True) to overwrite.\n",
+    "# To overwrite:\n",
+    "# parts.save(force=True)"
    ]
   },
   {
@@ -279,7 +280,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "id": "c5d6e7f8",
    "metadata": {},
    "outputs": [
@@ -306,15 +307,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 13,
    "id": "d6e7f8a9",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Replaying split snapshot from 2026-06-16T12:49:59.199793Z\n",
+      "Replaying split snapshot from 2026-06-19T16:37:40.622014Z\n",
       "{'val': 9, 'test': 4, 'train': 7}\n"
      ]
     }