fix: offset handling

**Bug fix**
Fixed an issue that occurred when specifying a floating-point value for `offset` during testing.

**New**
Added new configuration parameters for controlling gradient flow calculations, as well as a parameter defining the IoU threshold used for metric computation.

README.md

@@ -150,11 +150,19 @@ A brief overview of the key parameters you can adjust in your JSON config:
 * `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`).
 * `roi_size` (int): Defines the size of the square Region of Interest (ROI) used for cropping during training. This same size is also applied for the sliding window inference during validation and testing (default: `512`).
+* `iou_threshold` (float): Intersection over Union threshold used for metric computation. All detection and segmentation metrics are calculated based on this IoU value (default: `0.5`).
 * `remove_boundary_objects` (bool): Flag to remove boundary objects when testing (default: `True`).
 * `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: `""`).
 
+### Gradient Flow Settings (`gradient_flow`)
+
+* `prob_threshold` (float): Probability threshold for binarizing model outputs into masks. Pixels with probability values above this threshold are considered part of an object (default: `0.5`).
+* `flow_threshold` (float): Threshold for filtering unreliable flow vectors during instance mask reconstruction. Lower values allow more relaxed flow matching (default: `0.4`).
+* `num_iters` (int): Number of iterations used when following the flow field to reconstruct object instances (default: `200`).
+* `min_object_size` (int): Minimum area (in pixels) to keep an instance. Smaller regions are discarded as noise (default: `15`).
+
 ### Training Settings (`training`)
 
 * `is_split` (bool): Whether your data is already split (`true`) or needs splitting (`false`, default).

config/dataset_config.py

@@ -3,6 +3,30 @@ from typing import Any
 import os
 
 
+class GradientFlowConfig(BaseModel):
+    """
+    Configuration related to gradient flow and instance postprocessing.
+    """
+    prob_threshold: float = 0.5     # Threshold to binarize probability map
+    flow_threshold: float = 0.4     # Threshold for filtering bad flow masks
+    num_iters: int = 200            # Number of iterations for flow-following
+    min_object_size: int = 15       # Minimum area to keep small instances
+
+    @model_validator(mode="after")
+    def validate_gradient_flow(self) -> "GradientFlowConfig":
+        """
+        Validates gradient flow configuration values.
+        """
+        if not (0.0 <= self.prob_threshold <= 1.0):
+            raise ValueError("prob_threshold must be between 0 and 1")
+        if not (0.0 <= self.flow_threshold <= 1.0):
+            raise ValueError("flow_threshold must be between 0 and 1")
+        if self.num_iters <= 0:
+            raise ValueError("num_iters must be > 0")
+        if self.min_object_size <= 0:
+            raise ValueError("min_object_size must be > 0")
+        return self
+
 class DatasetCommonConfig(BaseModel):
     """
     Common configuration fields shared by both training and testing.
@@ -11,6 +35,8 @@ class DatasetCommonConfig(BaseModel):
     device: str = "cuda:0"          # Device used for training/testing (e.g., 'cpu' or 'cuda')
     use_amp: bool = True            # Flag to use Automatic Mixed Precision (AMP)
     roi_size: int = 512             # The size of the square window for cropping
+    iou_threshold: float = 0.5      # Threshold for calculating iou (all other metrics use iou for calculations)
+    gradient_flow: GradientFlowConfig = GradientFlowConfig()
     remove_boundary_objects: bool = True    # Flag to remove boundary objects when testing
     masks_subdir: str = ""          # Subdirectory where the required masks are located, e.g. 'masks/cars'
     predictions_dir: str = "."      # Directory to save predictions

core/segmentator.py

@@ -256,13 +256,20 @@ class CellSegmentator:
             test_images = os.path.join(self._dataset_setup.testing.test_dir, 'images')
             test_masks = os.path.join(self._dataset_setup.testing.test_dir, 'masks', self._dataset_setup.common.masks_subdir)
 
+            number_of_images = len(os.listdir(test_images))
+            test_offset = (
+                self._dataset_setup.training.test_offset
+                if isinstance(self._dataset_setup.training.test_offset, int)
+                else int(number_of_images * self._dataset_setup.training.test_offset)
+            )
+
             if test_transforms is not None:
                 test_dataset = self.__get_dataset(
                     images_dir=test_images,
                     masks_dir=test_masks,
                     transforms=test_transforms,
                     size=self._dataset_setup.testing.test_size,
-                    offset=self._dataset_setup.testing.test_offset,
+                    offset=test_offset,
                     shuffle=self._dataset_setup.testing.shuffle
                 )
                 self._test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)
@@ -274,7 +281,7 @@ class CellSegmentator:
                     masks_dir=None,
                     transforms=predict_transforms,
                     size=self._dataset_setup.testing.test_size,
-                    offset=self._dataset_setup.testing.test_offset,
+                    offset=test_offset,
                     shuffle=self._dataset_setup.testing.shuffle
                 )
                 self._predict_dataloader = DataLoader(predict_dataset, batch_size=1, shuffle=False)
@@ -982,7 +989,7 @@ class CellSegmentator:
                         tp, fp, fn, tp_masks, fp_masks, fn_masks = self.__compute_stats(
                             predicted_masks=preds,
                             ground_truth_masks=labels_post, # type: ignore
-                            iou_threshold=0.5,
+                            iou_threshold=self._dataset_setup.common.iou_threshold,
                             return_error_masks=(mode == "test") and save_results is True and not only_masks
                         )
                         all_tp.append(tp)
@@ -1120,9 +1127,10 @@ class CellSegmentator:
             instance_masks[idx] = self.__segment_instances(
                 probability_map=probabilities[idx],
                 flow=gradflow[idx],
-                prob_threshold=0.5,
-                flow_threshold=0.4,
-                min_object_size=15
+                prob_threshold=self._dataset_setup.common.gradient_flow.prob_threshold,
+                flow_threshold=self._dataset_setup.common.gradient_flow.flow_threshold,
+                num_iters=self._dataset_setup.common.gradient_flow.num_iters,
+                min_object_size=self._dataset_setup.common.gradient_flow.min_object_size
             )
 
         # Convert ground truth to numpy
@@ -2194,7 +2202,7 @@ class CellSegmentator:
         self,
         probability_map: np.ndarray,
         flow: np.ndarray,
-        prob_threshold: float = 0.0,
+        prob_threshold: float = 0.5,
         flow_threshold: float = 0.4,
         num_iters: int = 200,
         min_object_size: int = 15
@@ -2205,7 +2213,7 @@ class CellSegmentator:
         Args:
             probability_map: 3D array `(C, H, W)` of cell probabilities.
             flow: 3D array `(2*C, H, W)` of forward flow vectors.
-            prob_threshold: threshold to binarize probability_map. (Default 0.0)
+            prob_threshold: threshold to binarize probability_map. (Default 0.5)
             flow_threshold: threshold for filtering bad flow masks. (Default 0.4)
             num_iters: number of iterations for flow-following. (Default 200)
             min_object_size: minimum area to keep small instances. (Default 15)