From 9c5c750e6e6e9673f4faeb54f4ad5e45b552d2cd Mon Sep 17 00:00:00 2001
From: Robert Kiewisz <56911280+RRobert92@users.noreply.github.com>
Date: Mon, 8 Jul 2024 15:34:28 +0200
Subject: [PATCH] fix

---
 tardis_em/cnn/utils/build_cnn.py | 2 +-
 tardis_em/utils/predictor.py     | 5 +++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/tardis_em/cnn/utils/build_cnn.py b/tardis_em/cnn/utils/build_cnn.py
index 45cdcae..744f102 100644
--- a/tardis_em/cnn/utils/build_cnn.py
+++ b/tardis_em/cnn/utils/build_cnn.py
@@ -511,7 +511,7 @@ def __init__(
         self.sigmoid = sigmoid
 
         self.patch_sizes = [img_patch_size]
-        self.update_patch_size(img_patch_size)
+        self.update_patch_size(img_patch_size, self.sigmoid)
         self.build_cnn_model()
 
     def update_patch_size(self, img_patch_size, sigmoid):
diff --git a/tardis_em/utils/predictor.py b/tardis_em/utils/predictor.py
index fef355e..d66a69a 100644
--- a/tardis_em/utils/predictor.py
+++ b/tardis_em/utils/predictor.py
@@ -642,6 +642,11 @@ def postprocess_CNN(self, id_name: str):
                 )
             else:
                 if not self.output_format.startswith("return"):
+                    self.log_prediction.append(f"Semantic Prediction: {i[:-self.in_format]}"
+                                               f" | Number of pixels: {np.sum(self.image)}")
+                    with open(join("Predictions", "prediction_log.txt"), "w") as f:
+                        f.write(" \n".join(self.log_prediction))
+
                     tif.imwrite(
                         join(self.am_output, f"{id_name[:-self.in_format]}_CNN.tif"),
                         self.image,