1  #!/usr/bin/env python3
  2  # inspiration from https://github.com/nw-duncan/MRS-voxel-plot/
  3  # 20220324WF - init
  4  #  use nilearn glass brain to visualize slice position
  5  # colors from
  6  # https://matplotlib.org/3.5.1/tutorials/colors/colormaps.html
  7  # glass_brain
  8  # 'ortho', 'x', 'y', 'z', 'xz', 'yx', 'yz', 'l', 'r', 'lr', 'lzr', 'lyr', 'lzry', 'lyrz'.
  9  
 10  import sys
 11  import numpy as np
 12  import matplotlib.pyplot as plt
 13  import argparse
 14  import warnings
 15  warnings.simplefilter("ignore")
 16  import pandas as pd
 17  from nilearn import plotting, image
 18  
 19  
 20  def read_coords(fname='/Volumes/Hera/Projects/7TBrainMech/scripts/mri/MRSI_roi/roi_locations/labels_13MP20200207.txt'):
 21      labels = pd.read_csv(fname, sep="\t", header=None)
 22      # convert from LPI to "RAS+"
 23      centers = [[float(y) for y in x.split(" ")] for x in labels.iloc[:, 1]]
 24      # LPI to RAS+
 25      centers_ras = [(c[0], -1*c[1], c[2]) for c in centers]
 26      return centers_ras
 27  
 28  
 29  def roi_colors(n_roi=13):
 30      "here incase we want to change what roi gets what color"
 31      return plt.cm.get_cmap('tab20').colors[0:n_roi]
 32  
 33  
 34  def plot_rois(ax_roi, fname, roi_idxs=[1, 2, 7, 8, 9, 10], display_mode='xz'):
 35      centers_ras = read_coords(fname)
 36      colors = roi_colors(len(centers_ras))  # same colors despite subselect
 37      used_rois0 = [i-1 for i in roi_idxs]
 38      colors_used = [colors[i] for i in used_rois0]
 39      centers_used = [centers_ras[i] for i in used_rois0]
 40      # no lines between rois
 41      n_roi = len(centers_used)
 42      adjmat0 = np.zeros((n_roi, n_roi))
 43      plotting.plot_connectome(adjacency_matrix=adjmat0,
 44                               node_coords=centers_used,
 45                               annotate=False,
 46                               node_size=50, node_color=colors_used,
 47                               display_mode=display_mode, alpha=.5,
 48                               axes=ax_roi)
 49  
 50  def plot_nii(ax_roi, fname, display_mode='xz', digitize=True):
 51      """plot rois in nii. stat map needs its own function
 52      digitize takes eg roi mask/atlas with values 37,38 to 1,2"""
 53      # fname="/Volumes/Zeus/scratch/MNI_caez_ml_hpc37_38.nii.gz"
 54      #img = image.smooth_img(fname, 'fast')
 55      import nibabel as nib
 56      img = nib.load(fname)
 57      data = img.get_data()
 58  
 59      # if had nifti with roi vaules 7, 8 make them 1,2
 60      rois = np.unique(np.round(data))
 61      if digitize:
 62          dig = np.digitize(np.round(data), bins=rois)-1
 63          img = nib.Nifti1Image(dig, img.affine, img.header)
 64          rois = np.unique(dig)
 65  
 66      # TODO: what if we want 7 to always be the same color?
 67      # need to pass in maxroi?
 68      # TODO: tab20 not best choice? allow user to set
 69  
 70      plotting.plot_glass_brain(img,
 71                               cmap='tab20',
 72                               display_mode=display_mode, alpha=.5,
 73                               vmax=max(rois),
 74                               vmin=0,
 75                               axes=ax_roi)
 76  
 77  
 78  def main(args):
 79      "args display_mode save idx width height"
 80      # Plot the figure
 81      fig = plt.figure()
 82      fig.set_size_inches(args.width, args.height)
 83      ax_roi = plt.subplot(111)
 84      if args.roi_nii:
 85          plot_nii(ax_roi, args.roi_nii, args.display_mode)
 86      if not args.roi_txt in [None, "None", "", "NA"]:
 87          plot_rois(ax_roi, args.roi_txt, args.idx, args.display_mode)
 88      # display. save if not running interactively
 89      if hasattr(sys, 'ps1') or not args.save:
 90          plt.show()
 91      else:
 92          plt.savefig(args.save)
 93  
 94  
 95  if __name__ == "__main__":
 96      parser = argparse.ArgumentParser(description='')
 97      # USED_ROIS = [1, 2, 7, 8, 9, 10]
 98      parser.add_argument('idx', metavar='index', type=int, nargs='+',
 99                          help='1-based ROI indexes. probably want 1 2 7 8 9 10')
100      parser.add_argument('--save', metavar='save',
101                          help='save as', type=str, default=None)
102      parser.add_argument('--roi_txt', metavar='save',
103              help='ROI file with rows (LPI): "label:\tx y z"',
104              type=str, default=None)
105      parser.add_argument('--roi_nii', metavar='save',
106              help='ROI nifti file (unimplemented)',
107              type=str, default=None)
108      parser.add_argument('--width', type=float,
109                          dest='width',
110                          default=4,
111                          help="image width")
112      parser.add_argument('--height', type=float,
113                          dest='height',
114                          default=2,
115                          help="image height")
116      parser.add_argument('--display',  type=str,
117                          dest='display_mode',
118                          default='xz',
119                          help="'ortho', 'x', 'y', 'z', 'xz', 'yx', 'yz', 'l', 'r', 'lr', 'lzr', 'lyr', 'lzry', 'lyrz'.")
120  
121      args = parser.parse_args()
122      print(args)
123      main(args)
124  
125  # Local Variables:
126  # python-shell-interpreter: "ipython3"
127  # python-shell-interpreter-args: "-i --simple-prompt --InteractiveShell.display_page=True"
128  # End: