Organ Contours (using matplotlib, opencv, scipy)

contours_with_gt_pred.py

"""
This script shall plot contours using cv2 on an image
- Useful for plotting organ contours on a medical image
-- Note: Organ Mask should contain labels values in the [0,255] range 
"""

# Standard Libraries
import pdb
import cv2
import numpy as np
import matplotlib
from pathlib import Path
import SimpleITK as sitk
from matplotlib import cm
import matplotlib.pyplot as plt

# Libraries for smoothing contours
from scipy import ndimage
from scipy.interpolate import splprep, splev


###########################################################
#                      GLOBAL PARAMS                      #
###########################################################

SMOOTHING = 50 # contour smoothing condition
CONTOUR_PERC_POINTS = 0.7

CONTOUR_GT    = [(0,255,0), (0,0,255), (255,0,0), (241, 85, 230)]
CONTOUR_PRED  = [(0,255,0), (0,0,255), (255,0,0), (241, 85, 230)]

CONTOUR_ALPHA = 0.5
CONTOUR_WIDTH = 0.25

PLT_INTERP    = 'spline36'

cmap = cm.magma

FONTSIZE_CBAR = 5

###########################################################
#                      HELPER FUNCTIONS                   #
###########################################################

def get_smooth_contour(contour):
  # https://gist.github.com/shubhamwagh/b8148e65a8850a974efd37107ce3f2ec
  x = contour[0][:,0,0].tolist()
  y = contour[0][:,0,1].tolist()
  tck, u       = splprep([contour[0][:,0,0].tolist(), contour[0][:,0,1].tolist()], u=None, s=SMOOTHING, per=0) # higher the s value, more the smoothing
  u_new        = np.linspace(u.min(), u.max(), int(len(x) * CONTOUR_PERC_POINTS))
  x_new, y_new = splev(u_new, tck, der=0)
  contour_new  = np.array([[[int(i[0]), int(i[1])]] for i in zip(x_new,y_new)])
  return contour_new

###########################################################
#                           MAIN                          #
###########################################################

# Step 0.1 - Init volumes
vol_ct   = sitk.GetArrayFromImage(sitk.ReadImage(str(Path('')))) # [D,H,W]
vol_gt   = sitk.GetArrayFromImage(sitk.ReadImage(str(Path(''))))  # [D,H,W] 
vol_pred = sitk.GetArrayFromImage(sitk.ReadImage(str(Path('')))) # [D,H,W]
slice_id = 0
axis     = 0 # i.e. D=axial slices

if axis == 0:
    slice_ct   = vol_ct[slice_id]   # np.array - [H,W] - contains grayscale value representing Hounsfield units on a CT slice
    slice_gt   = vol_gt[slice_id]   # np.array - [H,W] - contains a binary map for an organ on a particular slice
    slice_pred = vol_pred[slice_id] # np.array - [H,W] - contains a binary map for an organ on a particular slice

# Step 0.2 - Init plot
f,axarr = plt.subplots(1, 1, figsize=(10,5), dpi=1000)
axarr.imshow(slice_ct, interpolation=PLT_INTERP, cmap='gray')

