Feat/mlflow integration

README.md

@@ -1,25 +1,74 @@
-
-# Datamint python API
+# Datamint Python API
 
 ![Build Status](https://github.com/SonanceAI/datamint-python-api/actions/workflows/run_test.yaml/badge.svg)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+
+A comprehensive Python SDK for interacting with the Datamint platform, providing seamless integration for medical imaging workflows, dataset management, and machine learning experiments.
+
+## 📋 Table of Contents
+
+- [Features](#-features)
+- [Installation](#-installation)
+- [Quick Setup](#-quick-setup)
+- [Documentation](#-documentation)
+- [Key Components](#-key-components)
+- [Command Line Tools](#️-command-line-tools)
+- [Examples](#-examples)
+- [Support](#-support)
+
+## 🚀 Features
+
+- **Dataset Management**: Download, upload, and manage medical imaging datasets
+- **Annotation Tools**: Create, upload, and manage annotations (segmentations, labels, measurements)
+- **Experiment Tracking**: Integrated MLflow support for experiment management
+- **PyTorch Lightning Integration**: Streamlined ML workflows with Lightning DataModules and callbacks
+- **DICOM Support**: Native handling of DICOM files with anonymization capabilities
+- **Multi-format Support**: PNG, JPEG, NIfTI, and other medical imaging formats
 
 See the full documentation at https://sonanceai.github.io/datamint-python-api/
 
-## Installation
+## 📦 Installation
+
+> [!NOTE]
+> We recommend using a virtual environment to avoid package conflicts.
 
-Datamint requires Python 3.10+.
-You can install/update Datamint and its dependencies using pip
+### From PyPI
+
+To be released soon
+
+### From Source
 
 ```bash
-pip install -U datamint
+pip install git+https://github.com/SonanceAI/datamint-python-api
 ```
 
+### Virtual Environment Setup
+
+<details>
+<summary>Click to expand virtual environment setup instructions</summary>
+
 We recommend that you install Datamint in a dedicated virtual environment, to avoid conflicting with your system packages.
-Create the enviroment once with `python3 -m venv datamint-env` and then activate it whenever you need it with:
-- `source datamint-env/bin/activate` (Linux/MAC)
-- `datamint-env\Scripts\activate.bat` (Windows CMD)
-- `datamint-env\Scripts\Activate.ps1` (Windows PowerShell)
+For instance, create the enviroment once with `python3 -m venv datamint-env` and then activate it whenever you need it with:
+
+1. **Create the environment** (one-time setup):
+   ```bash
+   python3 -m venv datamint-env
+   ```
+
+2. **Activate the environment** (run whenever you need it):
+
+   | Platform | Command |
+   |----------|---------|
+   | Linux/macOS | `source datamint-env/bin/activate` |
+   | Windows CMD | `datamint-env\Scripts\activate.bat` |
+   | Windows PowerShell | `datamint-env\Scripts\Activate.ps1` |
 
+3. **Install the package**:
+   ```bash
+   pip install git+https://github.com/SonanceAI/datamint-python-api
+   ```
+
+</details>
 
 ## Setup API key
 
@@ -45,29 +94,174 @@ import os
 os.environ["DATAMINT_API_KEY"] = "my_api_key"
 ```
 
-### Method 3: Api constructor
+## 📚 Documentation
+
+| Resource | Description |
+|----------|-------------|
+| [🚀 Getting Started](docs/source/getting_started.rst) | Step-by-step setup and basic usage |
+| [📖 API Reference](docs/source/client_api.rst) | Complete API documentation |
+| [🔥 PyTorch Integration](docs/source/pytorch_integration.rst) | ML workflow integration |
+| [💡 Examples](examples/) | Practical usage examples |
+
+## 🔗 Key Components
+
+### Dataset Management
+
+```python
+from datamint import Dataset
+
+# Load dataset with annotations
+dataset = Dataset(
+    project_name="medical-segmentation",
+)
+
+# Access data
+for sample in dataset:
+    image = sample['image']       # torch.Tensor
+    mask = sample['segmentation'] # torch.Tensor (if available)
+    metadata = sample['metainfo'] # dict
+```
+
+
+### PyTorch Lightning Integration
+
+```python
+import lightning as L
+from datamint.lightning import DatamintDataModule
+from datamint.mlflow.lightning.callbacks import MLFlowModelCheckpoint
+
+# Data module
+datamodule = DatamintDataModule(
+    project_name="your-project",
+    batch_size=16,
+    train_split=0.8
+)
+
+# ML tracking callback
+checkpoint_callback = MLFlowModelCheckpoint(
+    monitor="val_loss",
+    save_top_k=1,
+    register_model_name="best-model"
+)
+
+# Trainer with MLflow logging
+trainer = L.Trainer(
+    max_epochs=100,
+    callbacks=[checkpoint_callback],
+    logger=L.pytorch.loggers.MLFlowLogger(
+        experiment_name="medical-segmentation"
+    )
+)
+```
+
+
+### Annotation Management
 
-Specify API key in the Api constructor:
 
 ```python
-from datamint import Api
-api = Api(api_key='my_api_key')
+# Upload segmentation masks
+api.upload_segmentations(
+    resource_id="resource-123",
+    file_path="segmentation.nii.gz",
+    name="liver_segmentation",
+    frame_index=0
+)
+
+# Add categorical annotations
+api.add_image_category_annotation(
+    resource_id="resource-123",
+    identifier="diagnosis",
+    value="positive"
+)
+
+# Add geometric annotations
+api.add_line_annotation(
+    point1=(10, 20),
+    point2=(50, 80),
+    resource_id="resource-123",
+    identifier="measurement",
+    frame_index=5
+)
+```
+
+
+## 🛠️ Command Line Tools
+
+### Upload Resources
+
+**Upload DICOM files with anonymization:**
+```bash
+datamint-upload \
+    --path /path/to/dicoms \
+    --recursive \
+    --channel "training-data" \
+    --anonymize \
+    --publish
+```
+
+**Upload with segmentation masks:**
+```bash
+datamint-upload \
+    --path /path/to/images \
+    --segmentation_path /path/to/masks \
+    --segmentation_names segmentation_config.yaml
 ```
 
-## Tutorials
+### Configuration Management
 
+```bash
+# Interactive setup
+datamint-config
+
+# Set API key
+datamint-config --api-key "your-key"
+```
 
-You can find example notebooks in the `notebooks` folder:
+## 🔍 Examples
 
-- [Uploading your resources](notebooks/upload_data.ipynb)
-- [Uploading model segmentations](notebooks/upload_model_segmentations.ipynb)
+### Medical Image Segmentation Pipeline
 
-and example scripts in [examples](examples) folder:
+```python
+import torch
+import lightning as L
+from datamint.lightning import DatamintDataModule
+from datamint.mlflow.lightning.callbacks import MLFlowModelCheckpoint
+
+class SegmentationModel(L.LightningModule):
+    def __init__(self):
+        super().__init__()
+        # Model definition...
+
+    def training_step(self, batch, batch_idx):
+        # Training logic...
+        pass
+
+# Setup data
+datamodule = DatamintDataModule(
+    project_name="liver-segmentation",
+    batch_size=8,
+    train_split=0.8
+)
+
+# Setup model with MLflow tracking
+model = SegmentationModel()
+checkpoint_cb = MLFlowModelCheckpoint(
+    monitor="val_dice",
+    mode="max",
+    register_model_name="liver-segmentation-model"
+)
+
+# Train
+trainer = L.Trainer(
+    max_epochs=50,
+    callbacks=[checkpoint_cb],
+    logger=L.pytorch.loggers.MLFlowLogger()
+)
+trainer.fit(model, datamodule)
+```
 
-- [API usage examples](examples/api_usage.ipynb)
-- [Project and entity usage](examples/project_entity_usage.ipynb)
-- [Channels example](examples/channels_example.ipynb)
+## 🆘 Support
 
-## Full documentation
+[Full Documentation](https://datamint-python-api.readthedocs.io/) 
+[GitHub Issues](https://github.com/SonanceAI/datamint-python-api/issues)
 
-See all functionalities in the full documentation at https://sonanceai.github.io/datamint-python-api/

datamint/api/base_api.py

@@ -61,22 +61,56 @@ def __init__(self,
             client: Optional HTTP client instance. If None, a new one will be created.
         """
         self.config = config
-        self.client = client or self._create_client()
+        self._owns_client = client is None  # Track if we created the client
+        self.client = client or BaseApi._create_client(config)
         self.semaphore = asyncio.Semaphore(20)
         self._api_instance: 'Api | None' = None  # Injected by Api class
 
-    def _create_client(self) -> httpx.Client:
-        """Create and configure HTTP client with authentication and timeouts."""
-        headers = None
-        if self.config.api_key:
-            headers = {"apikey": self.config.api_key}
+    @staticmethod
+    def _create_client(config: ApiConfig) -> httpx.Client:
+        """Create and configure HTTP client with authentication and timeouts.
+
+        The client is designed to be long-lived and reused across multiple requests.
+        It maintains connection pooling for improved performance.
+        Default limits: max_keepalive_connections=20, max_connections=100
+        """
+        headers = {"apikey": config.api_key} if config.api_key else None
 
         return httpx.Client(
-            base_url=self.config.server_url,
+            base_url=config.server_url,
             headers=headers,
-            timeout=self.config.timeout
+            timeout=config.timeout,
+            limits=httpx.Limits(
+                max_keepalive_connections=5,  # Increased from default 20
+                max_connections=20,  # Increased from default 100
+                keepalive_expiry=8
+            )
         )
 
+    def close(self) -> None:
+        """Close the HTTP client and release resources.
+
+        Should be called when the API instance is no longer needed.
+        Only closes the client if it was created by this instance.
+        """
+        if self._owns_client and self.client is not None:
+            self.client.close()
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures client is closed."""
+        self.close()
+
+    def __del__(self):
+        """Destructor - ensures client is closed when instance is garbage collected."""
+        try:
+            self.close()
+        except Exception:
+            pass  # Ignore errors during cleanup
+
     def _stream_request(self, method: str, endpoint: str, **kwargs):
         """Make streaming HTTP request with error handling.
 

datamint/api/client.py

@@ -68,6 +68,8 @@ def check_connection(self):
                                     f" Please check your api_key and/or other configurations. {e}")
 
     def _get_endpoint(self, name: str):
+        if self._client is None:
+            self._client = BaseApi._create_client(self.config)
         if name not in self._endpoints:
             api_class = self._API_MAP[name]
             endpoint = api_class(self.config, self._client)

datamint/apihandler/base_api_handler.py

@@ -30,7 +30,6 @@
 
 _PAGE_LIMIT = 5000
 
-
 @deprecated(reason="Please use `from datamint import Api` instead.", version="2.0.0")
 class BaseAPIHandler:
     """

datamint/apihandler/dto/annotation_dto.py

@@ -178,6 +178,8 @@ def __init__(self,
         if model_id is not None:
             if is_model == False:
                 raise ValueError("model_id==False while self.model_id is provided.")
+            if not isinstance(model_id, str):
+                raise ValueError("model_id must be a string if provided.")
             is_model = True
         self.is_model = is_model
         self.geometry = geometry

datamint/dataset/base_dataset.py

@@ -307,6 +307,10 @@ def _setup_labels(self) -> None:
         self.image_lsets, self.image_lcodes = self._get_labels_set(framed=False)
         worklist_id = self.get_info()['worklist_id']
         groups: dict[str, dict] = self.api.annotationsets.get_segmentation_group(worklist_id)['groups']
+        if not groups:
+            self.seglabel_list = []
+            self.seglabel2code = {}
+            return
         # order by 'index' key
         max_index = max([g['index'] for g in groups.values()])
         self.seglabel_list : list[str] = ['UNKNOWN'] * max_index  # 1-based

datamint/lightning/__init__.py

@@ -0,0 +1 @@
+from .datamintdatamodule import DatamintDataModule