DiSTNet2D 0.2.2__tar.gz → 0.2.4__tar.gz

{distnet2d-0.2.2 → distnet2d-0.2.4}/DiSTNet2D.egg-info/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: DiSTNet2D
-Version: 0.2.2
+Version: 0.2.4
 Summary: tensorflow/keras implementation of DiSTNet 2D
 Home-page: https://github.com/jeanollion/distnet2d
-Download-URL: https://github.com/jeanollion/distnet2d/releases/download/v0.2.2/distnet2d-0.2.2.tar.gz
+Download-URL: https://github.com/jeanollion/distnet2d/releases/download/v0.2.4/distnet2d-0.2.4.tar.gz
 Author: Jean Ollion
 Author-email: jean.ollion@sabilab.fr
 Keywords: Segmentation,Tracking,Cell,Tensorflow,Keras
@@ -22,7 +22,7 @@ Requires-Dist: tensorflow>=2.7.1
 Requires-Dist: edt>=2.0.2
 Requires-Dist: scikit-fmm
 Requires-Dist: numba
-Requires-Dist: dataset_iterator>=0.5.5
+Requires-Dist: dataset_iterator>=0.5.7
 Requires-Dist: elasticdeform>=0.4.7
 Dynamic: author
 Dynamic: author-email

{distnet2d-0.2.2 → distnet2d-0.2.4}/DiSTNet2D.egg-info/requires.txt

@@ -4,5 +4,5 @@ tensorflow>=2.7.1
 edt>=2.0.2
 scikit-fmm
 numba
-dataset_iterator>=0.5.5
+dataset_iterator>=0.5.7
 elasticdeform>=0.4.7

{distnet2d-0.2.2 → distnet2d-0.2.4}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: DiSTNet2D
-Version: 0.2.2
+Version: 0.2.4
 Summary: tensorflow/keras implementation of DiSTNet 2D
 Home-page: https://github.com/jeanollion/distnet2d
-Download-URL: https://github.com/jeanollion/distnet2d/releases/download/v0.2.2/distnet2d-0.2.2.tar.gz
+Download-URL: https://github.com/jeanollion/distnet2d/releases/download/v0.2.4/distnet2d-0.2.4.tar.gz
 Author: Jean Ollion
 Author-email: jean.ollion@sabilab.fr
 Keywords: Segmentation,Tracking,Cell,Tensorflow,Keras
@@ -22,7 +22,7 @@ Requires-Dist: tensorflow>=2.7.1
 Requires-Dist: edt>=2.0.2
 Requires-Dist: scikit-fmm
 Requires-Dist: numba
-Requires-Dist: dataset_iterator>=0.5.5
+Requires-Dist: dataset_iterator>=0.5.7
 Requires-Dist: elasticdeform>=0.4.7
 Dynamic: author
 Dynamic: author-email

{distnet2d-0.2.2 → distnet2d-0.2.4}/distnet_2d/model/distnet_2d.py

@@ -1,3 +1,5 @@
+import contextlib
+
 import tensorflow as tf
 from .layers import ker_size_to_string, Combine, ResConv2D, Conv2DBNDrop, Conv2DTransposeBNDrop, WSConv2D, \
     BatchToChannel, SplitBatch, ChannelToBatch, NConvToBatch2D, SelectFeature, StopGradient, Stack
@@ -31,6 +33,7 @@ class DiSTNetModel(tf.keras.Model):
                  category_number:int=0, category_class_weights = None, category_max_class_weight=10,
                  print_gradients:bool=False,  # for optimization, available in eager mode only
                  accum_steps=1, use_agc=False, agc_clip_factor=0.1, agc_eps=1e-3, agc_exclude_output=False,  # lower clip factor clips more
+                 perform_test_step:bool=False,
                  **kwargs):
         super().__init__(*args, **kwargs)
         self.edm_weight = edm_loss_weight
@@ -87,41 +90,40 @@ class DiSTNetModel(tf.keras.Model):
         self.displacement_loss.reduction = tf.keras.losses.Reduction.NONE
 
         # metrics associated to losses for to display accurate loss in a distributed setting
-        self.edm_loss_metric = tf.keras.metrics.Mean(name="EDM")
-        self.center_loss_metric = tf.keras.metrics.Mean(name="CDM")
-        self.category_loss_metric = tf.keras.metrics.Mean(name="category") if self.category_weight > 0 else None
-        self.dx_loss_metric = tf.keras.metrics.Mean(name="dX")
-        self.dy_loss_metric = tf.keras.metrics.Mean(name="dY")
-        self.link_multiplicity_loss_metric = tf.keras.metrics.Mean(name="link_multiplicity")
-        self.loss_metric = tf.keras.metrics.Mean(name="loss")
-
-    @property
-    def metrics(self):
-        metrics = []
         if self.edm_weight > 0:
