Update package versions and include uv installation support

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 **/__pycache__/
 
 .vscode/
+.venv/
 *.json
 
 outputs/

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -0,0 +1,180 @@
+# Cell Segmentator
+
+## Overview
+
+This repository provides two main scripts to configure and run a cell segmentation workflow:
+
+* **generate\_config.py**: Interactive script to create JSON configuration files for training or prediction.
+* **main.py**: Entry point to train, test, or predict using the generated configuration.
+
+## Installation
+
+0. **Install uv**:
+Follow the official guide at [https://docs.astral.sh/uv/](https://docs.astral.sh/uv/)
+
+    **Linux / macOS**
+
+    ```bash
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    ```
+
+    **Windows**
+
+    ```powershell
+    powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
+    ```
+
+    ```bash
+    uv --version
+    ```
+
+1. **Clone the repository**:
+
+   ```bash
+   git clone https://git.ai.infran.ru/ilyukhin/model-v
+   cd model-v
+   ```
+2. **Install dependencies**:
+
+   ```bash
+   uv sync
+   ```
+
+## Dataset Structure
+
+Your data directory must follow this hierarchy:
+
+```
+path_to_data_folder/
+├── images/        # Input images
+│   ├── img1.tif
+│   ├── img2.tif
+│   └── ...
+└── masks/         # Ground-truth instance masks
+    ├── mask1.tif
+    ├── mask2.tif
+    └── ...
+```
+
+If your dataset contains multiple classes (e.g., class A and B) and you prefer not to duplicate images, you can organize masks into class-specific subdirectories:
+
+```
+path_to_data_folder/
+├── images/
+│   └── img1.tif
+└── masks/
+    ├── A/      # Masks for class A
+    │   ├── img1_mask.tif
+    │   └── ...
+    └── B/      # Masks for class B
+        ├── img1_mask.tif
+        └── ...
+```
+
+In this case, set the `masks_subdir` field in your dataset configuration to the name of the mask subdirectory (e.g., `"A"` or `"B"`).
+
+**Mask format**: Instance masks should be provided for multi-label segmentation with channel-last ordering, i.e., each mask array must have shape `(H, W, C)`.
+
+## generate\_config.py
+
+This script guides you through creating a JSON configuration for either training or prediction.
+
+### Usage
+
+```bash
+python generate_config.py
+```
+
+1. **Training mode?** Select `y` or `n`.
+2. **Model selection**: Choose from available models in the registry.
+3. **(If training)**
+
+   * Criterion selection
+   * Optimizer selection
+   * Scheduler selection
+4. Configuration is saved under `config/templates/train/` or `config/templates/predict/` with a unique filename.
+
+Generated config includes sections:
+
+* `model`: Model component and parameters
+* `dataset_config`: Paths, training flag, and mask subdirectory (if any)
+* `wandb_config`: Weights & Biases integration settings
+* *(If training)* `criterion`, `optimizer`, `scheduler`
+
+## main.py
+
+Entrypoint to run training, testing, or prediction using a config file.
+
+### Command-line Arguments
+
+```bash
+python main.py [-c CONFIG] [-m {train,test,predict}] [--no-save-masks] [--only-masks]
+```
+
+* `-c, --config` : Path to JSON config file (default: `config/templates/train/...json`).
+* `-m, --mode`   : `train`, `test`, or `predict` (default: `train`).
+* `--no-save-masks` : Disable saving predicted masks.
+* `--only-masks`    : Save only raw predicted masks (no visual overlays).
+
+### Workflow
+
+1. **Load config** and verify mode consistency.
+2. **Initialize** Weights & Biases if enabled.
+3. **Create** `CellSegmentator` and dataloaders with appropriate transforms.
+4. **Print** dataset info for the first batch.
+5. **Run** training or inference (`.run()`).
+6. **Save** model checkpoint and upload to W\&B if in training mode.
+
+## Configurable Parameters
+
+A brief overview of the key parameters you can adjust in your JSON config:
+
+### Common Settings (`common`)
+
+* `seed` (int): Random seed for data splitting and reproducibility (default: `0`).
+* `device` (str): Compute device to use, e.g., `'cuda:0'` or `'cpu'` (default: `'cuda:0'`).
+* `use_amp` (bool): Enable Automatic Mixed Precision for faster training (default: `false`).
+* `masks_subdir` (str): Name of subdirectory under `masks/` containing the instance masks (default: `""`).
+* `predictions_dir` (str): Output directory for saving predicted masks (default: `"."`).
+* `pretrained_weights` (str): Path to pretrained model weights (default: `""`).
+
+### Training Settings (`training`)
+
+* `is_split` (bool): Whether your data is already split (`true`) or needs splitting (`false`, default).
+* `split` / `pre_split`: Directories for data when pre-split or unsplit.
+* `train_size`, `valid_size`, `test_size` (int/float): Size or ratio of your splits (e.g., `0.7`, `0.1`, `0.2`).
+* `batch_size` (int): Number of samples per training batch (default: `1`).
+* `num_epochs` (int): Total training epochs (default: `100`).
+* `val_freq` (int): Frequency (in epochs) to run validation (default: `1`).
+
+### Testing Settings (`testing`)
+
+* `test_dir` (str): Directory containing test data (default: `"."`).
+* `test_size` (int/float): Portion or count of data for testing (default: `1.0`).
+* `shuffle` (bool): Shuffle test data before evaluation (default: `true`).
+
+> **Batch size note:** Validation, testing, and prediction runs always use a batch size of `1`, regardless of the `batch_size` setting in the training configuration.
+
+## Examples
+
+### Generate a training config
+
+```bash
+python generate_config.py
+# Follow prompts to select model, criterion, optimizer, scheduler
+# Output saved to config/templates/train/YourConfig.json
+```
+
+### Train a model
+
+```bash
+python main.py -c config/templates/train/YourConfig.json -m train
+```
+
+### Predict on new data
+
+```bash
+python main.py -c config/templates/predict/YourConfig.json -m predict
+```
+
+---

core/models/model_v.py

@@ -64,7 +64,7 @@ class ModelV(MAnet):
 
         # Encode the input and then decode it
         features = self.encoder(x)
-        decoder_output = self.decoder(*features)
+        decoder_output = self.decoder(features)
 
         # Generate masks for cell probability and gradient flows
         cellprob_mask = self.cellprob_head(decoder_output)

core/segmentator.py

@@ -943,19 +943,20 @@ class CellSegmentator:
                             predicted_masks=preds,
                             ground_truth_masks=labels_post, # type: ignore
                             iou_threshold=0.5,
-                            return_error_masks=(mode == "test") and save_results is True
+                            return_error_masks=(mode == "test") and save_results is True and not only_masks
                         )
                         all_tp.append(tp)
                         all_fp.append(fp)
                         all_fn.append(fn)
                         
                         if mode == "test" and save_results is True:
