mitkSegmentAnythingTool.h

Go to the documentation of this file.

/*============================================================================
 
The Medical Imaging Interaction Toolkit (MITK)
 
Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.
 
Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.
 
============================================================================*/
 
#ifndef mitkSegmentAnythingTool_h
#define mitkSegmentAnythingTool_h
 
#include "mitkSegWithPreviewTool.h"
#include "mitkPointSet.h"
#include "mitkProcessExecutor.h"
#include "mitkSegmentAnythingPythonService.h"
#include <MitkSegmentationExports.h>
#include <itkImage.h>
#include <mitkLevelWindow.h>
 
namespace us
{
  class ModuleResource;
}
 
namespace mitk
{
  class MITKSEGMENTATION_EXPORT SegmentAnythingTool : public SegWithPreviewTool
  {
  public:
    mitkClassMacro(SegmentAnythingTool, SegWithPreviewTool);
    itkFactorylessNewMacro(Self);
    itkCloneMacro(Self);
 
    const char **GetXPM() const override;
    const char *GetName() const override;
    us::ModuleResource GetIconResource() const override;
 
    void Activated() override;
    void Deactivated() override;
 
    virtual void ClearPicks();
 
    virtual bool HasPicks() const;
 
    itkSetMacro(MitkTempDir, std::string);
    itkGetConstMacro(MitkTempDir, std::string);
 
    itkSetMacro(PythonPath, std::string);
    itkGetConstMacro(PythonPath, std::string);
 
    itkSetMacro(ModelType, std::string);
    itkGetConstMacro(ModelType, std::string);
 
    itkSetMacro(CheckpointPath, std::string);
    itkGetConstMacro(CheckpointPath, std::string);
 
    itkSetMacro(Backend, std::string);
    itkGetConstMacro(Backend, std::string);
 
    itkSetMacro(GpuId, int);
    itkGetConstMacro(GpuId, int);
    
    itkSetMacro(TimeOutLimit, long);
    itkGetConstMacro(TimeOutLimit, long);
 
    itkSetMacro(IsReady, bool);
    itkGetConstMacro(IsReady, bool);
    itkBooleanMacro(IsReady);
 
    void InitSAMPythonProcess();
 
    bool IsPythonReady() const;
 
    Message1<const std::string&> SAMStatusMessageEvent;
 
  protected:
    SegmentAnythingTool();
    ~SegmentAnythingTool() = default;
 
    void ConnectActionsAndFunctions() override;
 
    /*
     * @brief Add positive seed point action of StateMachine pattern
     */
    virtual void OnAddPositivePoint(StateMachineAction*, InteractionEvent *interactionEvent);
    
    /*
     * @brief Add negative seed point action of StateMachine pattern
     */
    virtual void OnAddNegativePoint(StateMachineAction*, InteractionEvent *interactionEvent);
 
    /*
     * @brief Delete action of StateMachine pattern. The function deletes positive or negative points in 
       the reverse order of creation. This is done by finding & deleting the Point having the highest 
       PointIdentifier value from either of the PointSets m_PointSetPositive & m_PointSetNegative.
     */
    virtual void OnDelete(StateMachineAction*, InteractionEvent*);
 
    /*
     * @brief Clear all seed points and call UpdatePreview to reset the segmentation Preview
     */
    void ClearSeeds();
 
    void DoUpdatePreview(const Image *inputAtTimeStep, const Image *oldSegAtTimeStep, LabelSetImage *previewImage, TimeStepType timeStep) override;
 
    std::vector<std::pair<mitk::Point2D, std::string>> GetPointsAsVector(const mitk::BaseGeometry *baseGeometry) const;
 
    virtual std::string GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry) const;
 
    std::string GetHashForCurrentPlane(const mitk::LevelWindow &levelWindow) const;
 
    void EmitSAMStatusMessageEvent(const std::string &status);
 
    void ConfirmCleanUp() override;
 
    template <typename TPixel, unsigned int VImageDimension>
    void ITKWindowing(const itk::Image<TPixel, VImageDimension>*, mitk::Image*, ScalarType, ScalarType);
 
    static mitk::Point2D Get2DIndicesfrom3DWorld(const mitk::BaseGeometry *baseGeometry, const mitk::Point3D &point3d);
 
    std::unique_ptr<SegmentAnythingPythonService> m_PythonService;
 
  private:
    std::string m_MitkTempDir;
    std::string m_PythonPath;
    std::string m_ModelType;
    std::string m_CheckpointPath;
    std::string m_Backend;
    int m_GpuId = 0;
    PointSet::Pointer m_PointSetPositive;
    PointSet::Pointer m_PointSetNegative;
    DataNode::Pointer m_PointSetNodePositive;
    DataNode::Pointer m_PointSetNodeNegative;
    bool m_IsGenerateEmbeddings = true;
    bool m_IsReady = false;
    int m_PointSetCount = 0;
    long m_TimeOutLimit = -1;
    const Label::PixelType MASK_VALUE = 1;
  };
} // namespace
 
#endif