-            metrics.append(self.edm_loss_metric)
+            self.edm_loss_metric = tf.keras.metrics.Mean(name="EDM")
         if self.center_weight > 0:
-            metrics.append(self.center_loss_metric)
+            self.center_loss_metric = tf.keras.metrics.Mean(name="CDM")
+        if self.category_weight > 0 and self.category_number > 1:
+            self.category_loss_metric = tf.keras.metrics.Mean(name="category")
         if self.displacement_weight > 0:
-            metrics.append(self.dy_loss_metric)
-            metrics.append(self.dx_loss_metric)
+            self.dx_loss_metric = tf.keras.metrics.Mean(name="dX")
+            self.dy_loss_metric = tf.keras.metrics.Mean(name="dY")
         if self.link_multiplicity_weight > 0:
-            metrics.append(self.link_multiplicity_loss_metric)
-        if self.category_weight > 0:
-            metrics.append(self.category_loss_metric)
-        metrics.append(self.loss_metric)
-        if self._is_compiled:
-            if self.compiled_metrics is not None:
-                metrics += self.compiled_metrics.metrics
+            self.link_multiplicity_loss_metric = tf.keras.metrics.Mean(name="link_multiplicity")
+        self.loss_metric = tf.keras.metrics.Mean(name="loss")
+        self.perform_test_step=perform_test_step
 
-        for l in self._flatten_layers():
-            metrics.extend(l._metrics)
 
-        return metrics
+    @staticmethod
+    @contextlib.contextmanager
+    def nullcontext():
+        yield
+
+    def maybe_gradient_tape(self, training):
+        if training:
+            return tf.GradientTape(persistent=self.print_gradients)
+        return self.nullcontext()
 
 
     def train_step(self, data):
-        if self.use_grad_acc:
+        return self.step(data, training=True)
+
+    def test_step(self, data):
+        return self.step(data, False)
+
+    def step(self, data, training:bool):
+        if self.use_grad_acc and training:
             self.gradient_accumulator.init_train_step()
 
         fw = self.frame_window
@@ -144,7 +146,7 @@ class DiSTNetModel(tf.keras.Model):
         lm_idx = int(center_weight > 0) + 2 * int(displacement_weight > 0) + int(link_multiplicity_weight > 0)
         cat_idx = lm_idx + 1
 
-        with tf.GradientTape(persistent=self.print_gradients) as tape:
+        with self.maybe_gradient_tape(training) as tape:
             y_pred = self(x, training=True)  # Forward pass
             if self.predict_edm_derivatives:
                 edm, edm_dy, edm_dx = tf.split(y_pred[0], num_or_size_splits=3, axis=-1)
@@ -225,7 +227,7 @@ class DiSTNetModel(tf.keras.Model):
             if num_replicas > 1:
                 loss *= 1.0 / num_replicas
 
-        if self.print_gradients:
+        if self.print_gradients and training:
             trainable_vars_tape = [t for t in self.trainable_variables if (t.name.startswith("DecoderSegEDM") or t.name.startswith("DecoderCenterCDM") or t.name.startswith("DecoderTrackY0") or t.name.startswith("DecoderTrackX0") or t.name.startswith("DecoderLinkMultiplicity0") or t.name.startswith("FeatureSequence_Op4") or t.name.startswith("Attention")) and ("/kernel" in t.name or "/wv" in t.name) ]
             for loss_name, loss_value in losses.items():
                 if loss_name != "loss" :
@@ -246,31 +248,31 @@ class DiSTNetModel(tf.keras.Model):
                                 print(f"AGC: layer: {v.name}, loss: {loss_name}, grad: {tf.math.sqrt(tf.reduce_mean(tf.math.square(g))).numpy()}")
 
         # Compute gradients
