--- a/mesmer.xml	Thu Nov 10 20:38:34 2022 +0000
+++ b/mesmer.xml	Wed Dec 28 19:26:02 2022 +0000
@@ -7,63 +7,189 @@
 
     <expand macro="requirements"/>
     <expand macro="stdio"/>
-    <expand macro="version_cmd"/>
 
     <command detect_errors="exit_code"><![CDATA[
-    ln -s '$nuclear_image' 'input.tif' && 
-    @CMD_BEGIN@ 
-    --output-directory ./
-    --output-name 'mask.tif'
-    --nuclear-image 'input.tif'
-    --nuclear-channel $nuclear_channel
-    --compartment $compartment
-    --image-mpp $image_mpp
-    $squeeze
+    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 $dict.items($compartment_select.nuclear_options.adv_options)
+    '$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 $membrane_select.membrane_segment == "True":
-        --membrane-image '$membrane_select.membrane_image'
-        --membrane-channel '$membrane_select.membrane_channel'
+    #if $compartment_select.compartment != 'nuclear':
+    wc_kwargs = {
+        #for $key, $value in $dict.items($compartment_select.wc_options.adv_options)
+    '$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
-    ]]></command>
+    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="nuclear_image" type="data" format="tiff, ome.tiff" label="Image containing the nuclear marker(s)"/>
-        <param name="nuclear_channel" type="integer" value="0" label="The numerical index of the channel(s) from nuclear-image"/>
-        <param name="compartment" type="select"  label="Compartment for segmentation prediction: ">
-            <option selected="true" value="whole-cell">Whole cell</option>
-            <option value="nuclear">Nuclear</option>
-        </param>
+        <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" truevalue="--squeeze" falsevalue="" checked="false" label="Whether to np.squeeze the outputs before saving"/>
-        <conditional name="membrane_select">
-            <param name="membrane_segment" type="select" label="Segment with Cell Membrane">
-                <option selected="True" value="False">No</option>
-                <option value="True">Yes</option>
+        <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="True">
-                <param name="membrane_image" type="data" format="tiff, ome.tiff" label="The path to an image containing the membrane marker(s)"/>
-                <param name="membrane_channel" type="integer" value="0"  label="The numerical index of the channel(s) from membrane-image"/>
+            <when value="nuclear">
+                <expand macro="nuclear_options_macro"/>
             </when>
-            <when value="False" />
+            <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="mask" from_work_dir="mask.tif" label="${tool.name} on ${on_string}: Mask"/>
+        <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>
-            <param name="nuclear_image" value="test.tiff" />
+         <test>
+            <param name="image" value="deepcell_test.tiff" ftype="tiff"/>
             <param name="compartment" value="nuclear" />
-            <param name="membrane_segment" value="False" />
             <param name="image_mpp" value="0.65" />
-            <param name="squeeze" value="--squeeze" />
+            <param name="squeeze" value="True" />
             <output name="mask" ftype="tiff">
                 <assert_contents>
-                    <has_size value="360000" delta="1000" />
+                    <has_size value="1049000" delta="1000" />
                 </assert_contents>
             </output>
         </test>
-    </tests>
+         <test>
+            <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>
+            <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