cell_magic_wand.py

###########################################################################
# 
# Python implementation of the ImageJ Cell Magic Wand plugin
# (http://www.maxplanckflorida.org/fitzpatricklab/software/cellMagicWand/)
# with modifications to reduce variability due to seed point selection
# and to support edge detection using all z slices of a 3D image
#
# Author: Noah Apthorpe (apthorpe@cs.princeton.edu)
#
# Description: Draws a border within a specified radius
#              around a specified center "seed" point
#              using a polar transform and a dynamic
#              programming edge-following algorithm
#
# Usage: Import and call the cell_magic_wand() function
#        or cell_magic_wand_3d () function with
#        a source image, radius window, and location of center
#        point.  Other parameters set as optional arguments.
#        Returns a binary mask with 1s inside the detected edge and
#        a list of points along the detected edge.
#
###########################################################################

import numpy as np
from scipy.ndimage.interpolation import zoom
from scipy.ndimage.morphology import binary_fill_holes


def coord_polar_to_cart(r, theta, center):
    '''Converts polar coordinates around center to Cartesian'''
    x = r * np.cos(theta) + center[0]
    y = r * np.sin(theta) + center[1]
    return x, y


def coord_cart_to_polar(x, y, center):
    '''Converts Cartesian coordinates to polar'''
    r = np.sqrt((x-center[0])**2 + (y-center[1])**2)
    theta = np.arctan2((y-center[1]), (x-center[0]))
    return r, theta


def image_cart_to_polar(image, center, min_radius, max_radius, phase_width, zoom_factor=1):
    '''Converts an image from cartesian to polar coordinates around center'''

    # Upsample image
    if zoom_factor != 1:
        image = zoom(image, (zoom_factor, zoom_factor), order=4)
        center = (center[0]*zoom_factor + zoom_factor/2, center[1]*zoom_factor + zoom_factor/2)
        min_radius = min_radius * zoom_factor
        max_radius = max_radius * zoom_factor
    
    # pad if necessary
    max_x, max_y = image.shape[0], image.shape[1]
    pad_dist_x = np.max([(center[0] + max_radius) - max_x, -(center[0] - max_radius)])
    pad_dist_y = np.max([(center[1] + max_radius) - max_y, -(center[1] - max_radius)])
    pad_dist = int(np.max([0, pad_dist_x, pad_dist_y]))
    if pad_dist != 0:
        image = np.pad(image, pad_dist, 'constant')

    # coordinate conversion
    theta, r = np.meshgrid(np.linspace(0, 2*np.pi, phase_width),
                           np.arange(min_radius, max_radius))
    x, y = coord_polar_to_cart(r, theta, center)
    x, y = np.round(x), np.round(y)
    x, y = x.astype(int), y.astype(int)
    
    polar = image[x, y]
    polar.reshape((max_radius - min_radius, phase_width))

    return polar


def mask_polar_to_cart(mask, center, min_radius, max_radius, output_shape, zoom_factor=1):
    '''Converts a polar binary mask to Cartesian and places in an image of zeros'''

    # Account for upsampling 
    if zoom_factor != 1:
        center = (center[0]*zoom_factor + zoom_factor/2, center[1]*zoom_factor + zoom_factor/2)
        min_radius = min_radius * zoom_factor
        max_radius = max_radius * zoom_factor
        output_shape = map(lambda a: a * zoom_factor, output_shape)

    # new image
    image = np.zeros(output_shape)

    # coordinate conversion
    theta, r = np.meshgrid(np.linspace(0, 2*np.pi, mask.shape[1]),
                           np.arange(0, max_radius))
    x, y = coord_polar_to_cart(r, theta, center)
    x, y = np.round(x), np.round(y)
    x, y = x.astype(int), y.astype(int)

    x = np.clip(x, 0, image.shape[0]-1)
    y = np.clip(y, 0, image.shape[1]-1)
    ix,iy = np.meshgrid(np.arange(0,mask.shape[1]), np.arange(0,mask.shape[0]))    
    image[x,y] = mask

    # downsample image
    if zoom_factor != 1:
        zf = 1/float(zoom_factor)
        image = zoom(image, (zf, zf), order=4)

    # ensure image remains a filled binary mask
    image = (image > 0.5).astype(int)
    image = binary_fill_holes(image)
    return image
    

