core\refac: #98 full hp (scalar)

- full hps availability for scalar experiment
- added missing deps for optuna meaningful hps plotting

Author: blotero

.gitignore

@@ -11,4 +11,5 @@ core/dist
 __data__
 **/results/**
 tuner_results/
-**/vis/
+**/vis/
+**/optuna_results/

core/poetry.lock

@@ -1,7 +1,7 @@
 [project]
 description = "Framework for handling image segmentation in the context of multiple annotators"
 name = "seg_tgce"
-version = "0.3.4"
+version = "0.3.7"
 readme = "README.md"
 authors = [{ name = "Brandon Lotero", email = "blotero@gmail.com" }]
 maintainers = [{ name = "Brandon Lotero", email = "blotero@gmail.com" }]
@@ -15,7 +15,7 @@ Issues = "https://github.com/blotero/seg_tgce/issues"
 
 [tool.poetry]
 name = "seg_tgce"
-version = "0.3.4"
+version = "0.3.7"
 authors = ["Brandon Lotero <blotero@gmail.com>"]
 description = "A package for the SEG TGCE project"
 readme = "README.md"
@@ -44,6 +44,8 @@ albumentations = "^2.0.7"
 pandas = "^2.2.3"
 seaborn = "^0.13.2"
 optuna = "^4.4.0"
+plotly = "^6.2.0"
+kaleido = "^1.0.0"
 
 
 [tool.poetry.group.test.dependencies]

core/pyproject.toml

@@ -27,6 +27,10 @@
     "noise_tolerance": 0.5,
     "a": 0.3,
     "b": 0.7,
+    "c": 1.0,
+    "lambda_reg_weight": 0.1,
+    "lambda_entropy_weight": 0.1,
+    "lambda_sum_weight": 0.1,
 }
 
 
@@ -40,16 +44,27 @@ def build_model(hp: kt.HyperParameters | None = None) -> tf.keras.Model:
             "noise_tolerance": hp.Float(
                 "noise_tolerance", min_value=0.1, max_value=0.9, step=0.1
             ),
-            "b": hp.Float("b", min_value=0.1, max_value=1.0, step=0.1),
             "a": hp.Float("a", min_value=0.1, max_value=1.0, step=0.1),
+            "b": hp.Float("b", min_value=0.1, max_value=1.0, step=0.1),
+            "c": hp.Float("c", min_value=0.1, max_value=10.0, step=0.1),
+            "lambda_reg_weight": hp.Float(
+                "lambda_reg_weight", min_value=0.0, max_value=10.0, step=0.1
+            ),
+            "lambda_entropy_weight": hp.Float(
+                "lambda_entropy_weight", min_value=0.0, max_value=10.0, step=0.1
+            ),
         }
 
     return build_scalar_model_from_hparams(
         learning_rate=params["initial_learning_rate"],
         q=params["q"],
         noise_tolerance=params["noise_tolerance"],
-        b=params["b"],
         a=params["a"],
+        b=params["b"],
+        c=params["c"],
+        lambda_reg_weight=params["lambda_reg_weight"],
+        lambda_entropy_weight=params["lambda_entropy_weight"],
+        lambda_sum_weight=params["lambda_sum_weight"],
         num_classes=N_CLASSES,
         target_shape=TARGET_SHAPE,
         n_scorers=N_REAL_SCORERS,

core/seg_tgce/experiments/histology/scalar.py

@@ -30,6 +30,10 @@
     "noise_tolerance": 0.5,
     "a": 0.2,
     "b": 0.7,
+    "c": 1.0,
+    "lambda_reg_weight": 0.1,
+    "lambda_entropy_weight": 0.1,
+    "lambda_sum_weight": 0.1,
 }
 
 
@@ -40,7 +44,11 @@ def build_model_from_trial(trial: HpTunerTrial | None) -> Model:
             q=DEFAULT_HPARAMS["q"],
             noise_tolerance=DEFAULT_HPARAMS["noise_tolerance"],
             b=DEFAULT_HPARAMS["b"],
+            c=DEFAULT_HPARAMS["c"],
             a=DEFAULT_HPARAMS["a"],
+            lambda_reg_weight=DEFAULT_HPARAMS["lambda_reg_weight"],
+            lambda_entropy_weight=DEFAULT_HPARAMS["lambda_entropy_weight"],
+            lambda_sum_weight=DEFAULT_HPARAMS["lambda_sum_weight"],
             num_classes=NUM_CLASSES,
             target_shape=TARGET_SHAPE,
             n_scorers=NUM_SCORERS,
