mitkSegmentAnythingPythonService.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 mitkSegmentAnythingPythonService_h
#define mitkSegmentAnythingPythonService_h
 
#include <mitkSegmentAnythingProcessExecutor.h>
#include <MitkSegmentationExports.h>
#include <thread>
#include <future>
#include <mitkImage.h>
#include <mitkLabelSetImage.h>
#include <itkImage.h>
#include <mitkCommon.h>
 
namespace mitk
{
  class MITKSEGMENTATION_EXPORT SegmentAnythingPythonService : public itk::Object
  {
  public: 
    enum Status
    {
      READY,
      OFF,
      KILLED,
      CUDAError
    };
    mitkClassMacroItkParent(SegmentAnythingPythonService, itk::Object);
 
    SegmentAnythingPythonService(std::string workingDir, std::string modelType,
                                 std::string checkPointPath, unsigned int gpuId, std::string backend);
    
    ~SegmentAnythingPythonService();
    
    itkSetMacro(MitkTempDir, std::string);
    itkGetConstMacro(MitkTempDir, std::string);
    mitkNewMacro5Param(SegmentAnythingPythonService, std::string, std::string, std::string, unsigned int, std::string);
    static void onPythonProcessEvent(itk::Object*, const itk::EventObject&, void*);
 
    static bool CheckStatus() /*throw(mitk::Exception)*/;
 
    void StartAsyncProcess();
 
    void StopAsyncProcess();
 
    void TransferImageToProcess(const Image *inputAtTimeStep, std::string &UId);
 
    void TransferPointsToProcess(const std::string &triggerCSV) const;
 
    MultiLabelSegmentation::Pointer RetrieveImageFromProcess(long timeOut= -1) const;
 
    static Status CurrentStatus;
 
  private:
    void start_python_daemon() const;
 
    void WriteControlFile(const std::string &statusString) const;
 
    void CreateTempDirs(const std::string &dirPattern);
 
    template <typename TPixel, unsigned int VImageDimension>
    void ITKWriter(const itk::Image<TPixel, VImageDimension> *image, std::string& outputFilename) const;
 
 
    std::string m_MitkTempDir;
    std::string m_PythonPath;
    std::string m_ModelType;
    std::string m_CheckpointPath;
    std::string m_InDir, m_OutDir;
    std::string m_Backend;
    std::string m_CurrentUId;
    int m_GpuId = 0;
    const std::string PARENT_TEMP_DIR_PATTERN = "mitk-sam-XXXXXX";
    const std::string TRIGGER_FILENAME = "trigger.csv";
    const std::string SAM_PYTHON_FILE_NAME = "run_inference_daemon.py";
    std::future<void> m_Future;
    SegmentAnythingProcessExecutor::Pointer m_DaemonExec;
  };
 
  struct SIGNALCONSTANTS
  {
    static const std::string READY;
    static const std::string KILL;
    static const std::string OFF;
    static const std::string CUDA_OUT_OF_MEMORY_ERROR;
    static const std::string TIMEOUT_ERROR;
  };
 
} // namespace
 
#endif