PexSegment.py

"""Classes and methods for segmentation of spherical objects within cells."""


# IMPORT DEPENDENCIES
import matplotlib
import os
import sys
import pickle
import time
from operator import itemgetter
import numpy as np
import pandas as pd
from skimage import io
from skimage.morphology import watershed
from skimage.feature import canny
from scipy.ndimage.filters import gaussian_filter, maximum_filter
from scipy.ndimage.morphology import generate_binary_structure, binary_closing
from scipy.ndimage.morphology import distance_transform_edt
from scipy.ndimage.morphology import binary_erosion, binary_dilation
from scipy.ndimage.morphology import binary_fill_holes, binary_opening
from scipy.ndimage import generic_gradient_magnitude, sobel
import matplotlib.pyplot as plt
matplotlib.use('Agg')


class PexSegmentObj:
    '''A container class for objects generated by PexSegmenter.segment().

    Objects of class PexSegmentObj contain a raw multipage TIFF image in a
    numpy ndarray format, a similarly formatted image with the segmentation
    output from PexSegmenter.segment(), a number of segmentation intermediates
    that may be useful for diagnosis of segmentation problems, and relevant
    metadata that may be useful for analysis. Class methods for output of image
    data in TIFF format, for saving of object data in csv format, and for
    pickling the entire object are included.

    IMPORTANT: Do not call this class on its own. Objects of this class are
    generated as output from PexSegmenter.segment(), which provides all of the
    parameters described in __init__.

    Args !!!IMPORTANT: DO NOT CALL! Passed by PexSegmenter.segment()!!!
        f_directory (str): Path to the raw image used for segmentation.
        filename (str): raw image filename.
        raw_img (np.ndarray of ints): pixel intensity values of raw input image
            used as a starting point for segmentation. Each array position
            represents a single pixel.
        gaussian_img (np.ndarray of ints): pixel intensity values of gaussian
            filter output from PexSegmenter.segment().
        seg_method (str): Segmentation method provided to
            PexSegmenter.__init__().
        mode (str): Segmentation mode provided to PexSegmenter.__init__().
        threshold_img (binary np.ndarray): Binary array of pixels corresponding
            to a segmented object.
        dist_map (np.ndarray of ints): ndarray of the Euclidean distance of
            each pixel marked as 1 in threshold_img to a background pixel (a
            0). Generated by PexSegmenter.segment().
        smooth_dist_map (np.ndarray of ints): Smoothed distance map generated
            by PexSegmenter.segment().
        maxima (binary np.ndarray): Local maxima from the smooth_dist_map,
            generated by PexSegmenter.segment().
        labs (np.ndarray of ints): Labeled starting points for watershed
            segmentation. Generated by PexSegmenter.segment().
        watershed_output (np.ndarray of ints): Segmented objects, the final output
            from PexSegmenter.segment(). Each array position represents a
            single pixel, and all indices with a given integer value make up
            one segmented object. 0 represents background.
        obj_nums (list of ints): The list of numerical values assigned to
            segmented objects.
        volumes (dict): A dictionary with obj_num:volume pairs, with volume
            being the size of the segmented object in pixels.
        to_pdout (list of strings): The names of the attributes to be passed to
            to_csv() and to_pandas() for output. Assigned by
            PexSegmenter.segment().
        mode_params (dict): Dict of additional variables from
            PexSegmenter.segment() passed to PexSegmentObj for inclusion in
            attributes. Varies depending upon segmentation method. See
            PexSegmenter.segment() code for details.

    Attributes:
        f_directory (str): Path to the raw image used for segmentation.
        filename (str): raw image filename.
        raw_img (np.ndarray of ints): pixel intensity values of raw input image
            used as a starting point for segmentation. Each array position
            represents a single pixel.
        gaussian_img (np.ndarray of ints): pixel intensity values of gaussian
            filter output from PexSegmenter.segment().
        seg_method (str): Segmentation method provided to
            PexSegmenter.__init__().
        mode (str): Segmentation mode provided to PexSegmenter.__init__().
        threshold_img (binary np.ndarray): Binary array of pixels corresponding
            to a segmented object.
        dist_map (np.ndarray of ints): ndarray of the Euclidean distance of
            each pixel marked as 1 in threshold_img to a background pixel (a
            0). Generated by PexSegmenter.segment().
        smooth_dist_map (np.ndarray of ints): Smoothed distance map generated
            by PexSegmenter.segment().
        maxima (binary np.ndarray): Local maxima from the smooth_dist_map,
            generated by PexSegmenter.segment().
        labs (np.ndarray of ints): Labeled starting points for watershed
            segmentation. Generated by PexSegmenter.segment().
        peroxisomes (np.ndarray of ints): Segmented objects, the final output
            from PexSegmenter.segment(). Each array position represents a
            single pixel, and all indices with a given integer value make up
            one segmented object. 0 represents background.
        slices (int): The number of Z-slices in the image.
        height (int): The Y-direction image size in pixels.
        width (int): The X-direction image size in pixels.
        obj_nums (list of ints): The list of numerical values assigned to
            segmented objects.
        npexs (int): The number of segmented objects.
        volumes (dict): A dictionary with obj_num:volume pairs, with volume
            being the size of the segmented object in pixels.
        volumes_flag (str): The units for volume measurement. Always assigned
            as 'pixels'.
        pdout (list of strings): The names of the attributes to be passed to
            to_csv() and to_pandas() for output. Assigned by
            PexSegmenter.segment().
        border_rm_flag (bool): Indication of whether or not objects on the edge
            of the image are removed. Defaults to False.
    '''
    def __init__(self, f_directory, filename, raw_img, gaussian_img,
                 seg_method, mode, threshold_img, dist_map,
                 smooth_dist_map, maxima, labs, watershed_output,
                 obj_nums, volumes, to_pdout = [],
                 mode_params = {}):
        '''Initialize the PexSegmentObj with segmentation data.'''
        print('creating PexSegmentObj...')
        self.f_directory = f_directory
        self.filename = os.path.basename(filename).lower()
        self.raw_img = raw_img.astype('uint16')
        self.gaussian_img = gaussian_img.astype('uint16')
        self.seg_method = seg_method
        self.mode = mode
        self.threshold_img = threshold_img.astype('uint16')
        self.dist_map = dist_map.astype('uint16')
        self.smooth_dist_map = smooth_dist_map.astype('uint16')
        self.maxima = maxima.astype('uint16')
        self.labs = labs.astype('uint16')
        self.peroxisomes = watershed_output.astype('uint16')
        self.slices = self.raw_img.shape[0]
        self.height = self.raw_img.shape[1]
        self.width = self.raw_img.shape[2]
        self.obj_nums = obj_nums
        self.npexs = len(self.obj_nums)
        self.volumes = volumes
        self.volumes_flag = 'pixels'
        self.pdout = []
        self.border_rm_flag = False
        for key in mode_params:
            if hasattr(self, key):
                # raise an error if an attribute is somehow passed twice
                raise AttributeError('Two copies of the attribute ' + key +
                                     'were provided to PexSegmentObj.__init__()')
            setattr(self, key, mode_params[key])
        if to_pdout != []:
            for x in to_pdout:
                self.pdout.append(x)

    def __repr__(self):
        return 'PexSegmentObj '+ self.filename

    def plot_raw_img(self, display = False):
        '''Plot the raw image using matplotlib.'''
        self.plot_stack(self.raw_img, colormap = 'gray')
        if display == True:
            plt.show()
    def plot_gaussian_img(self, display = False):
        '''Plot the gaussian image using matplotlib.'''
        self.plot_stack(self.gaussian_img, colormap = 'gray')
        if display == True:
            plt.show()
    def plot_threshold_img(self, display = False):
        '''Plot the threshold image using matplotlib.'''
        self.plot_stack(self.threshold_img, colormap = 'gray')
        if display == True:
            plt.show()
    def plot_dist_map(self, display = False):
        '''Plot the distance map image using matplotlib.'''
        self.plot_stack(self.dist_map)
        if display == True:
            plt.show()
    def plot_smooth_dist_map(self, display = False):
        '''Plot the smoothed distance map image using matplotlib.'''
        self.plot_stack(self.smooth_dist_map)
        if display == True:
            plt.show()
    def plot_maxima(self, display = False):
        '''Plot the maxima image using matplotlib.'''
        # expand maxima to make them more easily visible in the output.
        vis_maxima = binary_dilation(self.maxima,
                                     structure = np.ones(shape = (1,5,5)))
        masked_maxima = np.ma.masked_where(vis_maxima == 0, vis_maxima)
        self.plot_maxima_stack(masked_maxima, self.smooth_dist_map)
        if display == True:
            plt.show()
    def plot_watershed(self, display = False):
         '''Plot the segmented objects image using matplotlib.'''
         self.plot_stack(self.peroxisomes)
         if display == True:
             plt.show()

    def output_all_images(self, output_dir = None):
        '''Write all images to a new directory.

        Write all images associated with the PexSegmentObj to a new directory.
        Name that directory according to the filename of the initial image that
        the object was derived from, unless an output directory is provided.

        Args:
            output_dir (str, optional): The directory to output image files to.
                If not provided, the images will be a subdirectory to the
                directory containing the initial raw image, which is named based
                on the filename of the raw image.
        '''
        if output_dir == None:
            # name as described in args
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        # make the output directory if it doesn't already exist
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        print('writing images...')
        # save numpy ndarrays as tif images using skimage.io.imsave
        io.imsave('raw_'+self.filename, self.raw_img)
        io.imsave('gaussian_'+self.filename, self.gaussian_img)
        io.imsave('threshold_'+self.filename, self.threshold_img)
        io.imsave('dist_'+self.filename, self.dist_map)
        io.imsave('smooth_dist_'+self.filename,self.smooth_dist_map)
        io.imsave('maxima_'+self.filename,self.maxima)
        io.imsave('wshed_'+self.filename,self.peroxisomes)
        if hasattr(self,'edges'):
            io.imsave('edges_'+self.filename,self.edges)

    def output_image(self, imageattr, output_dir = None):
        '''Write one specific image attribute to a new directory.

        Write an image associated with the PexSegmentObj to a new directory.
        Name that directory according to the filename of the initial image that
        the object was derived from, unless an output directory is provided.

        Args:
            imageattr (str): The name of the image attribute to be saved.
            output_dir (str, optional): The directory to output image files to.
                If not provided, the images will be a subdirectory to the
                directory containing the initial raw image, which is named based
                on the filename of the raw image.
        '''
        if output_dir == None:
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        # if the output directory doesn't already exist, make it.
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        print('writing image' + str(imageattr))
        # save the ndarray as a tif image using skimage.io.imsave
        io.imsave(str(imageattr)+self.filename, getattr(self,str(imageattr)))

    def output_plots(self):
        '''Write PDFs of slice-by-slice plots.

        Output: PDF plots of each image within PexSegmentObj in a directory
            named for the original filename they were generated from. Plots are
            generated using the plot_stack method and plotting methods defined
            here.
        '''
        os.chdir(self.f_directory)
        if not os.path.isdir(self.f_directory + '/' +
                             self.filename[0:self.filename.index('.tif')]):
            print('creating output directory...')
            os.mkdir(self.f_directory + '/' +
                     self.filename[0:self.filename.index('.tif')])
        os.chdir(self.f_directory + '/' +
                 self.filename[0:self.filename.index('.tif')])
        print('saving plots...')
        self.plot_raw_img()
        plt.savefig('praw_'+self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_gaussian_img()
        plt.savefig('pgaussian_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_threshold_img()
        plt.savefig('pthreshold_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        plt.savefig('pdist_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_smooth_dist_map()
        plt.savefig('psmooth_dist_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_maxima()
        plt.savefig('pmaxima_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
        self.plot_watershed()
        plt.savefig('pwshed_' +
                    self.filename[0:self.filename.index('.tif')]+'.pdf')
    def pickle(self, output_dir = None, filename = None):
        '''pickle the PexSegmentObj for later loading.'''
        if output_dir == None:
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        if filename == None:
            filename = self.filename[0:self.filename.index('.tif')] + '.pickle'
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        with open(filename, 'wb') as f:
            pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
        f.close()
    def output_all(self):
        '''Output images as tifs and plots as pdfs, and pickle object.'''
        os.chdir(self.f_directory)
        # make output directories if they don't exist
        if not os.path.isdir(self.f_directory + '/' +
                             self.filename[0:self.filename.index('.tif')]):
            os.mkdir(self.f_directory + '/' +
                     self.filename[0:self.filename.index('.tif')])
        os.chdir(self.f_directory + '/' +
                 self.filename[0:self.filename.index('.tif')])
        print('outputting all data...')
        self.output_plots()
        self.output_all_images()
        self.pickle()
        # TODO: UPDATE THIS METHOD TO INCLUDE PANDAS OUTPUT
    def to_csv(self, output_dir = None):
        '''Output attributes designated by self.pdout in csv format.

        Args:
            output_dir (str, optional): Directory to output csv to. If not
            provided, a new subdirectory is created within the directory
            containing the raw image named based on the raw image name, and the
            csv is saved there.
        '''
        os.chdir(self.f_directory)
        if output_dir == None:
            output_dir = self.f_directory + '/' + self.filename[0:self.filename.index('.tif')]
        if not os.path.isdir(output_dir):
            print('creating output directory...')
            os.mkdir(output_dir)
        os.chdir(output_dir)
        for_csv = self.to_pandas()
        for_csv.to_csv(path_or_buf = output_dir + '/' +
                       self.filename[0:self.filename.index('.tif')]+ '.csv',
                       index = True, header = True)
    def rm_border_objs(self, border = 1, z = True):
        '''remove all objects that contact the edge of the 3D stack.

        Args:
            border (int, optional): the size of the border around the edge which a pixel from
               the object must contact to be removed. Defaults to 1.
            z (bool, optional): should objects that contact the z-axis edge be eliminated? if
               True, any object with a pixel in the top or bottom image of the
               stack is removed. Defaults to True.

        Output: alters the objects within the PexSegmentObj. removes objects
        from the peroxisomes image, the obj_nums, and all other variables with
        elements of obj_nums as keys in a dict (parents, volumes, etc)
        '''

        border_mask = np.full(shape = self.peroxisomes.shape, fill_value = True,
                              dtype = bool)
        if z == True:
            border_mask[border:-border,border:-border,border:-border] = False
        elif z == False:
            border_mask[:,border:-border,border:-border] = False
        objs_to_rm = np.unique(self.peroxisomes[border_mask])
        objs_to_rm = objs_to_rm[objs_to_rm != 0]
        for x in objs_to_rm:
            self.peroxisomes[self.peroxisomes == x] = 0
            self.obj_nums.remove(x)
            self.volumes.pop(x, None)
            if hasattr(self, "parent"):
                self.parent.pop(x, None)
        self.npexs = len(self.obj_nums)
        self.border_rm_flag = True

    # HELPER METHODS #

    def to_pandas(self):
        '''create a pandas DataFrame of tabulated numeric data.

        the pdout attribute indicates which variables to include in the
        DataFrame.

        Helper method for self.to_csv().
        '''
        df_dict = {}
        for attr in self.pdout:
            df_dict[str(attr)] = pd.Series(getattr(self, attr))
        if 'volumes' in self.pdout:
            vflag_out = dict(zip(self.obj_nums,
                                 [self.volumes_flag]*len(self.obj_nums)))
            df_dict['volumes_flag'] = pd.Series(vflag_out)
        return pd.DataFrame(df_dict)
    def convert_volumes(self, z = 0.2, x = 0.0675):
        '''convert volumes from units of pixels to metric units.

        The default values provided correspond to the appropriate values for
        imaging using the Murray spinning disk confocal microscope with 0.2 um
        spacing in between slices.

        Args:
            z (float, optional): the distance between slices in the z-stack in
                units of microns.
            x (float, optional): the linear distance between adjacent pixels in
                each slice in units of microns. x is also used for y.

        Output: converts self.volumes to units of femtoliters, and changes the
            self.volumes_flag to 'femtoliters'.
        '''
        conv_factor = z*x*x
        for key, val in self.volumes:
            self.volumes[key] = float(self.volumes[key])*conv_factor
        self.volumes_flag = 'femtoliters'

    @staticmethod
    def plot_stack(stack_arr, colormap='jet'):
        ''' Create a matplotlib plot with each subplot containing a slice.

        Args:
            stack_arr (ndarray of ints): A numpy ndarray containing pixel
                intensity values.
            colormap: The colormap to be used when displaying pixel intensities.
                Defaults to jet.

        Yield: a pyplot object in which each slice from the image array
            is represented in a subplot. subplots are 4 columns
            across (when 4 or more slices are present) with rows to
            accommodate all slices.
        '''

        nimgs = stack_arr.shape[0] # z axis of array dictates number of slices
        # plot with 4 imgs across

        # determine how many rows and columns of images there are

        if nimgs < 5:
            f, axarr = plt.subplots(1,nimgs)

            for i in range(0,nimgs):
                axarr[i].imshow(stack_arr[i,:,:], cmap=colormap)
                axarr[i].xaxis.set_visible(False)
                axarr[i].yaxis.set_visible(False)
            f.set_figwidth(16)
            f.set_figheight(4)

        else:
            f, axarr = plt.subplots(int(np.ceil(nimgs/4)),4)

            for i in range(0,nimgs):
                r = int(np.floor(i/4))
                c = int(i % 4)
                axarr[r,c].imshow(stack_arr[i,:,:], cmap=colormap)
                axarr[r,c].xaxis.set_visible(False)
                axarr[r,c].yaxis.set_visible(False)
        # manage remainder images
            if nimgs%4 > 0:
                r = int(np.floor(nimgs/4))

                for c in range(nimgs%4,4):
                    axarr[r,c].axis('off')

            f.set_figwidth(16)
            f.set_figheight(4*np.ceil(nimgs/4))

    @staticmethod
    def plot_maxima_stack(masked_max, smooth_dist):

        ''' Creates a matplotlib plot object in which each slice from the image
        is displayed as a single subplot, in a 4-by-n matrix (n depends upon
        the number of slices in the image)'''

        nimgs = masked_max.shape[0] # z axis of array dictates number of slices
        # plot with 4 imgs across

        # determine how many rows and columns of images there are

        if nimgs < 5:
            f, axarr = plt.subplots(1,nimgs)

            for i in range(0,nimgs):
                axarr[i].imshow(smooth_dist[i,:,:], cmap='gray')
                axarr[i].imshow(masked_max[i,:,:], cmap='autumn')
                axarr[i].xaxis.set_visible(False)
                axarr[i].yaxis.set_visible(False)
            f.set_figwidth(16)
            f.set_figheight(4)

        else:
            f, axarr = plt.subplots(int(np.ceil(nimgs/4)),4)

            for i in range(0,nimgs):
                r = int(np.floor(i/4))
                c = int(i%4)
                axarr[r,c].imshow(smooth_dist[i,:,:], cmap='gray')
                axarr[r,c].imshow(masked_max[i,:,:], cmap='autumn')
                axarr[r,c].xaxis.set_visible(False)
                axarr[r,c].yaxis.set_visible(False)

            if nimgs%4 > 0:
                r = int(np.floor(nimgs/4))

                for c in range(nimgs%4, 4):
                    axarr[r,c].axis('off')

            f.set_figwidth(16)
            f.set_figheight(4*np.ceil(nimgs/4))


class PexSegmenter:
    '''Class and methods for segmenting foci from multipage TIFF images.

    When called, generates an object of class PexSegmenter with the
    segmentation parameters defined as per Args for the image indicated. The
    segment() method can then be called to perform segmentation, returning a
    PexSegmentObj.

    Args:
        filename (str): The filename of the multipage TIFF-format image to be
        segmented. This is defined with respect to the current working
        directory.
        seg_method (str, optional): The segmentation approach to use in a string
        format. This can be one of the following:
            threshold (default): Segmentation by setting an absolute cutoff for
                minimum pixel intensity of an object. Requires kwargs dependent
                upon the segmentation mode (see mode below).
            canny: Segmentation by Canny edge detection. Requires kwargs
                high_threshold and low_threshold to be set, or the default
                values 1000 and 500 will be used.
        mode (str, optional): The segmentation mode to be used if seg_method ==
            threshold. This can be one of the following:
            threshold (default): Uses a user-provided pixel intensity value as
                the segmentation threshold. Requires kwarg threshold.
            bg_scaled: Use the median value from all pixels that correspond to
                cells within a CellSegmentObj (see CellSegment.py in this
                module) segmented from another fluorescence channel as the
                background. Requires kwargs cells and bg_diff.
        g_xy(int, optional): Standard deviation in units of pixels (xy direction
            only) for the 3D gaussian smoothing prior to segmentation. Defaults
            to 1 pixel. May need empirical optimization depending upon the size
            of desired particles and background noise.
        g_z(int, optional): Standard deviation in units of Z-slices (z direction
            only) for the 3D gaussian smoothing prior to segmentation. Defaults
            to 1 slice. May need empirical optimization depending upon the
            spacing between slices. Setting to 0 results in only 2D smoothing.

        Additional kwargs:
        threshold (int): Only relevant if using seg_method threshold and mode
            threshold. The pixel intensity value to be used as a cutoff for
            segmented objects.
        cells (str): Only relevant if using seg_method threshold and mode
            bg_scaled. The name of a CellSegmentObj already present in the
            environment which corresponds to cells segmented from the same
            field and in a separate channel. See CellSegment.py for details.
        bg_diff (int): Only relevant if using seg_method threshold and mode
            bg_scaled. The desired pixel intensity units above the cells'
            median pixel intensity to set the threshold.
        high_threshold (int, optional): Only relevant if using seg_method
            "canny". The high threshold to be passed to skimage's canny edge
            detector. Default value is 1000. See skimage.feature.canny
            documentation for more details.
        low_threshold (int, optional): Only relevant if using seg_method
            "canny". The low threshold to be passed to skimage's canny edge
            detector. Default value is 500. See skimage.feature.canny
            documentation for more details.

    Attributes:
        filename (str): The filename provided in args.
        seg_method (str): The segmentation method provided in args, defaults to
            'threshold'.
        mode (str): The segmentation mode provided in args, defaults to
            'threshold'.
        g_xy (int): The gaussian filter xy standard deviation in pixels.
            Defaults to 1.
        g_z (int): The gaussian filter z standard deviation in slices.
            Defaults to 1.

        Additional attributes provided by **kwargs (see Args, Additional kwargs
            above).

    '''

    def __init__(self, filename='', src_data=None, seg_method='threshold',
                 mode='threshold', g_xy=1, g_z=1, **kwargs):
        # set attributes
        if src_data is None and filename == '':
            raise ValueError('no input image provided.')
        self.filename = filename
        self.src_data = src_data
        self.seg_method = seg_method
        self.mode = mode
        self.g_xy = g_xy
        self.g_z = g_z
        for key, value in kwargs.items():
            setattr(self, key, value)
        if self.seg_method == 'canny':
            if not self.mode == 'absolute' or self.mode == 'scaled':
                self.mode = 'absolute'  # if it wasn't set, use absolute
            self.high_threshold = kwargs.get('high_threshold',1000)
            self.low_threshold = kwargs.get('low_threshold',500)
        if self.seg_method == 'threshold':
            if mode == 'threshold':
                self.threshold = kwargs.get('threshold',float('nan'))
                if np.isnan(self.threshold):
                    raise ValueError('A threshold argument must be provided to segment with a constant threshold.')
            if mode == 'bg_scaled':
                self.cells = kwargs.get('cells', '')
                self.bg_diff = float(kwargs.get('bg_diff',float('nan')))
                if self.cells == '':
                    raise ValueError('A CellSegmentObj containing segmented cells is required if mode == bg_scaled.')
                if np.isnan(self.bg_diff):
                    raise ValueError('a bg_diff argument is needed if mode == bg_scaled.')
    def segment(self, fill_holes=False, edt_sampling=(3,1,1),
                edt_smooth=[1,3,3]):
        """Segment objects within the image according to attributes provided.

        Yields: a PexSegmentObj containing segmented objects as well as all
            images generated during segmentation (for post-hoc analysis) as
            well as relevant values, e.g. numbers and names of segmented
            particles. See PexSegmentObj documentation for more details.
        """
        starttime = time.time() # begin timing
        f_directory = os.getcwd()
        pdout = [] # list of PexSegmentObj attributes to pass to pandas for csv
        # data import
        if self.filename != '':
            print('reading' + self.filename)
            raw_img = io.imread(self.filename)
        elif self.src_data is not None:
            raw_img = self.src_data
        print('raw image imported.')
        if self.seg_method == 'pre-thresholded':
            gaussian_img = raw_img
        else:
            # gaussian filter
            print('performing gaussian filtering...')
            gaussian_img = gaussian_filter(raw_img,
                                           [self.g_z, self.g_xy, self.g_xy])
            print('Image smoothed.')
        print('preprocessing complete.')
        ## SEGMENTATION BY THRESHOLDING THE GAUSSIAN ##
        if self.seg_method == 'threshold':
            # binary thresholding and cleanup
            print('thresholding...')
            threshold_img = np.copy(gaussian_img)
            if self.mode == 'threshold':
                print('mode = threshold.')
                # make binary image
                threshold_img[threshold_img < self.threshold] = 0
                threshold_img[threshold_img > 0] = 1
                print('thresholding complete.')
                if fill_holes:
                    print('filling holes in objects.')
                    for i in range(0,threshold_img.shape[0]):
                        threshold_img[i, :, :] = binary_fill_holes(
                            threshold_img[i, :, :])
            elif self.mode == 'bg_scaled':
                print('mode = background-scaled.')
                self.thresholds = {}
                threshold_img = np.zeros(shape = raw_img.shape)
                for i in self.cells.obj_nums:
                    if i == 0:
                        pass
                    else:
                        print('thresholding cell ' + str(i))
                        # get median for the cell
                        cell_median = np.median(gaussian_img[self.cells.final_cells == i])
                        # generate the thresholded binary mask for each cell
                        threshold_img[np.logical_and(self.cells.final_cells == i,
                                      gaussian_img > cell_median + self.bg_diff)] = 1
                        self.thresholds[i] = cell_median + self.bg_diff #store val
                print('thresholding complete.')
            else:
                raise ValueError('mode parameter must be bg_scaled or threshold.')
            # distance and maxima transformation to find objects
            # next two steps assume 100x objective and 0.2 um slices
            print('generating distance map...')
            dist_map = distance_transform_edt(threshold_img, sampling = edt_sampling)
            print('distance map complete.')
            print('smoothing distance map...')
            # smooth the distance map
            smooth_dist = gaussian_filter(dist_map, edt_smooth)
            print('distance map smoothed.')
            print('identifying maxima...')
            # find local maxima in the smoothed distance map
            # these will be the watershed seeds
            max_strel = generate_binary_structure(3,2)
            maxima = maximum_filter(smooth_dist,
                                    footprint = max_strel) == smooth_dist
            # clean up background and edges
            bgrd_3d = smooth_dist == 0
            eroded_bgrd = binary_erosion(bgrd_3d, structure = max_strel,
                                         border_value = 1)
            maxima = np.logical_xor(maxima, eroded_bgrd)
            print('maxima identified.')
            # watershed segmentation
            labs = self.watershed_labels(maxima)
            print('watershedding...')
            peroxisomes = watershed(-smooth_dist, labs, mask = threshold_img)
            print('watershedding complete.')
            if self.mode == 'bg_scaled':
                # find cell boundaries and define objects that are on the
                # edges, then assign segmented objects to parent cells
                edge_struct = generate_binary_structure(3,1)
                self.c_edges = {}
                print('finding edges of cells...')
                for i in self.cells.obj_nums:
                    self.c_edges[i] = np.logical_xor(self.cells.final_cells == i,
                                                          binary_erosion(self.cells.final_cells== i,
                                                                         edge_struct))
                print('cell edges found.')
                self.primary_objs = [x for x in np.unique(peroxisomes) if x != 0]
                self.parent = {}
                self.obj_edges = {}
                self.on_edge = {}
                pex_mask = peroxisomes != 0
                for obj in self.primary_objs:
                    self.parent[obj] = self.cells.final_cells[labs == obj][0]
                    obj_mask = peroxisomes == obj
                    obj_edge = np.logical_xor(obj_mask,
                                              binary_erosion(obj_mask,
                                                             edge_struct))
                    self.obj_edges[obj] = obj_edge
                    # test if the object's edge and its cell's edge overlap
                    if np.any(np.logical_and(obj_edge,
                                             self.c_edges[self.parent[obj]])):
                        self.on_edge[obj] = True
                        print('object on the edge: ' + str(obj))
                        print('parent cell: ' + str(self.parent[obj]))
                        new_obj = obj_mask
                        search_obj = obj_mask
                        tester = 0
                        iteration = 1
                        while tester == 0:
                            # TODO: FIX THIS BLOCK OF CODE! GETTING STUCK WITHIN
                            # IT! NOT SURE HOW MANY ITERATIONS ITS DOING, OR FOR
                            # HOW MANY DIFFERENT PEROXISOMES.
                            new_px = binary_dilation(search_obj, edge_struct)
                            new_px[np.logical_or(new_obj, pex_mask)] = False
                            print('iteration: ' + str(iteration))
                            # print('new pixels for iteration ' + str(iteration) + \
                            #      ': ')
                            # print(np.nonzero(new_px))
                            if np.any(gaussian_img[new_px] >
                                      self.thresholds[self.parent[obj]]):
                                to_add = np.logical_and(new_px, gaussian_img >
                                                        self.thresholds[self.parent[obj]])
                                new_obj = np.logical_or(new_obj, to_add)
                            #    print('object pixels after iteration '
                            #          + str(iteration) + ': ')
                            #    print(np.nonzero(new_obj))
                                search_obj = to_add # only search from new pixels
                            else:
                                peroxisomes[new_obj] = obj
                                tester = 1
                            iteration = iteration + 1
                    else:
                        self.on_edge[obj] = False
        elif self.seg_method == 'canny':
            ## EDGE-DETECTION BASED SEGMENTATION ##
            threshold_img = np.empty_like(gaussian_img)
            edge_img = np.empty_like(gaussian_img)
            c_strel = generate_binary_structure(2,1)
            # perform canny edge detection on each slice s
            for s in range(0,gaussian_img.shape[0]):
                if self.mode == 'absolute':
                    c = canny(gaussian_img[s, :, :],
                              sigma=0,
                              low_threshold=self.low_threshold,
                              high_threshold=self.high_threshold)
                elif self.mode == 'scaled':
                    c = canny(gaussian_img[s, :, :],
                              sigma=0,
                              low_threshold=self.low_threshold,
                              high_threshold=self.high_threshold,
                              use_quantiles=True)
                # clean up object edges that have gaps
                c = binary_closing(c,c_strel)
                edge_img[s,:,:] = np.copy(c)
                # fill holes to generate binary mask of objects
                c = binary_fill_holes(c)
                c = binary_opening(c, c_strel) # eliminate incomplete lines
                threshold_img[s,:,:] = c
            print('generating distance map...')
            dist_map = distance_transform_edt(threshold_img, sampling = (3,1,1))
            print('distance map complete.')
            print('smoothing distance map...')
            smooth_dist = gaussian_filter(dist_map, [1,2,2])
            print('distance map smoothed.')
            print('identifying maxima...')
            max_strel = generate_binary_structure(3,2)
            # identify local maxima (these will be the seed points for
            # watershed segmentation)
            maxima = maximum_filter(smooth_dist,
                                    footprint = max_strel) == smooth_dist
            # clean up background and edges
            bgrd_3d = smooth_dist == 0
            eroded_bgrd = binary_erosion(bgrd_3d, structure = max_strel,
                                         border_value = 1)
            maxima = np.logical_xor(maxima, eroded_bgrd)
            print('maxima identified.')
            # watershed segmentation
            labs = self.watershed_labels(maxima)
            print('watershedding...')
            peroxisomes = watershed(-smooth_dist, labs, mask = threshold_img)
            print('watershedding complete.')
            if hasattr(self, 'cells'):
                # assign segmented objects to cells if a CellSegmentObj was
                # included
                self.primary_objs = [x for x in np.unique(peroxisomes) \
                                     if x != 0]
                self.parent = {}
                for obj in self.primary_objs:
                    o_parent = self.cells.final_cells[labs == obj][0]
                    if o_parent == 0:
                        self.primary_objs.remove(obj)
                    else:
                        self.parent[obj] = o_parent
        elif self.seg_method == 'pre-thresholded':
            threshold_img = np.copy(gaussian_img)
            if fill_holes:
                print('filling holes in objects.')
                for i in range(0, threshold_img.shape[0]):
                    threshold_img[i, :, :] = binary_fill_holes(
                        threshold_img[i, :, :])
                print('holes filled.')
            dist_map = distance_transform_edt(threshold_img,
                                              sampling=edt_sampling)
            print('distance map complete.')
            print('smoothing distance map...')
            # smooth the distance map
            smooth_dist = gaussian_filter(dist_map, edt_smooth)
            print('distance map smoothed.')
            print('identifying maxima...')
            # find local maxima in the smoothed distance map
            # these will be the watershed seeds
            max_strel = generate_binary_structure(3, 2)
            maxima = maximum_filter(smooth_dist,
                                    footprint=max_strel) == smooth_dist
            # clean up background and edges
            bgrd_3d = smooth_dist == 0
            eroded_bgrd = binary_erosion(bgrd_3d, structure= max_strel,
                                         border_value=1)
            maxima = np.logical_xor(maxima, eroded_bgrd)
            print('maxima identified.')
            # watershed segmentation
            labs = self.watershed_labels(maxima)
            print('watershedding...')
            peroxisomes = watershed(-smooth_dist, labs, mask=threshold_img)
            print('watershedding complete.')
        # Sometimes the watershedding algorithm inaccurately separates objects
        # on different Z-slices. The next section merges objects with
        # significant overlap
        for s in range(1,peroxisomes.shape[0]):
            cslice = peroxisomes[s,:,:]
            lslice = peroxisomes[s-1,:,:]
            for obj in np.unique(cslice)[np.unique(cslice)!= 0]:
                lslice_vals, cts = np.unique(lslice[cslice == obj],
                                             return_counts = True)
                lslice_vals = lslice_vals.tolist()
                cts = cts.tolist()
                ordered_by_ct = sorted(zip(lslice_vals, cts),
                                       key = itemgetter(1))
                if ordered_by_ct[-1][0] == 0 or ordered_by_ct[-1][0] == obj:
                    continue
                else:
                    # if >75% of pixels in the slice below obj are from another
                    # object, change obj to that object #
                    if float(ordered_by_ct[-1][1])/cslice[cslice == obj].size>0.5:
                        peroxisomes[s,:,:][cslice == obj] = ordered_by_ct[-1][0]
        obj_nums, volumes = np.unique(peroxisomes, return_counts=True)
        volumes = dict(zip(obj_nums.astype('uint16'), volumes))
        # remove the background
        del volumes[0]
        obj_nums = obj_nums.astype('uint16').tolist()
        obj_nums.remove(0)
        # generate dict of relevant parameters to pass to PexSegmentObj
        mode_params = {}
        if hasattr(self, 'parent'):
            pdout.append('parent')
            mode_params['parent'] = self.parent
        if self.seg_method == 'canny':
            mode_params['high_threshold'] = self.high_threshold
            mode_params['low_threshold'] = self.low_threshold
            mode_params['edges'] = edge_img
            pdout.append('volumes')
        if self.seg_method == 'threshold':
            if self.mode == 'threshold':
                mode_params['threshold'] = self.threshold
                pdout.append('volumes')
            elif self.mode == 'bg_scaled':
                mode_params['thresholds'] = self.thresholds
                mode_params['bg_diff'] = self.bg_diff
                mode_params['cells'] = self.cells
                mode_params['cell_edges'] = self.c_edges
                mode_params['cell_nums'] = self.cells.obj_nums
                mode_params['obj_edges'] = self.obj_edges
                mode_params['on_edge'] = self.on_edge
                for x in ['thresholds','on_edge','parent', 'volumes']:
                    pdout.append(x)
        return PexSegmentObj(f_directory, self.filename, raw_img,
                             gaussian_img, self.seg_method, self.mode,
                             threshold_img, dist_map, smooth_dist, maxima,
                             labs, peroxisomes, obj_nums, volumes,
                             to_pdout=pdout, mode_params=mode_params)

    ## HELPER METHODS ##
    @staticmethod
    def watershed_labels(maxima_img):
        '''Number local maxima in order for use in watershedding.

        Args:
            maxima_img (np.ndarray): A boolean array with local maxima labeled
            as true pixels.

        Yields:
            A numpy ndarray with maxima numbered sequentially.
        '''
        max_z, max_y, max_x = np.nonzero(maxima_img)
        label_output = np.zeros(maxima_img.shape)
        for i in range(0,len(max_y)):
            label_output[max_z[i],max_y[i],max_x[i]] = i+1
        return(label_output)