@@ -52,6 +60,12 @@ def build_model_from_trial(trial: HpTunerTrial | None) -> Model:
         noise_tolerance=trial.suggest_float("noise_tolerance", 0.1, 0.9, step=0.01),
         b=trial.suggest_float("b", 0.1, 1.0, step=0.01),
         a=trial.suggest_float("a", 0.1, 1.0, step=0.01),
+        c=trial.suggest_float("c", 0.1, 10.0, step=0.1),
+        lambda_reg_weight=trial.suggest_float("lambda_reg_weight", 0.0, 10.0, step=0.1),
+        lambda_entropy_weight=trial.suggest_float(
+            "lambda_entropy_weight", 0.0, 10.0, step=0.1
+        ),
+        lambda_sum_weight=trial.suggest_float("lambda_sum_weight", 0.0, 10.0, step=0.1),
         num_classes=NUM_CLASSES,
         target_shape=TARGET_SHAPE,
         n_scorers=NUM_SCORERS,

core/seg_tgce/experiments/pets/scalar.py

@@ -9,18 +9,23 @@
 
 
 def safe_divide(numerator: Tensor, denominator: Tensor, epsilon: float = 1e-8) -> Tensor:
-    """Safely divide two tensors, avoiding division by zero."""
     return tf.math.divide(
         numerator, tf.clip_by_value(denominator, epsilon, tf.reduce_max(denominator))
     )
 
 
 def safe_pow(x: Tensor, p: Tensor, epsilon: float = 1e-8) -> Tensor:
-    """Compute x^p safely by ensuring x is within a valid range."""
     return tf.pow(tf.clip_by_value(x, epsilon, 1.0 - epsilon), p)
 
 
-class TcgeScalar(Loss):
+def reliability_penalizer(
+    lms: Tensor, lambdas: Tensor, a: float, b: float, c: float
+) -> Tensor:
+    x = lambdas - lms
+    return c * tf.maximum(1 / (1 - a) * x * tf.exp((x - 1) / b), 0)
+
+
+class TgceScalar(Loss):
     """
     Truncated generalized cross entropy
     for semantic segmentation loss.
@@ -35,32 +40,40 @@ def __init__(  # pylint: disable=too-many-arguments
         noise_tolerance: float = 0.1,
         a: float = 0.7,
         b: float = 0.7,
+        c: float = 1.0,
+        lambda_reg_weight: float = 0.1,
+        lambda_entropy_weight: float = 0.1,
+        lambda_sum_weight: float = 0.1,
         epsilon: float = 1e-8,
     ) -> None:
         self.q = q
         self.num_classes = num_classes
         self.noise_tolerance = noise_tolerance
         self.a = a
         self.b = b
+        self.c = c
+        self.lambda_reg_weight = lambda_reg_weight
+        self.lambda_entropy_weight = lambda_entropy_weight
+        self.lambda_sum_weight = lambda_sum_weight
         self.epsilon = epsilon
         super().__init__(name=name)
 
-    def penalizer(self, lms: tf.Tensor, lambdas: tf.Tensor) -> tf.Tensor:
-        """Compute the penalizer term for reliability regularization."""
-        x = lambdas - lms
-        return tf.maximum(1 / (1 - self.a) * x * tf.exp((x - 1) / self.b), 0)
-
     def call(
         self,
         y_true: tf.Tensor,
         y_pred: tf.Tensor,
         lambda_r: tf.Tensor,
         labeler_mask: tf.Tensor,
     ) -> tf.Tensor:
+        # Cast inputs to target data type
+        y_true = tf.cast(y_true, TARGET_DATA_TYPE)
+        y_pred = tf.cast(y_pred, TARGET_DATA_TYPE)
+        lambda_r = tf.cast(lambda_r, TARGET_DATA_TYPE)
+
         y_pred = tf.clip_by_value(y_pred, self.epsilon, 1.0 - self.epsilon)
         lambda_r = tf.clip_by_value(lambda_r, self.epsilon, 1.0 - self.epsilon)
 
-        reg_term = self.penalizer(labeler_mask, lambda_r)
+        reg_term = reliability_penalizer(labeler_mask, lambda_r, self.a, self.b, self.c)
 
         y_pred_exp = tf.expand_dims(y_pred, axis=-1)
         y_pred_exp = tf.tile(y_pred_exp, [1, 1, 1, 1, tf.shape(y_true)[-1]])
@@ -78,7 +91,28 @@ def call(
             (1.0 - tf.pow(self.noise_tolerance, self.q)) / (self.q + self.epsilon)
         )
 
-        total_loss = tf.reduce_mean(term1 + term2) + reg_term
+        # Only compute regularization terms for valid labelers
+        valid_lambda_r = lambda_r * tf.expand_dims(tf.expand_dims(labeler_mask, 1), 1)
+        lambda_reg = self.lambda_reg_weight * tf.reduce_mean(
+            tf.square(valid_lambda_r - 0.5)
+        )
+
+        lambda_entropy = -self.lambda_entropy_weight * tf.reduce_mean(
+            valid_lambda_r * tf.math.log1p(valid_lambda_r)
+            + (1 - valid_lambda_r) * tf.math.log1p(1 - valid_lambda_r)
+        )
+
+        lambda_sum = self.lambda_sum_weight * tf.reduce_mean(
+            tf.square(tf.reduce_sum(valid_lambda_r, axis=-1) - 1.0)
+        )
+
+        total_loss = (
+            tf.reduce_mean(term1 + term2)
+            + reg_term
+            + lambda_reg
+            + lambda_entropy
+            + lambda_sum
+        )
 
         total_loss = tf.where(
             tf.math.is_nan(total_loss),
@@ -99,11 +133,14 @@ def get_config(
             **base_config,
             "q": self.q,
             "b": self.b,
+            "lambda_reg_weight": self.lambda_reg_weight,
+            "lambda_entropy_weight": self.lambda_entropy_weight,
+            "lambda_sum_weight": self.lambda_sum_weight,
             "epsilon": self.epsilon,
         }
 
 
-class TcgeFeatures(Loss):
+class TgceFeatures(Loss):
     """
     Truncated generalized cross entropy for semantic segmentation loss
     with feature-based reliability (reliability map from bottleneck features).
