mitknnInteractiveTool.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 mitknnInteractiveTool_h
#define mitknnInteractiveTool_h
 
#include <mitkSegWithPreviewTool.h>
#include <mitknnInteractiveEnums.h>
 
#include <optional>
#include <string>
#include <utility>
#include <vector>
 
namespace mitk::nnInteractive
{
  class Interactor;
}
 
namespace mitk
{
  class PythonContext;
 
  class MITKPYTHONSEGMENTATION_EXPORT nnInteractiveTool : public SegWithPreviewTool
  {
  public:
    struct CUDADeviceInfo
    {
      std::string Name;     
      int Major;            
      int Minor;            
      int TotalMemoryMB;    
    };
 
    struct SupportedInteractions
    {
      bool Point = true;    
      bool Box = true;      
      bool Scribble = true; 
      bool Lasso = true;    
      bool Mask = true;     
    };
 
    mitkClassMacro(nnInteractiveTool, SegWithPreviewTool)
    itkFactorylessNewMacro(Self)
 
    
    const char* GetName() const override;
 
    us::ModuleResource GetIconResource() const override;
 
    void Deactivated() override;
 
    using InteractorMap = std::unordered_map<nnInteractive::InteractionType, std::unique_ptr<nnInteractive::Interactor>>;
 
    const InteractorMap& GetInteractors() const;
 
    const nnInteractive::Interactor* GetInteractor(nnInteractive::InteractionType interactionType) const;
 
    void EnableInteractor(nnInteractive::InteractionType nextInteractionType, nnInteractive::PromptType promptType);
 
    void DisableInteractor(std::optional<nnInteractive::InteractionType> interactionType = std::nullopt);
 
    void ResetInteractions();
 
    bool HasInteractions() const;
 
    void UndoLastInteraction();
 
    bool CanUndo() const;
 
    bool SupportsUndo() const;
 
    bool GetAutoZoom() const;
 
    void SetAutoZoom(bool autoZoom);
 
    bool GetAutoRefine() const;
 
    void SetAutoRefine(bool autoRefine);
 
    std::optional<nnInteractive::Backend> GetBackend() const;
 
    std::string GetVirtualEnvName() const;
 
    bool CreatePythonContext();
 
    PythonContext* GetPythonContext() const;
 
    bool IsInstalled() const;
 
    bool IsLocalInferenceAvailable() const;
 
    static const std::vector<std::pair<std::string, std::string>>& GetSessionDefiningPreferences();
 
    bool GetCUDADeviceInfo(CUDADeviceInfo& info) const;
 
    void StartSession();
 
    void EndSession();
 
    void AbortSession();
 
    bool IsSessionRunning() const;
 
    bool IsRemoteSession() const;
 
    void Heartbeat();
 
    int GetHeartbeatIntervalMs() const;
 
    SupportedInteractions GetSupportedInteractions() const;
 
    std::optional<std::string> GetModelLicense() const;
 
    void InitializeSessionWithMask(Image* mask);
 
    Message1<bool> ConfirmCleanUpEvent;
 
    Message<> PreviewUpdatedEvent;
 
    Message<> DeactivatedEvent;
 
    Message<> SessionEndedEvent;
 
    Message<> SessionExpiredEvent;
 
  protected:
    nnInteractiveTool();
 
    ~nnInteractiveTool() override;
 
    void SetToolManager(ToolManager* toolManager) override;
 
    void DoUpdatePreview(const Image* inputAtTimeStep, const Image* oldSegAtTimeStep, MultiLabelSegmentation* previewImage, TimeStepType timeStep) override;
 
    void OnTimePointChanged() override;
 
    void Notify(InteractionEvent* event, bool isHandled) override;
 
    void ConfirmCleanUp() override;
 
    void SetPreviewLabel(MultiLabelSegmentation::LabelValueType value, const Color& color);
 
  private:
    void ConstructLocalSession();
 
    void ConstructRemoteSession();
 
    void BindSessionImageAndTargetBuffer();
 
    bool IsRemoteConnectionError(const std::string& message) const;
 
    bool HandleSessionError(const std::string& errorMessage);
 
    class Impl;
    std::unique_ptr<Impl> m_Impl;
  };
}
 
#endif