Add optional postprocessing

bayesian_cnn_prometheus/analysis/analyse_lesions.py

@@ -16,8 +16,8 @@ def analyse_lesions():
 
     model_config = load_config(Path(config_path))
     prediction_options = PredictionOptions(
-        chunk_size=model_config['preprocessing']['create_chunks']['chunk_size'],
-        stride=model_config['preprocessing']['create_chunks']['stride'],
+        chunk_size=model_config['evaluation']['chunk_size'],
+        stride=model_config['evaluation']['stride'],
         mc_sample=model_config['mc_samples'])
 
     lesions_analyzer = LesionsAnalyzer(model_path, input_path, prediction_options, patients_to_analysis, output_path)

bayesian_cnn_prometheus/analysis/apply_variance_postprocess.py

@@ -0,0 +1,23 @@
+import cv2
+import numpy as np
+
+
+def is_contour_bad(c):
+    _, radius = cv2.minEnclosingCircle(c)
+    return radius > 10
+
+
+def getXFromRect(item):
+    return item[0]
+
+
+def apply_variance_postprocess(variance: np.ndarray):
+    lower_bound = np.percentile(variance, 99.6)
+    variance[lower_bound > variance] = 0
+
+    variance[np.nonzero(variance)] = 255
+    variance = variance.astype(np.uint8)
+    for i in range(variance.shape[0]):
+        im = cv2.erode(variance[i], np.ones((3, 3)))
+        variance[i, :, :] = cv2.dilate(im, np.ones((3, 3)))
+    return variance

bayesian_cnn_prometheus/analysis/lesions_analyzer.py

@@ -8,6 +8,7 @@
 import numpy as np
 from tqdm import tqdm
 
+from bayesian_cnn_prometheus.analysis.apply_variance_postprocess import apply_variance_postprocess
 from bayesian_cnn_prometheus.constants import Paths
 from bayesian_cnn_prometheus.evaluation.bayesian_model_evaluator import BayesianModelEvaluator
 from bayesian_cnn_prometheus.evaluation.evaluate_model import PredictionOptions, crop_image_to_bounding_box_with_lungs
@@ -35,7 +36,7 @@ def __init__(self, model_path: str, input_path: str, prediction_options: Predict
         self.results = {'chunks_analyse': {'overall': {'tp': 0, 'tn': 0, 'fp': 0, 'fn': 0}},
                         'voxels_analyse': {'overall': {'tp': 0, 'tn': 0, 'fp': 0, 'fn': 0}}}
 
-    def run_analysis(self):
+    def run_analysis(self, should_postprocess_variance: bool = False):
         for image_path in tqdm(glob.glob(os.path.join(self.input_path, Paths.IMAGES_DIR, '*.nii.gz'))):
             patient_idx = int(get_patient_index(image_path))
 
@@ -66,6 +67,9 @@ def run_analysis(self):
                 variance = self.get_variance(cropped_image, cropped_segmentation, variance_path, nifti.affine,
                                              nifti.header)
 
+            if should_postprocess_variance:
+                variance = apply_variance_postprocess(variance)
+
             self._analyze_chunks(patient_idx, cropped_mask, variance)
             self._analyze_voxels(patient_idx, cropped_mask, variance)