@@ -210,7 +247,7 @@ def get_config(
         }
 
 
-class TcgePixel(Loss):
+class TgcePixel(Loss):
     """
     Truncated generalized cross entropy for semantic segmentation loss
     with pixel-wise reliability (full resolution reliability map).

core/seg_tgce/loss/tgce.py

@@ -1,7 +1,7 @@
 from keras import Model
 from keras.optimizers import Adam
 
-from seg_tgce.loss.tgce import TcgeFeatures, TcgePixel, TcgeScalar, TgceBaseline
+from seg_tgce.loss.tgce import TgceBaseline, TgceFeatures, TgcePixel, TgceScalar
 from seg_tgce.metrics import DiceCoefficient, JaccardCoefficient
 from seg_tgce.models.unet import (
     unet_baseline,
@@ -55,39 +55,32 @@ def build_baseline_model_from_hparams(
 
 
 def build_scalar_model_from_hparams(
+    *,
     learning_rate: float,
     q: float,
     noise_tolerance: float,
-    b: float,
     a: float,
+    b: float,
+    c: float,
+    lambda_reg_weight: float,
+    lambda_entropy_weight: float,
+    lambda_sum_weight: float,
     num_classes: int,
     target_shape: tuple,
     n_scorers: int,
 ) -> Model:
-    """Build the scalar model with direct hyperparameter values.
-
-    Args:
-        learning_rate: Learning rate for the optimizer
-        q: q parameter for TGCE loss
-        noise_tolerance: Noise tolerance parameter for TGCE loss
-        lambda_reg_weight: Regularization weight for TGCE loss
-        lambda_entropy_weight: Entropy weight for TGCE loss
-        lambda_sum_weight: Sum weight for TGCE loss
-        num_classes: Number of classes in the segmentation
-        target_shape: Target shape of input images
-        n_scorers: Number of annotators/scorers
-
-    Returns:
-        Compiled Keras model
-    """
     optimizer = Adam(learning_rate=learning_rate)
 
-    loss_fn = TcgeScalar(
+    loss_fn = TgceScalar(
         num_classes=num_classes,
         q=q,
         noise_tolerance=noise_tolerance,
-        b=b,
         a=a,
+        b=b,
+        c=c,
+        lambda_reg_weight=lambda_reg_weight,
+        lambda_entropy_weight=lambda_entropy_weight,
+        lambda_sum_weight=lambda_sum_weight,
         name="TGCE",
     )
 
@@ -145,7 +138,7 @@ def build_features_model_from_hparams(
     """
     optimizer = Adam(learning_rate=learning_rate)
 
-    loss_fn = TcgeFeatures(
+    loss_fn = TgceFeatures(
         num_classes=num_classes,
         q=q,
         noise_tolerance=noise_tolerance,
@@ -193,7 +186,7 @@ def build_pixel_model_from_hparams(
 ) -> Model:
     optimizer = Adam(learning_rate=learning_rate)
 
-    loss_fn = TcgePixel(
+    loss_fn = TgcePixel(
         num_classes=num_classes,
         q=q,
         noise_tolerance=noise_tolerance,