+                            masks = (tp_masks, fp_masks, fn_masks) if not only_masks else None
                             self.__save_prediction_masks(
                                 sample=batch,
                                 predicted_mask=preds,
                                 start_index=batch_counter,
                                 only_masks=only_masks,
-                                masks=(tp_masks, fp_masks, fn_masks) # type: ignore
+                                masks=masks # type: ignore
                             )
 
                 # Backpropagation and optimizer step in training

pyproject.toml

@@ -0,0 +1,37 @@
+[project]
+name = "model-v"
+version = "1.0.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "colorlog>=6.9.0",
+    "fastremap>=1.16.1",
+    "fill-voids>=2.0.8",
+    "imagecodecs>=2025.3.30",
+    "matplotlib>=3.10.3",
+    "monai>=1.4.0",
+    "numba>=0.61.2",
+    "numpy>=1.26.4",
+    "pydantic>=2.11.4",
+    "scikit-image>=0.25.2",
+    "scipy>=1.15.3",
+    "segmentation-models-pytorch>=0.5.0",
+    "tabulate>=0.9.0",
+    "tifffile>=2025.3.30",
+    "torch==2.5.1+cu124",
+    "torchaudio==2.5.1+cu124",
+    "torchvision==0.20.1+cu124",
+    "tqdm>=4.67.1",
+    "wandb>=0.19.11",
+]
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu124" }
+torchvision = { index = "pytorch-cu124" }
+torchaudio = { index = "pytorch-cu124" }
+
+[[tool.uv.index]]
+name = "pytorch-cu124"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true