Merge pull request #26 from Zhitao-Li/master

Add colormap to plot.ImagePlot
mikgroup · Jan 31, 2020 · 6364088 · 6364088
2 parents 348f77f + ca9a32b
commit 6364088
Showing 1 changed file with 22 additions and 5 deletions.
diff --git a/sigpy/plot.py b/sigpy/plot.py
@@ -73,6 +73,9 @@ def __init__(
             c=None,
             hide_axes=False,
             mode=None,
+            colormap=None,
+            vmin_mag=None,
+            vmax_mag=None,
             title='',
             interpolation='nearest',
             save_basename='Figure',
@@ -100,9 +103,12 @@ def __init__(
         self.title = title
         self.interpolation = interpolation
         self.mode = mode
+        self.colormap = colormap
         self.entering_slice = False
-        self.vmin = None
-        self.vmax = None
+        self.vmin = vmin_mag
+        self.vmax = vmax_mag
+        self.vmin_mag = vmin_mag
+        self.vmax_mag = vmax_mag
         self.save_basename = save_basename
         self.fps = fps
         self.help_text = None
@@ -251,8 +257,12 @@ def key_press(self, event):
             self.fig.canvas.draw()
 
         elif event.key in ['m', 'p', 'r', 'i', 'l']:
-            self.vmin = None
-            self.vmax = None
+            if event.key in ['m', 'r', 'i']:
+                self.vmin = self.vmin_mag
+                self.vmax = self.vmax_mag
+            else:
+                self.vmin = None
+                self.vmax = None
             self.mode = event.key
 
             self.update_axes()
@@ -423,11 +433,15 @@ def update_image(self):
             self.vmax = imv.max()
 
         if self.axim is None:
+            if self.colormap is None:
+                colormap = 'gray'
+            else:
+                colormap = self.colormap
             self.axim = self.ax.imshow(
                 imv,
                 vmin=self.vmin,
                 vmax=self.vmax,
-                cmap='gray',
+                cmap=colormap,
                 origin='lower',
                 interpolation=self.interpolation,
                 aspect=1.0,
@@ -437,6 +451,9 @@ def update_image(self):
                     0,
                     imv.shape[0]])
 
+            if self.colormap is not None:
+                self.fig.colorbar(self.axim)
+
         else:
             self.axim.set_data(imv)
             self.axim.set_extent([0, imv.shape[1], 0, imv.shape[0]])