-        gradients = tape.gradient(loss, self.trainable_variables)
-        if mixed_precision:
-            gradients = self.optimizer.get_unscaled_gradients(gradients)
-        if self.use_agc:
-            gradients = adaptive_clip_grad(self.trainable_variables, gradients, clip_factor=self.agc_clip_factor, eps=self.agc_eps, exclude_keywords=self.agc_exclude_keywords)
-        if not self.use_grad_acc:
-            self.optimizer.apply_gradients(zip(gradients, self.trainable_variables)) #Update weights
-        else:
-            self.gradient_accumulator.accumulate_gradients(gradients)
-            self.gradient_accumulator.apply_gradients()
+        if training:
+            gradients = tape.gradient(loss, self.trainable_variables)
+            if mixed_precision:
+                gradients = self.optimizer.get_unscaled_gradients(gradients)
+            if self.use_agc:
+                gradients = adaptive_clip_grad(self.trainable_variables, gradients, clip_factor=self.agc_clip_factor, eps=self.agc_eps, exclude_keywords=self.agc_exclude_keywords)
+            if not self.use_grad_acc:
+                self.optimizer.apply_gradients(zip(gradients, self.trainable_variables)) #Update weights
+            else:
+                self.gradient_accumulator.accumulate_gradients(gradients)
+                self.gradient_accumulator.apply_gradients()
 
         # Update metrics state
-        if edm_weight > 0:
+        if self.edm_weight > 0:
             self.edm_loss_metric.update_state(losses["EDM"], sample_weight=batch_dim)
-        if center_weight > 0:
+        if self.center_weight > 0:
             self.center_loss_metric.update_state(losses["CDM"], sample_weight=batch_dim)
-        if displacement_weight>0:
+        if self.displacement_weight > 0:
             self.dx_loss_metric.update_state(losses["dX"], sample_weight=batch_dim)
             self.dy_loss_metric.update_state(losses["dY"], sample_weight=batch_dim)
-        if link_multiplicity_weight >0:
+        if self.link_multiplicity_weight > 0:
             self.link_multiplicity_loss_metric.update_state(losses["link_multiplicity"], sample_weight=batch_dim)
-        if self.category_loss_metric is not None:
+        if self.category_weight > 0 and self.category_number > 1:
             self.category_loss_metric.update_state(losses["category"], sample_weight=batch_dim)
         self.loss_metric.update_state(losses["loss"], sample_weight=batch_dim)
-
         return self.compute_metrics(x, y, y_pred, None)
 
     def _compute_displacement_loss(self, y, y_pred, cell_mask):

{distnet2d-0.2.2 → distnet2d-0.2.4}/distnet_2d/utils/metrics_tf.py

@@ -1,16 +1,16 @@
 import tensorflow as tf
 from .objectwise_computation_tf import get_max_by_object_fun, coord_distance_fun, get_argmax_2d_by_object_fun, \
     get_mean_by_object_fun, get_label_size, IoU, objectwise_compute, objectwise_compute_channel, reduce_pop_size, \
-    FPR
+    FP
 
 
-def get_metrics_fun(center_scale: float, max_objects_number: int = 0, category:bool = False, tracking:bool=True):
+def get_metrics_fun(scale: float, max_objects_number: int = 0, category:bool = False, tracking:bool=True):
     """
     return metric function for disnet2D
     assumes iterator in return_central_only= True mode (thus framewindow = 1 and next = true)
     Parameters
     ----------
-    center_scale
+    scale
     max_objects_number
     reduce
 
@@ -19,7 +19,6 @@ def get_metrics_fun(center_scale: float, max_objects_number: int = 0, category:b
 
     """
 
-    scale = tf.cast(center_scale, tf.float32)
     coord_distance_function = coord_distance_fun(max=True, sqrt=True, pop_fraction=0.25)
     spa_max_fun = get_argmax_2d_by_object_fun()
     mean_fun = get_mean_by_object_fun()
@@ -41,7 +40,7 @@ def get_metrics_fun(center_scale: float, max_objects_number: int = 0, category:b
         labels = tf.transpose(labels, perm=[2, 0, 1])  # (1, Y, X)
         edm = tf.transpose(edm, perm=[2, 0, 1]) # (1, Y, X)
         gdcm = tf.transpose(gdcm, perm=[2, 0, 1])  # (1, Y, X)
-        center_values = tf.math.exp(-tf.math.square(tf.math.divide(gdcm, scale)))
+        center_values = tf.math.exp(-tf.math.square(tf.math.divide(gdcm, tf.cast(scale/2., tf.float32))))
         ids, sizes, N = get_label_size(labels, max_objects_number)  # (1, N), (1, N)
         ids = ids[0]
         sizes = sizes[0]
@@ -61,14 +60,14 @@ def get_metrics_fun(center_scale: float, max_objects_number: int = 0, category:b
         metrics = []
 
         # EDM : foreground/background IoU
-        pred_foreground = tf.math.greater(edm, tf.cast(0.5, edm.dtype))
+        pred_foreground = tf.math.greater(edm, tf.cast(0, edm.dtype))
         true_foreground = tf.math.greater(labels, tf.cast(0, labels.dtype))
-        edm_IoU = IoU(true_foreground, pred_foreground, tolerance=False)
+        edm_IoU = IoU(true_foreground, pred_foreground, tolerance_radius=scale / 8.) #
         metrics.append(edm_IoU)
 
-        # Surface-based False Positive Rate (FPR) based on EDM
-        #fpr = FPR(true_foreground, pred_foreground, tolerance=True)
-        #metrics.append(-fpr)
+        # Surface-based False Positive Density (FPD) based on EDM
+        fp = FP(true_foreground, pred_foreground, rate=False, tolerance_radius=scale / 4.) #
+        metrics.append(-fp)
 
         # contour IoU : problem: true positive contours are usually not precise enough.
         #pred_contours = tf.math.logical_and(tf.math.greater(edm, tf.cast(0.5, edm.dtype)), tf.math.less_equal(edm, tf.cast(1.5, edm.dtype)))

{distnet2d-0.2.2 → distnet2d-0.2.4}/distnet_2d/utils/objectwise_computation_tf.py

@@ -211,32 +211,71 @@ def _generate_kernel(sizeY, sizeX, C=1, O=0):
     return kernel
 
 
-def IoU(true_foreground, pred_foreground, tolerance:bool=False):
-    true_inter = _dilate_mask(true_foreground) if tolerance else true_foreground
+def IoU(true_foreground, pred_foreground, tolerance_radius:float=0):
+    true_inter = _dilate_mask(true_foreground, radius=tolerance_radius, symmetric_padding=True) if tolerance_radius>=1 else true_foreground
     intersection = tf.math.count_nonzero(tf.math.logical_and(true_inter, pred_foreground), keepdims=False)
     union = tf.math.count_nonzero(tf.math.logical_or(true_foreground, pred_foreground), keepdims=False)
     return tf.cond(tf.math.equal(union, tf.cast(0, union.dtype)), lambda: tf.cast(1., tf.float32), lambda: tf.math.divide(tf.cast(intersection, tf.float32), tf.cast(union, tf.float32)))  # if union is null -> metric is 1
 
 
-def FPR(true_foreground, pred_foreground, tolerance:bool=False):
+def FP(true_foreground, pred_foreground, rate:bool = False, tolerance_radius:float=0):
     true_background = tf.math.logical_not(true_foreground)
     false_positives = tf.logical_and(pred_foreground, true_background)
-    false_positives = _erode_mask(false_positives) if tolerance else false_positives
-    num_fp = tf.reduce_sum(tf.cast(false_positives, tf.float32))
-    num_tn = tf.reduce_sum(tf.cast(true_background, tf.float32))
-    return tf.math.divide_no_nan(num_fp, num_tn)
+    false_positives = _erode_mask(false_positives, radius=tolerance_radius, symmetric_padding=False) if tolerance_radius>=1 else false_positives
+    fp = tf.math.count_nonzero(false_positives, keepdims=False)
+    if rate: # FPR
+        tn = tf.math.count_nonzero(true_background, keepdims=False) # for FRP
+        return tf.math.divide_no_nan(tf.cast(fp, tf.float32), tf.cast(tn, tf.float32))
+    else: # FPD
+        #return tf.cast(fp, tf.float32)
+        npix = tf.reduce_prod(tf.shape(true_background)) # for FPD
+        return tf.math.divide(tf.cast(fp, tf.float32), tf.cast(npix, tf.float32))
+
+
+def _dilate_mask(maskBYX, radius:float=1.5, tolerance:float=0.25, symmetric_padding:bool=True):
+    assert 0<=tolerance<0.5
+    maskBYX = tf.cast(maskBYX, tf.int32)
+    ker, rad = circular_kernel(radius)
+    thld = tf.math.floor(tf.cast(tf.math.reduce_sum(ker), tf.float32) * tf.cast(tolerance, tf.float32))
+    conv = _convolve(maskBYX, ker, rad, symmetric_padding=symmetric_padding)
+    return tf.math.greater(conv, tf.cast(thld, tf.int32))
 
 
-def _dilate_mask(maskBYX):
+def _erode_mask(maskBYX, radius:float=1.5, tolerance:float=0.25, symmetric_padding:bool=False):
+    assert 0 <= tolerance < 0.5
     maskBYX = tf.cast(maskBYX, tf.int32)
-    conv = _convolve(maskBYX, tf.ones(shape=[3, 3], dtype=tf.int32))
-    return tf.math.greater(conv, tf.cast(2, tf.int32))
+    ker, rad = circular_kernel(radius)
+    thld = tf.math.ceil(tf.cast(tf.math.reduce_sum(ker), tf.float32) * tf.cast(1 - tolerance, tf.float32))
+    conv = _convolve(maskBYX, ker, rad, symmetric_padding=symmetric_padding)
+    return tf.math.greater_equal(conv, tf.cast(thld, tf.int32))
 
 
-def _erode_mask(maskBYX):
-    maskBYX = tf.cast(maskBYX, tf.int32)
-    conv = _convolve(maskBYX, tf.ones(shape=[3, 3], dtype=tf.int32))
-    return tf.math.greater_equal(conv, tf.cast(7, tf.int32))
+
+def circular_kernel(radius: float) :
+    """
+    Create a circular 2D kernel of ones with a given float radius.
+    Args:
+        radius: The radius of the circle (float).
+    Returns:
+        A 2D TensorFlow tensor representing the circular kernel (dtype: tf.int32).
+    """
+    radius_int = tf.cast(radius, tf.int32)
+    diameter = tf.cast(2 * radius_int + 1, tf.int32)
+    center = diameter // 2
+
+    # Create a grid of coordinates
+    y = tf.range(-center, diameter - center, dtype=tf.float32)
+    x = tf.range(-center, diameter - center, dtype=tf.float32)
+    y_grid, x_grid = tf.meshgrid(y, x, indexing='ij')
+
+    # Compute the distance from the center
+    distance = tf.math.sqrt(x_grid**2 + y_grid**2)
+
+    # Create the circular kernel
+    kernel = tf.zeros((diameter, diameter), dtype=tf.int32)
+    kernel = tf.where(distance <= radius, tf.ones_like(kernel, dtype=tf.int32), kernel)
+
+    return kernel, radius_int
 
 
 def _contour_IoU_fun(pred_contour, mask, size):
@@ -245,13 +284,15 @@ def _contour_IoU_fun(pred_contour, mask, size):
 
 
 def _compute_contours(maskBYX):
-    kernel = [[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]]
-    conv = _convolve(maskBYX, kernel)
-    return tf.math.greater(conv, 0) # detect at least one zero in the neighborhood
+    kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]])
+    conv = _convolve(maskBYX, kernel, 1, symmetric_padding=True)
+    return tf.math.greater(conv, tf.cast(0, conv.dtype)) # detect at least one zero in the neighborhood
 
 
-def _convolve(imageBYX, kernel):
-    padded = tf.pad(imageBYX, [[0, 0], [1, 1], [1, 1]], 'SYMMETRIC')
-    input = padded[..., tf.newaxis]
-    conv = tf.nn.conv2d(input, kernel[:, :, tf.newaxis, tf.newaxis], strides=1, padding='VALID')
+def _convolve(imageBYX, kernel, radius, symmetric_padding:bool):
+    if symmetric_padding:
+        imageBYX = tf.pad(imageBYX, [[0, 0], [radius, radius], [radius, radius]], 'SYMMETRIC')
+    imageBYX = imageBYX[..., tf.newaxis]
+    kernel = tf.cast(kernel, imageBYX.dtype)
+    conv = tf.nn.conv2d(imageBYX, kernel[:, :, tf.newaxis, tf.newaxis], strides=1, padding='VALID' if symmetric_padding else "SAME")
     return conv[..., 0]

