From 3c0caae1e0f449ae4e97a877a4c833b10b62792b Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Wed, 22 Feb 2017 16:57:06 -0500
Subject: [PATCH] Make coordinate display work for NIFTI files

---
 glue_medical/medical_coordinates.py | 41 ++++++++++++++++++-----------
 glue_medical/nifti/nifti_factory.py | 28 +++++++++++++++-----
 2 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/glue_medical/medical_coordinates.py b/glue_medical/medical_coordinates.py
index aa14e0d..f1c32b7 100644
--- a/glue_medical/medical_coordinates.py
+++ b/glue_medical/medical_coordinates.py
@@ -1,27 +1,36 @@
 import numpy as np
-
+from astropy.wcs import WCS
 from glue.core.coordinates import Coordinates
 
-class Coordinates4DMatrix(Coordinates):
-
-	def __init__(self, matrix, axis_labels):
 
-		# Other code to translate header into labels depending on file format.
-
-		self.matrix = matrix
-		self.matrix_invert = np.linalg.inv(matrix)
-		self.axis_labels = axis_labels
-
-	def axis_label(self, axis):
-		return self.axis_labels[axis]
+class Coordinates4DMatrix(Coordinates):
 
-	def pixel2world(self, *args):
+    def __init__(self, matrix, axis_labels):
 
-		return np.matmul(self.matrix, args + (np.zeros_like(args[0]),))[:-1]
+        # Other code to translate header into labels depending on file format.
 
-	def world2pixel(self, *args):
+        self.matrix = matrix
+        self.matrix_invert = np.linalg.inv(matrix)
+        self.axis_labels = axis_labels
 
-		return np.matmul(self.matrix_invert, args + (np.zeros_like(args[0]),))[:-1]		
+    def axis_label(self, axis):
+        return self.axis_labels[axis]
 
+    def pixel2world(self, *args):
+        return np.matmul(self.matrix, args + (np.zeros_like(args[0]),))[:-1]
 
+    def world2pixel(self, *args):
+        return np.matmul(self.matrix_invert, args + (np.zeros_like(args[0]),))[:-1]
 
+    @property
+    def wcs(self):
+        # We provide a wcs property since this can then be used by glue
+        # to display world coordinates. In this case, the transformation matrix
+        # is in the same order as the WCS convention so we don't need to swap
+        # anything.
+        wcs = WCS(naxis=self.matrix.shape[0] - 1)
+        wcs.wcs.cd = self.matrix[:-1, :-1]
+        wcs.wcs.crpix = np.zeros(wcs.naxis)
+        wcs.wcs.crval = self.matrix[:-1, -1]
+        wcs.wcs.ctype = self.axis_labels[::-1]
+        return wcs
diff --git a/glue_medical/nifti/nifti_factory.py b/glue_medical/nifti/nifti_factory.py
index ce2799d..49e4eea 100644
--- a/glue_medical/nifti/nifti_factory.py
+++ b/glue_medical/nifti/nifti_factory.py
@@ -1,16 +1,13 @@
 from __future__ import absolute_import, division, print_function
 
 import os
-import glob
 
 import nibabel as nib
-import numpy as np
 
-from glue.logger import logger
 from glue.core.data import Data
 from glue.config import data_factory
 from ..medical_coordinates import Coordinates4DMatrix
-
+from ..utils import flip
 
 __all__ = ['is_nifti_file', 'nifti_reader']
 
@@ -39,22 +36,41 @@ def nifti_label(filename):
     label = label.split('.nii')[0]
     return label
 
+
 @data_factory(
     label='NIFTI file (.nii or .nii.gz)',
     identifier=is_nifti_file,
     priority=100,
 )
 def nifti_reader(filepath):
-
     """
     Reads in a NIFTI file. Uses an affine matrix extracted from nibabel to perform coordinate changes.
     """
 
     nifti_data = nib.load(filepath)
     matrix = nifti_data.affine
-    axis_labels = ['Saggital','Coronal','Axial']
+    axis_labels = ['Axial', 'Coronal', 'Saggital']
     array = nifti_data.get_data()
 
+    # NIFTI files are stored with the array indices transposed from the usual
+    # Numpy convention of having the first dimension as the last index. But
+    # before we transpose the data, we should go through and check whether
+    # the matrix has negative diagonal elements, since for those we can flip
+    # the data array and adjust the transformation matrix - this will ensure
+    # that coordinates always increase to the right/top.
+    # Check for negative diagonal matrix elements and flip the data and adjust the
+    # transformation
+    for idx in range(array.ndim):
+        if matrix[idx, idx] < 0:
+            matrix[idx, :] = -matrix[idx, :]
+            matrix[-1, :] += array.shape[idx] * matrix[idx, :]
+            array = flip(array, idx)
+
+    # We now transpose the array so that it has the expected Numpy axis order
+    array = array.transpose()
+
+    # Set up the coordinate object using the matrix - for this we keep the
+    # matrix in the original order (as well as the axis labels)
     coords = Coordinates4DMatrix(matrix, axis_labels)
 
     data = [Data(label=nifti_label(filepath))]