Add background channel

pyplati · Apr 1, 2024 · ec18dbc · ec18dbc
1 parent f20acf6
commit ec18dbc
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/platipy/imaging/cnn/train.py b/platipy/imaging/cnn/train.py
@@ -360,6 +360,11 @@ def infer(
                     s = s.unsqueeze(0)
                     s = s.unsqueeze(0)
 
+                    # Add in background channel
+                    not_s = 1 - sample_strategy.max(axis=1).values
+                    not_s = torch.unsqueeze(not_s, dim=1)
+                    s = torch.cat((not_s, s), dim=1).float()
+
                 if self.hparams.prob_type == "prob":
                     if seg is not None:
                         self.prob_unet.forward(x, seg=s, training=True)