Allow cutoff predictions, support empty background pixels and tests

.dockerignore

@@ -1,3 +1,5 @@
+**/tmp*
+test_data
 .ruff_cache
 .tox
 *.egg-info

.gitignore

@@ -0,0 +1 @@
+**/benchmark_results

superpixel_classification/SuperpixelClassification/SuperpixelClassification.xml

@@ -61,7 +61,7 @@
     <description>Superpixel parameters</description>
     <boolean>
       <name>gensuperpixels</name>
-      <longflag>generate-superpxiels</longflag>
+      <longflag>generate-superpixels</longflag>
       <description>If an image does not have an annotation with superpixels, generate one</description>
       <label>Generate superpixels</label>
       <default>true</default>
@@ -100,6 +100,13 @@
       <label>Train model</label>
       <default>true</default>
     </boolean>
+    <boolean>
+      <name>useCuda</name>
+      <longflag>usecuda</longflag>
+      <description>Whether or not to use GPU/cuda (true) or cpu (false).</description>
+      <label>Use CUDA</label>
+      <default>false</default>
+    </boolean>
     <integer>
       <name>batchSize</name>
       <longflag>batchsize</longflag>
@@ -198,5 +205,12 @@
       <default>4</default>
       <description>The number of worker threads for superpixel and feature generation</description>
     </integer>
+    <integer>
+      <name>cutoff</name>
+      <longflag>cutoff</longflag>
+      <label>Number of annotations per slide</label>
+      <default>500</default>
+      <description>Number of unannotated superpixels to use per slide for features, training and predictions</description>
+    </integer>
   </parameters>
 </executable>

superpixel_classification/SuperpixelClassification/SuperpixelClassificationTensorflow.py

@@ -3,6 +3,7 @@
 from typing import Optional
 
 import h5py
+import numpy as np
 import tensorflow as tf
 from SuperpixelClassificationBase import SuperpixelClassificationBase
 
@@ -35,33 +36,56 @@ class SuperpixelClassificationTensorflow(SuperpixelClassificationBase):
     def __init__(self):
         self.training_optimal_batchsize: Optional[int] = None
         self.prediction_optimal_batchsize: Optional[int] = None
+        self.use_cuda = False
 
     def trainModelDetails(self, record, annotationName, batchSize, epochs, itemsAndAnnot, prog,
-                          tempdir, trainingSplit):
-        # print(f'Tensorflow trainModelDetails(batchSize={batchSize}, ...)')
-        # make model
-        num_classes = len(record['labels'])
-        model = tf.keras.Sequential([
-            tf.keras.layers.Rescaling(1.0 / 255),
-            tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
-            tf.keras.layers.MaxPooling2D(),
-            tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
-            tf.keras.layers.MaxPooling2D(),
-            tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
-            tf.keras.layers.MaxPooling2D(),
-            tf.keras.layers.Flatten(),
-            # tf.keras.layers.Dropout(0.2),
-            tf.keras.layers.Dense(128, activation='relu'),
-            tf.keras.layers.Dense(num_classes)])
-        prog.progress(0.2)
-        model.compile(optimizer='adam',
-                      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-                      metrics=['accuracy'])
+                          tempdir, trainingSplit, use_cuda):
+        self.use_cuda = use_cuda
+
+        # Enable GPU memory growth globally to avoid precondition errors
+        gpus = tf.config.list_physical_devices('GPU')
+        if gpus and self.use_cuda:
+            try:
+                for gpu in gpus:
+                    tf.config.experimental.set_memory_growth(gpu, True)
+            except RuntimeError as e:
+                print(f"Could not set memory growth: {e}")
+        if not self.use_cuda:
+            tf.config.set_visible_devices([], 'GPU')
+        device = "gpu" if use_cuda else "cpu"
+        print(f"Using device: {device}")
+
+        # Dataset preparation (outside strategy scope)
+        ds_h5 = record['ds']
+        labelds_h5 = record['labelds']
+        # Fully load to memory and break h5py reference
+        ds_numpy = np.array(ds_h5[:])
+        labelds_numpy = np.array(labelds_h5[:])
+
+        strategy = tf.distribute.MirroredStrategy()
+        with strategy.scope():
+            num_classes = len(record['labels'])
+            model = tf.keras.Sequential([
+                tf.keras.layers.Rescaling(1.0 / 255),
+                tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
+                tf.keras.layers.MaxPooling2D(),
+                tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
+                tf.keras.layers.MaxPooling2D(),
+                tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
+                tf.keras.layers.MaxPooling2D(),
+                tf.keras.layers.Flatten(),
+                tf.keras.layers.Dense(128, activation='relu'),
+                tf.keras.layers.Dense(num_classes)])
+            prog.progress(0.2)
+            model.compile(optimizer='adam',
+                          loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+                          metrics=['accuracy'])
+
         prog.progress(0.7)
