<tool id="mesmer" name="Mesmer" version="@TOOL_VERSION@+galaxy@VERSION_SUFFIX@" profile="19.01">
    <description>Mesmer for whole-cell segmentation of multiplexed tissue imaging data</description>

    <macros>
        <import>macros.xml</import>
    </macros>

    <expand macro="requirements"/>
    <expand macro="stdio"/>

    <command detect_errors="exit_code"><![CDATA[
    python $script
    ]]></command>
    <configfiles>
        <configfile name="script">
import argparse
import os 
import sys

import numpy as np
import tifffile
import zarr

from deepcell.applications import Mesmer

level = 0
is_ome=False
#if $image.file_ext == "ome.tiff":
is_ome=True
#end if

## Grab params
nuc_kwargs = {}
wc_kwargs = {}


with tifffile.TiffFile("$image", is_ome=is_ome) as tiff:

    # Read single pyramid level
    level_array = zarr.open(tiff.aszarr(series = 0, level = level))

    ## grab the nuclear and membrane channels based on their indices add the markers along the channel axis
    ## Tifffile should always read as (channel,X,Y)
    #if $compartment_select.compartment != 'whole-cell':
    nuc_kwargs = {
        #if $compartment_select.nuclear_options.pixel_expansion != '':
    'pixel_expansion': $compartment_select.nuclear_options.pixel_expansion,
        #end if
        #for $key, $value in $compartment_select.nuclear_options.adv_options.items():
    '$key': $value,
        #end for
    'maxima_threshold': $compartment_select.nuclear_options.maxima_threshold
    }
    #end if

    nuclear_indices = [int(x) for x in $nuclear_channels.split(',')]
    nuclear_channels = level_array.oindex[nuclear_indices, :, :]
    nuclear_channels = np.sum(nuclear_channels, axis = 0)

    #if $compartment_select.compartment != 'nuclear':
    wc_kwargs = {
        #for $key, $value in $compartment_select.wc_options.adv_options.items():
    '$key': $value,
        #end for
    'maxima_threshold': $compartment_select.wc_options.maxima_threshold
    }

    membrane_indices = [int(x) for x in $compartment_select.wc_options.wc_channels.split(',')]
    membrane_channels = level_array.oindex[membrane_indices, :, :]
    membrane_channels = np.sum(membrane_channels, axis = 0)  
    #end if

    ## stack the nuclear and membrane composite channels with nuclear in channel 0 and mem in channel 1
    ## mesmer expects dimensions to be (X,Y,Channel) so axis = -1
    #if $compartment_select.compartment == 'nuclear':
    membrane_channels = np.zeros(nuclear_channels.shape)
    #end if
    formatted_image = np.stack((nuclear_channels,membrane_channels), axis=-1)

    ## add batch dimension. Will have to be squeezed out later
    formatted_image = np.expand_dims(formatted_image, 0)

## Create the application
app = Mesmer()

## Run segmentation
mask = app.predict(
    formatted_image, 
    image_mpp = $image_mpp, 
    compartment = "$compartment_select.compartment",
    pad_mode = 'constant',
    postprocess_kwargs_whole_cell = wc_kwargs, 
    postprocess_kwargs_nuclear = nuc_kwargs)

#if $squeeze:
mask = np.squeeze(mask)
#end if

#if $compartment_select.compartment == 'both':
## split the two-channel mask into separate outputs
    #if $squeeze:
tifffile.imsave( "WC_output_mask.tif", mask[:,:,0])
tifffile.imsave("NU_output_mask.tif", mask[:,:,1])
    #else:
tifffile.imsave( "WC_output_mask.tif", mask[:,:,:,0])
tifffile.imsave("NU_output_mask.tif", mask[:,:,:,1])
    #end if
