Add utils for models exploration

Author: ViiSkor

src/metrics.py

@@ -133,3 +133,27 @@ def mean_iou(y_true, y_pred):
       y_pred = K.cast(K.greater(y_pred, t), dtype='float32')
       score = iou(y_true, y_pred)
       prec.append(score)
+
+
+def mean_dice(y_true, y_pred):
+    """Mean of dice score.
+    Thresholds for mask are from 0.5, to 1.0 with a step 0.5.
+  
+    Args:
+      y_true (numpy.array/tf.tensor):
+        b x X x Y( x Z...) x c One hot encoding of ground truth.
+      y_pred (numpy.array/tf.tensor):
+        b x X x Y( x Z...) x c Network output, must sum to 1 over c channel.
+    Returns:
+      (tf.tensor): Mean IoU coefficient.
+    """
+
+    prec = []
+    for t in np.arange(0.5, 1.0, 0.05):
+      y_pred = K.cast(K.greater(y_pred, t), dtype='float32')
+      y_pred = K.constant(y_pred)
+      y_true = K.constant(y_true)
+      score = dice_coefficient(y_true, y_pred)
+      prec.append(score)
+
+    return K.mean(K.stack(prec), axis=0)

src/utils.py

@@ -170,3 +170,92 @@ def get_dataset_split(data_paths, train_ratio=0.7, seed=42, shuffle=True):
   test_data_paths = test_data_paths[n_test:]
 
   return train_data_paths, test_data_paths, val_data_paths
+
+
+def stack_2D_2_3D(samples_sep, arr, dim, n_channels):
+  stacked_pred = []
+  prev_idx = 0
+  for sep in samples_sep:
+    pred = arr[prev_idx:prev_idx+sep]
+    if pred.shape[0] > dim[0]:
+      depth_shift = int((pred.shape[0] - dim[0]) // 2)+1
+      pred = pred[depth_shift:-depth_shift,:,:,:]
+    pred = np.moveaxis(pred, -1, 0)
+    _, pred = pad(pred, pred, pred.shape[1:],  dim, n_channels, n_channels)
+    pred = np.moveaxis(pred, 0, -1)
+    stacked_pred.append(pred)
+    prev_idx += sep + 1
+
+  return np.array(stacked_pred)
+
+
+def load_sample(path, dim, scan_types, classes, merge_classes, mode="3D"):
+  dim_before_axes_swap = (dim[-1], dim[1], dim[0])
+  if mode=="3D":
+    masks = preprocess_label(np.asanyarray(nib.load(path['seg']).dataobj), output_classes=classes, merge_classes=merge_classes)
+    imgs = np.array([np.asanyarray(nib.load(path[m]).dataobj) for m in scan_types])
+  elif mode=="2D":
+    masks = preprocess_label(np.load(path['seg']), output_classes=classes, merge_classes=merge_classes)
+    imgs = np.array([np.load(path[m]) for m in scan_types], dtype=np.float16)
+    imgs = np.moveaxis(imgs, [0, 1, 2, 3], [0, 3, 2, 1])
+    masks = np.moveaxis(masks, [0, 1, 2, 3], [0, 3, 2, 1])
+
+  imgs, masks = crop(imgs, masks, depth=dim[0])
+  imgs, masks = pad(imgs, masks, masks.shape[1:],  dim_before_axes_swap, n_channels, n_classes)
+
+  imgs = change_orientation(imgs)
+  masks = change_orientation(masks)
+
+  return masks, imgs
+
+
+def evaluate(data_paths, prediction, metric, dim, scan_types, classes, merge_classes, mode="3D"):
+  scores = {'class': [], 'score': []}
+  for path, pred in zip(data_paths, prediction):
+    if merge_classes:
+      load = ['mask']
+    else:
+      load = classes
+    for cls in load:
+      if merge_classes:
+        cls = classes
+        cls_name = 'mask'
+      else:
+        cls_name = cls
+        cls = [cls]
+      mask, _ = load_sample(path=path, dim=dim, scan_types=scan_types, classes=cls, merge_classes=merge_classes, mode=mode)
+      mask = np.array([mask])
+      pred = np.array([pred])
+      score = metric(mask, pred)
+      scores['class'] = scores['class'] + [cls_name]
+      scores['score'] = scores['score'] + [score.numpy()]
+  return scores
+
+
+def get_data_paths4existing_slit(data_dir, splitted_data, mode="2D"):
+  data_paths = []
+  for modalities in splitted_data:
+    curr_case = {}
+    for name, path in modalities.items():
+      path = path.split('/')[-2:]
+      if mode=="2D":
+        ext = '.npy'
+      elif mode=="3D":
+        ext = '.nii.gz'
+      path[-1] = path[-1].split('.')[0] + ext
+      curr_case[name] = os.path.join(data_dir, *path)
+    data_paths.append(curr_case)
+  return data_paths
+
+
+def get_2D_sep_slices(paths_unnpacked):
+  samples_sep = []
+  prev_slice_idx = 0
+  for i, path in enumerate(paths_unnpacked):
+    if path['seg'][1] == 0:
+      samples_sep.append(prev_slice_idx)
+    prev_slice_idx = path['seg'][1]
+
+  samples_sep.append(prev_slice_idx)
+  samples_sep.pop(0)
+  return samples_sep

src/visualization.py

@@ -6,19 +6,20 @@
 plt.style.use('seaborn-pastel')
 
 
-def animate_scan(scan, mask):
-  fig = plt.figure(figsize=(16, 8))
-  ax1 = fig.add_subplot(1,2,1)
-  ax2 = fig.add_subplot(1,2,2)
+def animate_scan(images:list, titles:list, figsize=(16, 8)):
+  fig = plt.figure(figsize=figsize)
+  plots_num = len(titles)
+  axis = [fig.add_subplot(1, plots_num, i) for i in range(1, plots_num+1)]
 
   myimages = []
-  for i in range(scan.shape[0]):
-      ax1.axis('off')
-      ax1.set_title('Scan', fontsize='medium')
-      ax2.axis('off')
-      ax2.set_title('Mask', fontsize='medium')
+  for i in range(images[0].shape[0]):
+    curr_slice = []
+    for ax_idx, (img, title) in enumerate(zip(images, titles)):
+      axis[ax_idx].axis('off')
+      axis[ax_idx].set_title(title, fontsize=18)
+      curr_slice.append(axis[ax_idx].imshow(img[i], cmap='Greys_r'))
 
-      myimages.append([ax1.imshow(scan[i], cmap='Greys_r'), ax2.imshow(mask[i], cmap='Greys_r')])
+    myimages.append(curr_slice)
 
   anim = animation.ArtistAnimation(fig, myimages, interval=1000, blit=True, repeat_delay=1000)
   return anim