From e481165c36e1bf90675119b29f03d47b3044c4bd Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Fri, 19 Jun 2026 14:43:20 -0300 Subject: [PATCH] add split saving --- datamint/dataset/__init__.py | 2 + datamint/dataset/base.py | 20 ++++++--- datamint/dataset/split_result.py | 57 ++++++++++++++++++++++++ notebooks/patient_wise_split.ipynb | 69 +++++++++++++++--------------- 4 files changed, 107 insertions(+), 41 deletions(-) create mode 100644 datamint/dataset/split_result.py diff --git a/datamint/dataset/__init__.py b/datamint/dataset/__init__.py index d54a5170..6ea11f59 100644 --- a/datamint/dataset/__init__.py +++ b/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', ] \ No newline at end of file diff --git a/datamint/dataset/base.py b/datamint/dataset/base.py index a4351f8d..08e1d2b0 100644 --- a/datamint/dataset/base.py +++ b/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, diff --git a/datamint/dataset/split_result.py b/datamint/dataset/split_result.py new file mode 100644 index 00000000..4d9d2aee --- /dev/null +++ b/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) diff --git a/notebooks/patient_wise_split.ipynb b/notebooks/patient_wise_split.ipynb index 0475a25a..133d9c07 100644 --- a/notebooks/patient_wise_split.ipynb +++ b/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,7 +121,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 8, "id": "e5f6a7b8", "metadata": {}, "outputs": [ @@ -119,11 +129,11 @@ "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,7 +307,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 13, "id": "d6e7f8a9", "metadata": {}, "outputs": [ @@ -314,7 +315,7 @@ "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" ] }