From 7a68b3233c46d7a9cd91e8a9ada09a8273d1229d Mon Sep 17 00:00:00 2001 From: laynholt Date: Tue, 14 Oct 2025 15:15:16 +0000 Subject: [PATCH] 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 | 8 ++++++++ config/dataset_config.py | 26 ++++++++++++++++++++++++++ core/segmentator.py | 24 ++++++++++++++++-------- 3 files changed, 50 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 2a319aa..746cb3f 100644 --- a/README.md +++ b/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). diff --git a/config/dataset_config.py b/config/dataset_config.py index 0b2ac7a..af9282a 100644 --- a/config/dataset_config.py +++ b/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 diff --git a/core/segmentator.py b/core/segmentator.py index c390f43..54812b8 100644 --- a/core/segmentator.py +++ b/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)