#else:
## save single-channel mask outputs as a tiff
tifffile.imsave("mask.tif", mask)
#end if
        </configfile>
    </configfiles>
    <inputs>
        <param name="image" type="data" format="tiff, ome.tiff" label="Image containing the all marker(s)"/>
        <param name="nuclear_channels" type="text" value="0" label="The numerical indices of the channel(s) for the nuclear markers" help="No quotes, separated by comma. e.g. 0, 1."/>
        <param name="image_mpp" type="float" value="0.5" label="Resolution of the image in microns-per-pixel"/>
        <param name="squeeze" type="boolean" checked="true" label="Whether to np.squeeze the outputs before saving"/>
        <conditional name="compartment_select">
            <param name="compartment" type="select"  label="Compartment for segmentation prediction: ">
                <option selected="true" value="whole-cell">Whole cell</option>
                <option value="nuclear">Nuclear</option>
                <option value="both">Both</option>
            </param>
            <when value="nuclear">
                <expand macro="nuclear_options_macro"/>
            </when>
            <when value="whole-cell">
                <expand macro="wc_options_macro"/>
            </when>
            <when value="both">
                <expand macro="nuclear_options_macro"/>
                <expand macro="wc_options_macro"/>
            </when>
        </conditional>
    </inputs>

    <outputs>
        <data format="tiff" name="nu_mask" from_work_dir="NU_output_mask.tif" label="${tool.name} on ${on_string}: Nuclear Mask">
            <filter>compartment_select['compartment'] == 'both'</filter>
        </data>
        <data format="tiff" name="wc_mask" from_work_dir="WC_output_mask.tif" label="${tool.name} on ${on_string}: Whole Cell Mask">
            <filter>compartment_select['compartment'] == 'both'</filter>
        </data>
        <data format="tiff" name="mask" from_work_dir="mask.tif" label="${tool.name} on ${on_string}: Mask">
            <filter>compartment_select['compartment'] != 'both'</filter>
        </data>
    </outputs>
    <tests>
         <test expect_num_outputs="1">
            <param name="image" value="deepcell_test.tiff" ftype="tiff"/>
            <param name="compartment" value="nuclear" />
            <param name="image_mpp" value="0.65" />
            <param name="squeeze" value="True" />
            <output name="mask" ftype="tiff">
                <assert_contents>
                    <has_size value="1049000" delta="1000" />
                </assert_contents>
            </output>
        </test>
         <test expect_num_outputs="1">
            <param name="image" value="deepcell_test.tiff" ftype="tiff"/>
            <param name="compartment" value="whole-cell" />
            <param name="wc_channels" value="0, 1" />
            <param name="maxima_threshold" value="0.075" />
            <param name="image_mpp" value="0.65" />
            <param name="squeeze" value="True" />
            <output name="mask" ftype="tiff">
                <assert_contents>
                    <has_size value="1049000" delta="1000" />
                </assert_contents>
            </output>
        </test>
         <test expect_num_outputs="2">
            <param name="image" value="deepcell_test.tiff" ftype="tiff"/>
            <param name="compartment" value="both" />
            <param name="image_mpp" value="0.65" />
            <param name="wc_channels" value="1" />
            <param name="nuclear_channels" value="0" />
            <param name="squeeze" value="True" />
            <output name="wc_mask" ftype="tiff">
                <assert_contents>
                    <has_size value="1049000" delta="1000" />
                </assert_contents>
            </output>
            <output name="nu_mask" ftype="tiff">
                <assert_contents>
                    <has_size value="1049000" delta="1000" />
                </assert_contents>
            </output>
        </test>
 </tests>
    <help><![CDATA[
------
Mesmer
------

**Mesmer** is a DeepCell application that uses a pre-trained model for predicting nuclear and whole cell segmentation from multiplex tissue data.

**Inputs**
1. a multiple tissue image that includes a nuclear marker and optionally cytoplasm/membrane markers
        
**Outputs**
1. the nuclear or whole cell segmentation mask 
    ]]></help>
    <expand macro="citations" />
</tool>