#==============================================================================
#
# Project:   SharpImage
# Module:    VesselPathCost.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.FastMorphologicalFilters")
clr.AddReference("ManagedITK.ThresholdFilters")
clr.AddReference("ManagedITK.PixelMathFilters")
clr.AddReference("ManagedITK.CastFilters")
clr.AddReference("ManagedITK.NoiseFilters")
clr.AddReference("ManagedITK.IntensityFilters")
from itk import *

class VesselPathCostScript(ImageToImageScriptObject):
    # -------------------------------------------------------------------------
    Name = "VesselPathCost"
    Help = """Computes the cost function for minimal path extraction of vessels.
              The cost function is computed by morphological top-hat by opening,
              thresholding, and then smoothing. The result is a floating point image."""
    Parameters = """(int) OpenRadius = The radius of the box structuring element. (8)
                    (itkPixel) Threshold1 = The lower threshold to remove noise. (100)
                    (itkPixel) Threshold2 = The upper threshold to remove artefacts. (800)
                    (bool) Smooth = Indicates if smoothing is applied. (True)
                    (float) TimeStep = Curvature Flow time-step parameter. (0.175)
                    (int) Iterations = Curvature Flow iterations parameter. (2)"""
    OpenRadius = 8
    OpenKernelRadius = None
    OpenKernel = None
    OpenAlgorithm = FastMorphologyEnum.VHGW
    Threshold1 = 100
    Threshold2 = 800
    ThresholdOutside = 0
    Smooth = True
    SmoothTimeStep = 0.175
    SmoothIterations = 2
    TempImage = None
    # -------------------------------------------------------------------------

    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.TempImage = itkImage.New( itkPixelType.F, self.Input.Dimension )
            self.Output = itkImage.New( itkPixelType.F, 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 morphological opening
            filterOpen = itkGrayscaleMorphologicalOpeningImageFilter.New( self.Input, self.Input )
            self.AddEventHandlersToProcessObject( filterOpen )
            filterOpen.SetInput( self.Input )
            if (self.OpenKernelRadius == None):
                self.OpenKernelRadius = itkSize( self.Input.Dimension )
                self.OpenKernelRadius.Fill( self.OpenRadius )
            if (self.OpenKernel == None):
                self.OpenKernel = itkFlatStructuringElement.Box( self.OpenKernelRadius );
            filterOpen.Kernel = self.OpenKernel
            filterOpen.Algorithm = self.OpenAlgorithm
            self.WriteLineToConsole( "OpenKernel=" + self.OpenKernel.ToString() )

            # Apply subtract to compute top-hat
            filterSub = itkSubtractImageFilter.New( self.Input, self.Input, self.Input )
            self.AddEventHandlersToProcessObject( filterSub )
            filterSub.SetInput1( self.Input )
            filterSub.SetInput2( filterOpen.GetOutput() )

            # Apply thresholding
            filterThreshold = itkThresholdImageFilter.New( self.Input )
            self.AddEventHandlersToProcessObject( filterThreshold )
            filterThreshold.SetInput( filterSub.GetOutput() )
            filterThreshold.OutsideValue = itkPixel( self.Input.PixelType, self.ThresholdOutside )
            pixelLower = itkPixel( self.Input.PixelType, self.Threshold1 )
            pixelUpper = itkPixel( self.Input.PixelType, self.Threshold2 )
            self.WriteLineToConsole( "Threshold1=" + pixelLower.ToString() )
            self.WriteLineToConsole( "Threshold2=" + pixelUpper.ToString() )
            filterThreshold.ThresholdOutside( pixelLower, pixelUpper )

            # Cast to float
            filterCast = itkCastImageFilter.New( self.Input, self.Output )
            self.AddEventHandlersToProcessObject( filterCast )
            filterCast.SetInput( filterThreshold.GetOutput() )

            # Smooth
            filterFlow = itkCurvatureFlowImageFilter.New( self.Output, self.Output )
            if ( self.Smooth ):
                self.AddEventHandlersToProcessObject( filterFlow )
                filterFlow.SetInput( filterCast.GetOutput() )
                filterFlow.TimeStep = self.SmoothTimeStep
                filterFlow.NumberOfIterations = self.SmoothIterations
                self.WriteLineToConsole( "SmoothTimeStep=" + self.SmoothTimeStep.ToString() )
                self.WriteLineToConsole( "SmoothIterations=" + self.SmoothIterations.ToString() )
                filterFlow.GetOutput( self.TempImage )
            else:
                filterCast.GetOutput( self.TempImage )

            # Rescale intensity
            filterRescale = itkRescaleIntensityImageFilter.New( self.Output, self.Output )
            self.AddEventHandlersToProcessObject( filterRescale )
            filterRescale.SetInput( self.TempImage )
            filterRescale.OutputMinimum = itkPixel( self.Output.PixelType, 0.0 )
            filterRescale.OutputMaximum = itkPixel( self.Output.PixelType, 1.0 )
            filterRescale.UpdateLargestPossibleRegion()
            filterRescale.GetOutput( self.Output )

            # Clean up
            self.DisconnectInputAndOutput()
            self.DisposeOfObject( filterOpen )
            self.DisposeOfObject( filterSub )
            self.DisposeOfObject( filterThreshold )
            self.DisposeOfObject( filterCast )
            self.DisposeOfObject( filterFlow )
            self.DisposeOfObject( filterRescale )

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

        except Exception, ex:
            self.HandleException( ex )
            self.FinishedWork( False )
            self.Finalise( )