# Step 1 - Extract GT contours
if slice_gt is not None:
  for label_id in np.unique(slice_gt):
      if label_id > 0: # assuming background label is 0
        contours_gt, _    = cv2.findContours((slice_gt == label_id).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # [cv2.RETR_EXTERNAL, cv2.RETR_CCOMP]
        for contour_gt_ in contours_gt:
            if len(contour_gt_) > 2:
              coords_gt     = get_smooth_contour([contour_gt_])
              coords_gt     = np.append(coords_gt, [coords_gt[0]], axis=0)
              coords_gt     = coords_gt[:,0,:].tolist()
              xs_gt, ys_gt  = zip(*coords_gt)
              axarr.plot(xs_gt,ys_gt    , color=np.array(CONTOUR_GT[int(label_id)-1])/255.0, linewidth=CONTOUR_WIDTH, alpha=CONTOUR_ALPHA)

# Step 2 - Extract Pred Contours
if slice_pred is not None:
  for label_id in np.unique(slice_pred):
    if label_id > 0: # assuming background label is 0
      contours_pred, _  = cv2.findContours((slice_pred == label_id).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # [cv2.RETR_EXTERNAL, cv2.RETR_CCOMP]
      for contour_pred_ in contours_pred:
        if len(contour_pred_) > 2:
          coords_pred      = get_smooth_contour([contour_pred_])
          coords_pred      = np.append(coords_pred, [coords_pred[0]], axis=0)
          coords_pred      = coords_pred[:,0,:].tolist()
          xs_pred, ys_pred = zip(*coords_pred)
          axarr.plot(xs_pred,ys_pred, color=np.array(CONTOUR_PRED[int(label_id)-1])/255.0, linewidth=CONTOUR_WIDTH, alpha=CONTOUR_ALPHA, linestyle='dashed')

# Step 3 - Some beautification
axarr.axis('off')

# Step 4 - Finish
# plt.show()
filename = 'contours_with_gt_pred.png'
plt.savefig(filename, bbox_inches='tight')

contours_with_uncertainty.py

"""
This script shall plot contours using cv2 on an image
- Useful for plotting organ contours on a medical image
"""

# Standard Libraries
import pdb
import cv2
import numpy as np
import matplotlib
from pathlib import Path
import SimpleITK as sitk
from matplotlib import cm
import matplotlib.pyplot as plt

# Libraries for smoothing contours
from scipy import ndimage
from scipy.interpolate import splprep, splev


###########################################################
#                      GLOBAL PARAMS                      #
###########################################################

SMOOTHING = 50 # contour smoothing condition
CONTOUR_PERC_POINTS = 0.7

CONTOUR_GT    = (0,255,0) # green
CONTOUR_PRED  = (0,0,255) # blue
CONTOUR_ALPHA = 1.0
CONTOUR_WIDTH = 0.25

UNC_VMIN      = 0.0
UNC_VMAX      = 0.5
UNC_THRESH    = 0.2

PLT_INTERP    = 'spline36'

cmap = cm.magma

FONTSIZE_CBAR = 5

###########################################################
#                      HELPER FUNCTIONS                   #
###########################################################

def get_smooth_contour(contour):
  x = contour[0][:,0,0].tolist()
  y = contour[0][:,0,1].tolist()
  tck, u       = splprep([contour[0][:,0,0].tolist(), contour[0][:,0,1].tolist()], u=None, s=SMOOTHING, per=0) # higher the s value, more the smoothing
  u_new        = np.linspace(u.min(), u.max(), int(len(x) * CONTOUR_PERC_POINTS))
  x_new, y_new = splev(u_new, tck, der=0)
  contour_new  = np.array([[[int(i[0]), int(i[1])]] for i in zip(x_new,y_new)])
  return contour_new

###########################################################
#                           MAIN                          #
###########################################################

# Step 0 - Init
vol_ct   = sitk.GetArrayFromImage(sitk.ReadImage(str(Path('')))) # [D,H,W]
vol_gt   = sitk.GetArrayFromImage(sitk.ReadImage(str(Path(''))))  # [D,H,W]
vol_pred = None
vol_unc  = None
vol_gt[vol_gt != 2] = 0
vol_gt[vol_gt == 2] = 1
slice_id = None 

slice_ct   = vol_ct[slice_id] # np.array - [H,W] - contains grayscale value representing Hounsfield units on a CT slice
slice_gt   = vol_gt[slice_id] # np.array - [H,W] - contains a binary map for an organ on a particular slice (only a single blob)
slice_pred = None # np.array - [H,W] - contains a binary map for an organ on a particular slice (only a single blob)
slice_unc  = None # np.array - [H,W] - contains values ranging from [0,inf] for uncertainty on a particular slice

# Step 1 - Extract GT contours
if slice_gt is not None:
  contour_gt, _ = cv2.findContours(slice_gt.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # [cv2.RETR_EXTERNAL, cv2.RETR_CCOMP]
  coords_gt     = get_smooth_contour(contour_gt)
  coords_gt     = np.append(coords_gt, [coords_gt[0]], axis=0)
  coords_gt     = coords_gt[:,0,:].tolist()
  xs_gt, ys_gt  = zip(*coords_gt)

# Step 2 - Extract Pred Contours
if slice_pred is not None:
  contour_pred, _  = cv2.findContours(slice_pred.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # [cv2.RETR_EXTERNAL, cv2.RETR_CCOMP]
  coords_pred      = get_smooth_contour(contour_pred)
  coords_pred.append(coords_pred[0])
  xs_pred, ys_pred = zip(*coords_pred)

# Step 3.1 - Plot GT + Pred
f,axarr = plt.subplots(1, 2, figsize=(20,10), gridspec_kw = {'wspace':0.05, 'hspace':0.01}, dpi=1000)
axarr[0].imshow(slice_ct, interpolation=PLT_INTERP, cmap='gray')
if slice_pred is not None:
  axarr[0].plot(xs_gt,ys_gt    , color=np.array(CONTOUR_GT)/255.0, linewidth=CONTOUR_WIDTH, alpha=CONTOUR_ALPHA)
if slice_pred is not None:
  axarr[0].plot(xs_pred,ys_pred, color=np.array(CONTOUR_GT)/255.0, linewidth=CONTOUR_WIDTH, alpha=CONTOUR_ALPHA)

# Step 3.2 - Plot unc
if slice_unc is not None:
  slice_unc_cmap = np.array(slice_unc, copy=True)           # [H,W]
  slice_unc_cmap[slice_unc_cmap < UNC_VMIN] = UNC_VMIN      # [H,W]
  slice_unc_cmap[slice_unc_cmap > UNC_VMAX] = UNC_VMAX      # [H,W]
  slice_unc_cmap = ((slice_unc_cmap - UNC_VMIN) / (UNC_VMAX-UNC_VMIN))  # [H,W]
  slice_unc_cmap = cmap(slice_unc_cmap)                 # [H,W,4]
  slice_unc_alpha = np.array(slice_unc, copy=True)
  slice_unc_alpha[slice_unc_alpha < UNC_THRESH] = 0
  slice_unc_alpha[slice_unc_alpha >= UNC_THRESH] = 1

  slice_unc_alpha = ndimage.gaussian_filter(slice_unc_alpha, sigma=2)
  slice_unc_cmap[:,:,3] = slice_unc_alpha
  axarr[1].imshow(slice_ct, interpolation=PLT_INTERP, cmap='gray')
  axarr[1].imshow(slice_unc_cmap, interpolation=PLT_INTERP)
  axarr[1].plot(xs_gt,ys_gt    , color=np.array(CONTOUR_GT)/255.0, linewidth=CONTOUR_WIDTH, alpha=CONTOUR_ALPHA)
  axarr[1].plot(xs_pred,ys_pred, color=np.array(CONTOUR_GT)/255.0, linewidth=CONTOUR_WIDTH, alpha=CONTOUR_ALPHA)

# Step 4 - Some beautification
_ = [axarr[ax_id].axis('off') for ax_id in range(len(axarr))]

# Step 5 - Colorbar
norm_cmap=matplotlib.colors.Normalize(vmin=UNC_VMIN, vmax=UNC_VMAX)
cb = f.colorbar(cm.ScalarMappable(cmap=cmap.name, norm=norm_cmap), ax=axarr.ravel().tolist(), pad=0.01, shrink=0.96)
cb.ax.set_yticklabels(["{:.2f}".format(i) for i in cb.get_ticks()])
for t in cb.ax.get_yticklabels():t.set_fontsize(FONTSIZE_CBAR)

# Step 6 - Finish
plt.show(block=False)
filename = 'contours_with_uncertainty.png'
plt.savefig(filename, bbox_inches='tight')