{distnet2d-0.2.2 → distnet2d-0.2.4}/setup.py

@@ -5,14 +5,14 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="DiSTNet2D",
-    version="0.2.2",
+    version="0.2.4",
     author="Jean Ollion",
     author_email="jean.ollion@sabilab.fr",
     description="tensorflow/keras implementation of DiSTNet 2D",
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/jeanollion/distnet2d",
-    download_url='https://github.com/jeanollion/distnet2d/releases/download/v0.2.2/distnet2d-0.2.2.tar.gz',
+    download_url='https://github.com/jeanollion/distnet2d/releases/download/v0.2.4/distnet2d-0.2.4.tar.gz',
     packages=setuptools.find_packages(),
     keywords=['Segmentation', 'Tracking', 'Cell', 'Tensorflow', 'Keras'],
     classifiers=[
@@ -24,5 +24,5 @@ setuptools.setup(
         'Programming Language :: Python :: 3',
     ],
     python_requires='>=3',
-    install_requires=['numpy', 'scipy', 'tensorflow>=2.7.1', 'edt>=2.0.2', 'scikit-fmm', 'numba', 'dataset_iterator>=0.5.5', 'elasticdeform>=0.4.7']
+    install_requires=['numpy', 'scipy', 'tensorflow>=2.7.1', 'edt>=2.0.2', 'scikit-fmm', 'numba', 'dataset_iterator>=0.5.7', 'elasticdeform>=0.4.7']
 )