def find_edge_2d(polar, min_radius):
    '''Dynamic programming algorithm to find edge given polar image'''
    if len(polar.shape) != 2:
        raise ValueError("argument to find_edge_2d must be 2D")

    # Dynamic programming phase
    values_right_shift = np.pad(polar, ((0, 0), (0, 1)), 'constant', constant_values=0)[:, 1:]
    values_closeright_shift = np.pad(polar, ((1, 0),(0, 1)), 'constant', constant_values=0)[:-1, 1:]
    values_awayright_shift = np.pad(polar, ((0, 1), (0, 1)), 'constant', constant_values=0)[1:, 1:]

    values_move = np.zeros((polar.shape[0], polar.shape[1], 3))
    values_move[:, :, 2] = np.add(polar, values_closeright_shift)  # closeright
    values_move[:, :, 1] = np.add(polar, values_right_shift)  # right
    values_move[:, :, 0] = np.add(polar, values_awayright_shift)  # awayright
    values = np.amax(values_move, axis=2)

    directions = np.argmax(values_move, axis=2)
    directions = np.subtract(directions, 1)
    directions = np.negative(directions)
        
    # Edge following phase
    edge = []
    mask = np.zeros(polar.shape)
    r_max, r = 0, 0
    for i,v in enumerate(values[:,0]):
        if v >= r_max:
            r, r_max = i, v
    edge.append((r+min_radius, 0))
    mask[0:r+1, 0] = 1
    for t in range(1,polar.shape[1]):
        r += directions[r, t-1]
        if r >= directions.shape[0]: r = directions.shape[0]-1
        if r < 0: r = 0
        edge.append((r+min_radius, t))
        mask[0:r+1, t] = 1

    # add to inside of mask accounting for min_radius
    new_mask = np.ones((min_radius+mask.shape[0], mask.shape[1]))
    new_mask[min_radius:, :] = mask
    
    return np.array(edge), new_mask


def edge_polar_to_cart(edge, center):
    '''Converts a list of polar edge points to a list of cartesian edge points'''
    cart_edge = [] 
    for (r,t) in edge:
        x, y = coord_polar_to_cart(r, t, center)
        cart_edge.append((round(x), round(y)))
    return cart_edge


def cell_magic_wand_single_point(image, center, min_radius, max_radius,
                                 roughness=2, zoom_factor=1):
    '''Draws a border within a specified radius around a specified center "seed" point
    using a polar transform and a dynamic programming edge-following algorithm.

    Returns a binary mask with 1s inside the detected edge and
    a list of points along the detected edge.'''
    if roughness < 1:
        roughness = 1
        print("roughness must be >= 1, setting roughness to 1")
    if min_radius < 0:
        min_radius = 0
        print("min_radius must be >=0, setting min_radius to 0")
    if max_radius <= min_radius:
        max_radius = min_radius + 1
        print("max_radius must be larger than min_radius, setting max_radius to " + str(max_radius))
    if zoom_factor <= 0:
        zoom_factor = 1
        print("negative zoom_factor not allowed, setting zoom_factor to 1")
    phase_width = int(2 * np.pi * max_radius * roughness)
    polar_image = image_cart_to_polar(image, center, min_radius, max_radius,
                                      phase_width=phase_width, zoom_factor=zoom_factor)
    polar_edge, polar_mask = find_edge_2d(polar_image, min_radius)
    cart_edge = edge_polar_to_cart(polar_edge, center)
    cart_mask = mask_polar_to_cart(polar_mask, center, min_radius, max_radius,
                                   image.shape, zoom_factor=zoom_factor)
    return cart_mask, cart_edge


def cell_magic_wand(image, center, min_radius, max_radius,
                    roughness=2, zoom_factor=1, center_range=2):
    '''Runs the cell magic wand tool on multiple points near the supplied center and 
    combines the results for a more robust edge detection then provided by the vanilla wand tool.

    Returns a binary mask with 1s inside detected edge'''
    
    centers = []
    for i in [-center_range, 0, center_range]:
        for j in [-center_range, 0, center_range]:
            centers.append((center[0]+i, center[1]+j))
    masks = np.zeros((image.shape[0], image.shape[1], len(centers)))
    for i, c in enumerate(centers):
        mask, edge = cell_magic_wand_single_point(image, c, min_radius, max_radius,
                                                  roughness=roughness, zoom_factor=zoom_factor)
        masks[:,:,i] = mask
    mean_mask = np.mean(masks, axis=2)
    final_mask = (mean_mask > 0.5).astype(int)
    return final_mask


def cell_magic_wand_3d(image_3d, center, min_radius, max_radius,
                       roughness=2, zoom_factor=1, center_range=2, z_step=1):
    '''Robust cell magic wand tool for 3D images with dimensions (z, x, y) - default for tifffile.load.
    This functions runs the robust wand tool on each z slice in the image and returns the mean mask
    thresholded to 0.5'''
    masks = np.zeros((int(image_3d.shape[0]/z_step), image_3d.shape[1], image_3d.shape[2]))
    for s in range(int(image_3d.shape[0]/z_step)):
        mask = cell_magic_wand(image_3d[s*z_step,:,:], center, min_radius, max_radius,
                               roughness=roughness, zoom_factor=zoom_factor,
                               center_range=center_range)
        masks[s,:,:] = mask
    mean_mask = np.mean(masks, axis=0)
    final_mask = (mean_mask > 0.5).astype(int)
    return final_mask