-        # generate split
-        full_ds = tf.data.Dataset.from_tensor_slices((record['ds'], record['labelds']))
-        full_ds = full_ds.shuffle(1000)  # add seed=123 ?
-        count = len(full_ds)
+        # generate split using numpy arrays
+        full_ds = tf.data.Dataset.from_tensor_slices((ds_numpy, labelds_numpy))
+        full_ds = full_ds.shuffle(1000)
+        count = len(ds_numpy)
         train_size = int(count * trainingSplit)
         if batchSize < 1:
             batchSize = self.findOptimalBatchSize(model, full_ds, training=True)
@@ -85,24 +109,53 @@ def trainModelDetails(self, record, annotationName, batchSize, epochs, itemsAndA
         self.saveModel(model, modelPath)
         return history, modelPath
 
+    def _get_device(self, use_cuda):
+        if tf.config.list_physical_devices('GPU') and use_cuda:
+            return '/GPU:0'
+        return '/CPU:0'
+
     def predictLabelsForItemDetails(
-        self, batchSize, ds: h5py._hl.dataset.Dataset, item, model, prog,
+            self, batchSize, ds: h5py._hl.dataset.Dataset, indices, item, model, use_cuda, prog,
     ):
-        # print(f'Tensorflow predictLabelsForItemDetails(batchSize={batchSize}, ...)')
         if batchSize < 1:
             batchSize = self.findOptimalBatchSize(
                 model, tf.data.Dataset.from_tensor_slices(ds), training=False,
             )
             print(f'Optimal batch size for prediction = {batchSize}')
-        predictions = model.predict(
-            ds,
-            batch_size=batchSize,
-            callbacks=[_LogTensorflowProgress(
-                prog, (ds.shape[0] + batchSize - 1) // batchSize, 0.05, 0.35, item)])
-        prog.item_progress(item, 0.4)
-        # softmax to scale to 0 to 1
-        catWeights = tf.nn.softmax(predictions)
-        return catWeights, predictions
+
+        device = self._get_device(use_cuda)
+        with tf.device(device):
+            # Create a dataset that pairs the data with their indices
+            dataset = tf.data.Dataset.from_tensor_slices((ds, indices))
+            dataset = dataset.batch(batchSize)
+
+            # Initialize arrays to store results
+            all_predictions = []
+            all_cat_weights = []
+            all_indices = []
+
+            # Iterate through batches manually to keep track of indices
+            for data, batch_indices in dataset:
+                batch_predictions = model.predict(
+                    data,
+                    batch_size=batchSize,
+                    verbose=0)  # Set verbose=0 to avoid multiple progress bars
+
+                # Apply softmax to scale to 0 to 1
+                batch_cat_weights = tf.nn.softmax(batch_predictions)
+
+                all_predictions.append(batch_predictions)
+                all_cat_weights.append(batch_cat_weights)
+                all_indices.append(batch_indices)
+
+                prog.item_progress(item, 0.4)
+
+            # Concatenate all results
+            predictions = tf.concat(all_predictions, axis=0)
+            catWeights = tf.concat(all_cat_weights, axis=0)
+            final_indices = tf.concat(all_indices, axis=0)
+
+            return catWeights.numpy(), predictions.numpy(), final_indices.numpy().astype(np.int64)
 
     def findOptimalBatchSize(self, model, ds, training) -> int:
         if training and self.training_optimal_batchsize is not None:

superpixel_classification/SuperpixelClassification/SuperpixelClassificationTorch.py

@@ -66,12 +66,10 @@ class _BayesianPatchTorchModel(bbald.consistent_mc_dropout.BayesianModule):
     # A Bayesian model that takes patches (2-dimensional shape) rather than vectors
     # (1-dimensional shape) as input.  It is useful when feature != 'vector' and
     # SuperpixelClassificationBase.certainty == 'batchbald'.
-    def __init__(self, num_classes: int) -> None:
+    def __init__(self, num_classes: int, device : torch.device) -> None:
         # Set `self.device` as early as possible so that other code does not lock out
         # what we want.
-        self.device: str = torch.device(
-            ('cuda' if torch.cuda.is_available() and torch.cuda.device_count() > 0 else 'cpu'),
-        )
+        self.device : torch.device = device
         # print(f'Initial model.device = {self.device}')
         super(_BayesianPatchTorchModel, self).__init__()
 
@@ -134,18 +132,16 @@ class _VectorTorchModel(torch.nn.Module):
     # (2-dimensional shape) as input.  It is useful when feature == 'vector' and
     # SuperpixelClassificationBase.certainty != 'batchbald'.
 
-    def __init__(self, input_dim: int, num_classes: int) -> None:
+    def __init__(self, input_dim: int, num_classes: int, device : torch.device) -> None:
         # Set `self.device` as early as possible so that other code does not lock out
         # what we want.
-        self.device: str = torch.device(
-            ('cuda' if torch.cuda.is_available() and torch.cuda.device_count() > 0 else 'cpu'),
-        )
+        self.device: torch.device = device
         # print(f'Initial model.device = {self.device}')
         super(_VectorTorchModel, self).__init__()
 
         self.input_dim: int = input_dim
         self.num_classes: int = num_classes
-        self.fc: torch.Module = torch.nn.Linear(input_dim, num_classes)
+        self.fc: torch.Linear = torch.nn.Linear(input_dim, num_classes)
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         # TODO: Is torch.mul appropriate here?
@@ -161,20 +157,18 @@ class _BayesianVectorTorchModel(bbald.consistent_mc_dropout.BayesianModule):
     # (2-dimensional shape) as input.  It is useful when feature == 'vector' and
     # SuperpixelClassificationBase.certainty == 'batchbald'.
 
-    def __init__(self, input_dim: int, num_classes: int) -> None:
+    def __init__(self, input_dim: int, num_classes: int, device : torch.device) -> None:
         # Set `self.device` as early as possible so that other code does not lock out
         # what we want.
-        self.device: str = torch.device(
-            ('cuda' if torch.cuda.is_available() and torch.cuda.device_count() > 0 else 'cpu'),
-        )
+        self.device = device
         # print(f'Initial model.device = {self.device}')
         super(_BayesianVectorTorchModel, self).__init__()
 
         self.input_dim: int = input_dim
         self.num_classes: int = num_classes
         self.bayesian_samples: int = 12
-        self.fc: torch.Module = torch.nn.Linear(input_dim, num_classes)
-        self.fc_drop: torch.Module = bbald.consistent_mc_dropout.ConsistentMCDropout()
+        self.fc: torch.Linear = torch.nn.Linear(input_dim, num_classes)
+        self.fc_drop: torch.ConsistentMCDropout = bbald.consistent_mc_dropout.ConsistentMCDropout()
 
     def mc_forward_impl(self, input: torch.Tensor) -> torch.Tensor:
         # TODO: Is torch.mul appropriate here?
@@ -311,24 +305,27 @@ def trainModelDetails(
         prog: ProgressHelper,
         tempdir: str,
         trainingSplit: float,
+        cuda : bool,
     ):
+        device = torch.device("cuda" if cuda else "cpu")
+        print(f"Using device: {device}")
         # make model
         num_classes: int = len(record['labels'])
         model: torch.nn.Module
         if self.feature_is_image:
             # Feature is patch
             if self.certainty == 'batchbald':
-                model = _BayesianPatchTorchModel(num_classes)
+                model = _BayesianPatchTorchModel(num_classes, device)
             else:
                 mesg = 'Expected torch model for input of type image to be Bayesian'
                 raise ValueError(mesg)
         else:
             # Feature is vector
             input_dim: int = record['ds'].shape[1]
             if self.certainty == 'batchbald':
-                model = _BayesianVectorTorchModel(input_dim, num_classes)
+                model = _BayesianVectorTorchModel(input_dim, num_classes, device)
             else:
-                model = _VectorTorchModel(input_dim, num_classes)
+                model = _VectorTorchModel(input_dim, num_classes, device)
         model.to(model.device)
 
         # print(f'Torch trainModelDetails(batchSize={batchSize}, ...)')
@@ -348,6 +345,7 @@ def trainModelDetails(
         val_ds: torch.utils.data.TensorDataset
         train_dl: torch.utils.data.DataLoader
         val_dl: torch.utils.data.DataLoader
+        prog.message('Loading features for model training')
         train_arg1 = (
             torch.from_numpy(record['ds'][train_indices].transpose((0, 3, 2, 1)))
             if self.feature_is_image
@@ -507,7 +505,7 @@ def fitModel(
         return history
 
     def predictLabelsForItemDetails(
-        self, batchSize: int, ds_h5, item, model: torch.nn.Module, prog: ProgressHelper,
+        self, batchSize: int, ds_h5, indices, item, model: torch.nn.Module, use_cuda : bool, prog: ProgressHelper,
     ):
         # print(f'Torch predictLabelsForItemDetails(batchSize={batchSize}, ...)')
         num_superpixels: int = ds_h5.shape[0]
@@ -517,6 +515,9 @@ def predictLabelsForItemDetails(
         num_classes: int = model.num_classes
         # print(f'{num_classes = }')
 
+        # also set on model.device, ideally
+        #device = torch.device("cuda" if use_cuda else "cpu")
+
         callbacks = [
             _LogTorchProgress(prog, 1 + (num_superpixels - 1) // batchSize, 0.05, 0.35, item),
         ]
@@ -532,19 +533,21 @@ def predictLabelsForItemDetails(
         for cb in callbacks:
             cb.on_predict_begin(logs=logs)
 
+        # ds also needs to have information about the indices so that we can shuffle the data but still link it to an index
         ds: torch.utils.data.TensorDataset = torch.utils.data.TensorDataset(
             (
                 torch.from_numpy(np.array(ds_h5).transpose((0, 3, 2, 1)))
                 if self.feature_is_image
                 else torch.from_numpy(np.array(ds_h5))
-            ),
+            ), torch.from_numpy(indices),
         )
         if batchSize < 1:
             batchSize = self.findOptimalBatchSize(model, ds, training=False)
             print(f'Optimal batch size for prediction (device = {model.device}) = {batchSize}')
         dl: torch.utils.data.DataLoader = torch.utils.data.DataLoader(ds, batch_size=batchSize)
         predictions: NDArray[np.float_] = np.zeros((num_superpixels, bayesian_samples, num_classes))
         catWeights: NDArray[np.float_] = np.zeros((num_superpixels, bayesian_samples, num_classes))
+        outIndices: NDArray[np.int64] = np.zeros(num_superpixels, dtype=np.int64)
         with torch.no_grad():
             model.eval()  # Tell torch that we will be doing predictions
             row: int = 0
@@ -567,14 +570,16 @@ def predictLabelsForItemDetails(
                 catWeights_raw = torch.nn.functional.softmax(predictions_raw, dim=-1)
                 predictions[row:new_row, :, :] = predictions_raw.detach().cpu().numpy()
                 catWeights[row:new_row, :, :] = catWeights_raw.detach().cpu().numpy()
+                outIndices[row:new_row] = data[1].detach().cpu().numpy().astype(np.int64)[:]
+
                 row = new_row
                 for cb in callbacks:
                     cb.on_predict_batch_end(i)
         for cb in callbacks:
             cb.on_predict_end({'outputs': predictions})
         prog.item_progress(item, 0.4)
         # scale to units
-        return catWeights, predictions
+        return catWeights, predictions, outIndices
 
     def findOptimalBatchSize(
         self, model: torch.nn.Module, ds: torch.utils.data.TensorDataset, training: bool,
@@ -651,9 +656,14 @@ def add_safe_globals(self):
 
     def loadModel(self, modelPath):
         self.add_safe_globals()
-        model = torch.load(modelPath)
-        model.eval()
-        return model
+        try:
+            model = torch.load(modelPath, weights_only=False)
+            model.eval()
+            return model
+        except Exception as e:
+            print(f"Unable to load {modelPath}")
+            raise
+
 
     def saveModel(self, model, modelPath):
         self.add_safe_globals()