feat: add data loading

.gitignore

@@ -0,0 +1,183 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm-project.org/#use-with-ide
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# ruff
+.ruff_cache/
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# VSCode
+.vscode/
+
+# MacOS
+.DS_Store
+
+# Mlflow
+mlruns/
+mlflow_store/
+mlflow_artifacts/
+
+# Pytorch Lightning
+lightning_logs/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+
+# Project
+/data
+/demo_*

configs/data/datasets/stain_normalization/predict.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - /data/normalize@normalize: default
+
+_target_: stain_normalization.data.datasets.PredictDataset
+uris: ["mlflow-artifacts:/79/ce91ce61ac8e4b4784b76a7193e59da9/artifacts/FNBrno_400/2"]

configs/data/datasets/stain_normalization/test.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - /data/modify@modify: test
+  - /data/normalize@normalize: default
+
+_target_: stain_normalization.data.datasets.TestDataset
+uris: ["mlflow-artifacts:/65/c47cd842fbf14bf399cd4aca43af3127/artifacts/MMCI TL Test 512"]

configs/data/datasets/stain_normalization/train.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - /data/modify@modify: train
+  - /data/normalize@normalize: default
+
+_target_: stain_normalization.data.datasets.TrainDataset
+uris: ["mlflow-artifacts:/65/b587398845c54e1c84e8b60f11c42b15/artifacts/MMCI TL Train 512"]

configs/data/datasets/stain_normalization/val.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - /data/modify@modify: train
+  - /data/normalize@normalize: default
+
+_target_: stain_normalization.data.datasets.TrainDataset
+uris: ["mlflow-artifacts:/65/a43da470f89848abb3a32e59c504ee1a/artifacts/MMCI TL Val 512"]
+

configs/data/normalize/default.yaml

@@ -0,0 +1,13 @@
+_target_: albumentations.Normalize
+# From PC-A-42 MOU prostate dataset
+# Original 255-scale values: 
+# mean=[228.5544, 178.8584, 219.8793], std=[27.8285, 51.4639, 26.4458]
+mean:
+  - 0.896291765
+  - 0.70140549
+  - 0.862271765
+std:
+  - 0.109131373
+  - 0.201819216
+  - 0.10370902
+max_pixel_value: 1

pyproject.toml

@@ -0,0 +1,29 @@
+[project]
+name = "stain-normalization"
+version = "0.1.0"
+authors = [{name = "Adam Lopatka"}]
+requires-python = ">=3.12.5,<3.14"
+readme = "README.md"
+license = { file = "LICENSE" }
+dependencies = [
+    "lightning>=2.0.0",
+    "albumentations>=1.4.14",
+    "torchmetrics>=1.4.3",
+    "torchvision>=0.17.2",
+    "torch>=2.2.2",
+    "tqdm>=4.66.5",
+    "rationai-mlkit @ git+https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/mlkit.git",
+    "rationai-masks @ git+https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/masks.git",
+    "rationai-tiling @ git+https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/tiling.git",
+    "scikit-image>=0.25.2",
+    "rationai-staining @ git+https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/staining.git",
+    "openslide-python>=1.4.3",
+]
+
+[dependency-groups]
+dev = ["mypy", "ruff"]
+
+[tool.uv]
+environments = ["sys_platform == 'linux'"]
+
+override-dependencies = ["mlflow>=2.15.1"]

stain_normalization/data/__init__.py

@@ -0,0 +1,4 @@
+from stain_normalization.data.data_module import DataModule
+
+
+__all__ = ["DataModule"]

stain_normalization/data/data_module.py

@@ -0,0 +1,71 @@
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from hydra.utils import instantiate
+from lightning import LightningDataModule
+from omegaconf import DictConfig
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+from stain_normalization.type_aliases import Batch, PredictBatch
+
+
+def collate_fn(batch: list[tuple[Tensor, Any]]) -> tuple[Tensor, list[Any]]:
+    return torch.stack([x[0] for x in batch]), [x[1] for x in batch]
+
+
+class DataModule(LightningDataModule):
+    def __init__(
+        self, batch_size: int, num_workers: int = 0, **datasets: DictConfig
+    ) -> None:
+        super().__init__()
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.datasets = datasets
+
+    def setup(self, stage: str) -> None:
+        match stage:
+            case "fit":
+                self.train = instantiate(self.datasets["train"])
+                self.val = instantiate(self.datasets["val"])
+            case "validate":
+                self.val = instantiate(self.datasets["val"])
+            case "test":
+                self.test = instantiate(self.datasets["test"])
+            case "predict":
+                self.predict = instantiate(self.datasets["predict"])
+
+    def train_dataloader(self) -> Iterable[Batch]:
+        return DataLoader(
+            self.train,
+            batch_size=self.batch_size,
+            shuffle=True,
+            drop_last=True,
+            num_workers=self.num_workers,
+            persistent_workers=self.num_workers > 0,
+        )
+
+    def val_dataloader(self) -> Iterable[Batch]:
+        return DataLoader(
+            self.val,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            persistent_workers=self.num_workers > 0,
+        )
+
+    def test_dataloader(self) -> Iterable[PredictBatch]:
+        return DataLoader(
+            self.test,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=collate_fn,
+        )
+
+    def predict_dataloader(self) -> Iterable[PredictBatch]:
+        return DataLoader(
+            self.predict,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=collate_fn,
+        )

stain_normalization/data/datasets/__init__.py

@@ -0,0 +1,6 @@
+from stain_normalization.data.datasets.predict_dataset import PredictDataset
+from stain_normalization.data.datasets.test_dataset import TestDataset
+from stain_normalization.data.datasets.train_dataset import TrainDataset
+
+
+__all__ = ["PredictDataset", "TestDataset", "TrainDataset"]