155-overlay_segmentation-function-is-missing

miscnn/utils/visualizer.py

@@ -24,6 +24,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
+import seaborn as sns
 import numpy as np
 import os
 import random
@@ -267,21 +268,21 @@ def visualize_samples(sample_list, out_dir="vis", mask_seg=False, mask_pred=True
         out_path = os.path.join(out_dir, file_name)
 
         if mask_seg and mask_pred:
-            vol_truth = overlay_segmentation_greyscale(sample.img_data, sample.seg_data)
-            vol_pred = overlay_segmentation_greyscale(sample.img_data, sample.pred_data)
+            vol_truth = overlay_segmentation(sample.img_data, sample.seg_data, sample.classes)
+            vol_pred = overlay_segmentation(sample.img_data, sample.pred_data, sample.classes)
             fig, (ax1, ax2) = plt.subplots(1, 2)
             ax1.set_title("Ground Truth")
             ax2.set_title("Prediction")
             display(out_path, fig, [ax1, ax2], "", "", np.stack([vol_truth, vol_pred], axis = 0))
 
 
         elif mask_seg:
-            vol_truth = overlay_segmentation_greyscale(sample.img_data, sample.seg_data)
+            vol_truth = overlay_segmentation(sample.img_data, sample.seg_data, sample.classes)
             fig, ax = plt.subplots()
             ax.set_title("Segmentation")
             display(out_path, fig, [ax], "", "", np.expand_dims(vol_truth, axis = 0))
         elif mask_pred:
-            vol_pred = overlay_segmentation_greyscale(sample.img_data, sample.pred_data)
+            vol_pred = overlay_segmentation(sample.img_data, sample.pred_data, sample.classes)
             fig, ax = plt.subplots()
             # Initialize the two subplots (axes) with an empty 512x512 image
             ax.set_title("Segmentation")
@@ -362,6 +363,35 @@ def update(i):
 #-----------------------------------------------------#
 #                     Subroutines                     #
 #-----------------------------------------------------#
+
+def overlay_segmentation(vol, seg, num_classes, cm="hls", alpha=0.3):
+    # Convert volume to RGB
+    vol_rgb = np.stack([vol, vol, vol], axis=-1)
+    # Initialize segmentation in RGB
+    shp = seg.shape
+    seg_rgb = np.zeros((shp[0], shp[1], shp[2], 3), dtype=np.int)
+
+    # Get unique classes from the sample
+    unique_classes, tmp = np.unique(seg, return_counts=True)
+    # Get color palette for all classes
+    color_palette = sns.color_palette(cm, num_classes)
+    for i, label in np.ndenumerate(unique_classes):
+        label = int(label)
+        seg_rgb[np.equal(seg, label)] = (color_palette[label][0]*255, color_palette[label][1]*255, color_palette[label][2]*255)
+
+    # Get binary array for places where an ROI lives
+    segbin = np.greater(seg, 0)
+    repeated_segbin = np.stack((segbin, segbin, segbin), axis=-1)
+
+    # Weighted sum where there's a value to overlay
+    vol_overlayed = np.where(
+        repeated_segbin,
+        np.round(alpha*seg_rgb+(1-alpha)*vol_rgb).astype(np.uint8),
+        np.round(vol_rgb).astype(np.uint8)
+    )
+    # Return final volume with segmentation overlay
+    return vol_overlayed
+
 # Based on: https://github.com/neheller/kits19/blob/master/starter_code/visualize.py
 def overlay_segmentation_greyscale(vol, seg, cm="viridis", alpha=0.3, min_tolerance = -0.1):
 

requirements.txt

@@ -10,3 +10,4 @@ pydicom==2.0.0
 SimpleITK==2.0.2
 scikit-image==0.18.2
 tqdm==4.51.0
+seaborn==0.11.2

setup.py

@@ -28,7 +28,8 @@
                      'batchgenerators==0.21',
                      'pydicom>=2.0.0',
                      'SimpleITK>=2.0.2',
-                     'scikit-image>=0.18.2'],
+                     'scikit-image>=0.18.2',
+                     'seaborn>=0.11.2'],
    classifiers=["Programming Language :: Python :: 3",
                 "Programming Language :: Python :: 3.7",
                 "Programming Language :: Python :: 3.8",