#==============================================================================
#
# Project:   SharpImage
# Module:    Watershed.py
# Language:  IronPython
# Author:    Dan Mueller
# Date:      $Date: 2007-07-06 10:57:00 +1000 (Fri, 06 Jul 2007) $
# Revision:  $Revision: 2 $
#
# Copyright (c) Queensland University of Technology (QUT) 2007.
# All rights reserved.
#
# This software is distributed WITHOUT ANY WARRANTY; without even
# the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE.  See the above copyright notices for more information.
#
#==============================================================================

# Import the base script class
import ImageToImageScript
from ImageToImageScript import *

# Add references and import
clr.AddReference("ManagedITK.CastFilters")
clr.AddReference("ManagedITK.MorphologicalFilters")
clr.AddReference("ManagedITK.ThresholdFitlers")
from itk import *

class WatershedScript(ImageToImageScriptObject):
    # -------------------------------------------------------------------------
    Name = "Watershed"
    Help = """Applies the watershed segmentation algorithm, producing a heirarchy
              of labels. The input height function (typically a smoothed GradientMagnitude
              image) is segmented into regions or catchment basins. The output label image
              describes each segmented region with a different label value."""
    Parameters = """(double) Threshold = the minimum percentage height for processing. (0.05)
                    (double) Level = the depth of metaphorical flooding, between 0.0 and 1.0. (0.2)
                    (bool) OverlayResult = the resultant output is overlaid on the input. (False)"""
    LabelImage = None
    Threshold = 0.05
    Level = 0.2
    OverlayResult = False
    # -------------------------------------------------------------------------

    def Initialise(self):
        """ Initialise the environment for running this script. """
        """ NOTE: This function is invoked on the main UI thread. """
        try:
            ThreadedScriptObject.Initialise( self )
            self.Renderer = self.ParentApplication.CurrentRenderer
            self.Input = self.Renderer.Inputs[0]
            self.LabelImage = itkImage.New( itkPixelType.UL, self.Input.Dimension )
            self.Output = itkImage.New( itkPixelType.UC, self.Input.Dimension )
        except Exception, ex:
            self.HandleException( ex )
            self.FinishedWork( False )
            ScriptObject.Finalise( self )

    def ThreadedDoWork(self):
        """ Perform the main functions of the script on a background thread. """
        try:
            self.StartedWork()
            self.WriteInputName()
            
            # Apply watershed
            filterWatershed = itkWatershedImageFilter.New( self.Input )
            self.AddEventHandlersToProcessObject( filterWatershed )
            filterWatershed.SetInput( self.Input )
            filterWatershed.Threshold = self.Threshold
            filterWatershed.Level = self.Level
            self.WriteLineToConsole( "Threshold=" + self.Threshold.ToString() )
            self.WriteLineToConsole( "Level=" + self.Level.ToString() )
            filterWatershed.UpdateLargestPossibleRegion()

            # Relabel
            self.WriteLineToConsole( "WARNING: Relabeling UL as UC: Potiential loss of labels!" )
            filterRelabel = itkRelabelComponentImageFilter.New( self.LabelImage, self.Output )
            self.AddEventHandlersToProcessObject( filterRelabel )
            filterRelabel.SetInput( filterWatershed.GetOutput() )
            filterRelabel.UpdateLargestPossibleRegion()
            filterRelabel.GetOutput( self.Output )
            self.DisconnectInputAndOutput()

            # Clean up
            self.DisposeOfObject( filterWatershed )
            self.DisposeOfObject( filterRelabel )

            # Finish
            self.WriteOutputName()
            self.FinishedWork( True )

        except Exception, ex:
            self.HandleException( ex )
            self.FinishedWork( False )
            self.Finalise( )
    
    def Finalise(self):
        """ Finalise the environment after running this script. """
        """ NOTE: This function is invoked on the main UI thread. """
        ThreadedScriptObject.Finalise(self)

        # Setup LUT
        lut = siGdiLookupTable( )
        lut.SetTableRange( self.Output, 256 )

        if (self.OverlayResult):
            # Add to current Renderer
            lut.SetTableToPsuedoRandomColor( 0.4 )
            self.Renderer.AddInputAsLabel( self.Output, lut )
            self.Renderer.InitialiseInputs()
            self.Renderer.Repaint()
            self.Renderer.Focus()
        else:
            # Display in a new Renderer
            newRenderer = siGdiSliceRenderer( self.ParentApplication )    
            lut.SetTableToPsuedoRandomColor( 1.0 )
            newRenderer.AddInputAsLabel( self.Output, lut )
            newRenderer.Initialise()
            self.ParentApplication.ShowRenderer( newRenderer )