blink_vols.py

import numpy as np
from os.path import basename
from glob import glob
from skimage.color import label2rgb
import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from tomopy.misc.phantom import shepp3d


class VolumeBlinker(object):
    
    def __init__(self, vols, label_vols=None, alpha=0.2):
        self.vols = vols
        self.label_vols = label_vols
        self.alpha = alpha

        self.dim = 0
        self.slice_nums = [int(round(0.5*s)) for s in vols[0].shape]
        self.current_vol_num = 0

        # create initial plot
        self.figure = plt.figure(figsize=(7,7))
        self.ax = plt.subplot(111)
        self.update_image()
        
        # disable default keybindings
        manager, canvas = self.figure.canvas.manager, self.figure.canvas  
        canvas.mpl_disconnect(manager.key_press_handler_id)

        # set callbacks
        self.cid = self.figure.canvas.mpl_connect('key_press_event', self._on_keypress)
        self.cid2 = self.figure.canvas.mpl_connect('scroll_event', self._on_mouse_scroll)

        plt.ion()
        plt.show()
        self.figure.canvas.start_event_loop(timeout=0)
        

    def slice_volumes(self):
        slc_num = self.slice_nums[self.dim]
        inds = [slice(None, None, None), 
                slice(None, None, None), 
                slice(None, None, None)]
        inds[self.dim] = self.slice_nums[self.dim]
        vol = self.vols[self.current_vol_num]
        label_vol = self.label_vols[self.current_vol_num]
        slc = vol[inds]
        
        if label_vol is not None:
            label_slc = label_vol[inds]
            slc = label2rgb(label_slc, image=slc, alpha=self.alpha, bg_label=0)
        return slc

    def update_image(self):
        slc = self.slice_volumes()
        try:
            self.im_plot.set_data(slc)
        except AttributeError:
            self.im_plot = self.ax.imshow(slc, interpolation='none', cmap='gray')
            plt.axis('off')
            self.ax.axis('off')

    def _on_mouse_scroll(self, event):
        self.slice_nums[self.dim] += event.step
        self.slice_nums[self.dim] = self.slice_nums[self.dim] % self.vols[0].shape[self.dim]
        self.update_image()
    
    def next_slice(self):
        self.slice_nums[self.dim] += 1
        self.slice_nums[self.dim] = self.slice_nums[self.dim] % self.vols[0].shape[self.dim]
        self.update_image()
    
    def previous_slice(self):
        self.slice_nums[self.dim] -= 1
        self.slice_nums[self.dim] = self.slice_nums[self.dim] % self.vols[0].shape[self.dim]
        self.update_image()
    
    def next_volume(self):
        self.current_vol_num += 1
        self.current_vol_num = self.current_vol_num % len(self.vols)
        self.update_image()
    
    def previous_volume(self):
        self.current_vol_num -= 1
        self.current_vol_num = self.current_vol_num % len(self.vols)
        self.update_image()
    
    def toggle_dimension(self):
        self.dim = (self.dim + 1) % 3
        self.update_image()
        
    def _on_keypress(self, event):
        if event.key.lower() in ['q', 'escape']:
            plt.close()
            self.figure.canvas.stop_event_loop()
        elif event.key == 'up':
            self.next_slice()
        elif event.key == 'down':
            self.previous_slice()
        elif event.key == 'right':
            self.next_volume()
        elif event.key == 'left':
            self.previous_volume()
        elif event.key == 'd':
            self.toggle_dimension()
        else:
            pass
            # print 'pressed', event.key



if __name__ == '__main__':

    from skimage.filters import threshold_otsu
    from skimage import measure

    vol_shape = (128, 128, 128)
    vol1 = np.random.randn(*vol_shape)
    vol2 = shepp3d(vol_shape)
    vols = [vol1, vol2]

    labels2 = measure.label(vol2 > 0.1)
    labels = [None, labels2]

    comparer = VolumeBlinker(vols, labels)  # blocks until quit