mitknnUnetTool.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 mitknnUnetTool_h
#define mitknnUnetTool_h
 
#include "mitkSegWithPreviewTool.h"
#include "mitkCommon.h"
#include "mitkToolManager.h"
#include <MitkSegmentationExports.h>
#include <mitkStandardFileLocations.h>
#include <numeric>
#include <utility>
 
namespace us
{
  class ModuleResource;
}
 
namespace mitk
{
  struct ModelParams
  {
    std::string task;
    std::vector<std::string> folds;
    std::string model;
    std::string trainer;
    std::string planId;
    std::string outputDir;
    std::string inputName;
    std::string timeStamp;
 
    size_t generateHash() const
    {
      std::string toHash;
      std::string foldsConcatenated = std::accumulate(folds.begin(), folds.end(), std::string(""));
      toHash += this->task;
      toHash += this->model;
      toHash += this->inputName;
      toHash += foldsConcatenated;
      toHash += this->timeStamp;
      size_t hashVal = std::hash<std::string>{}(toHash);
      return hashVal;
    }
  };
 
  class MITKSEGMENTATION_EXPORT nnUNetTool : public SegWithPreviewTool
  {
  public:
    mitkClassMacro(nnUNetTool, SegWithPreviewTool);
    itkFactorylessNewMacro(Self);
    itkCloneMacro(Self);
 
    const char **GetXPM() const override;
    const char *GetName() const override;
    us::ModuleResource GetIconResource() const override;
 
    void Activated() override;
 
    itkSetMacro(nnUNetDirectory, std::string);
    itkGetConstMacro(nnUNetDirectory, std::string);
 
    itkSetMacro(ModelDirectory, std::string);
    itkGetConstMacro(ModelDirectory, std::string);
 
    itkSetMacro(PythonPath, std::string);
    itkGetConstMacro(PythonPath, std::string);
 
    itkSetMacro(MitkTempDir, std::string);
    itkGetConstMacro(MitkTempDir, std::string);
 
    itkSetMacro(PostProcessingJsonDirectory, std::string);
    itkGetConstMacro(PostProcessingJsonDirectory, std::string);
 
    itkSetMacro(MixedPrecision, bool);
    itkGetConstMacro(MixedPrecision, bool);
    itkBooleanMacro(MixedPrecision);
 
    itkSetMacro(Mirror, bool);
    itkGetConstMacro(Mirror, bool);
    itkBooleanMacro(Mirror);
 
    itkSetMacro(MultiModal, bool);
    itkGetConstMacro(MultiModal, bool);
    itkBooleanMacro(MultiModal);
 
    itkSetMacro(NoPip, bool);
    itkGetConstMacro(NoPip, bool);
    itkBooleanMacro(NoPip);
 
    itkSetMacro(Ensemble, bool);
    itkGetConstMacro(Ensemble, bool);
    itkBooleanMacro(Ensemble);
 
    itkSetMacro(Predict, bool);
    itkGetConstMacro(Predict, bool);
    itkBooleanMacro(Predict);
 
    itkSetMacro(GpuId, unsigned int);
    itkGetConstMacro(GpuId, unsigned int);
 
    std::vector<ModelParams> m_ParamQ;
 
    std::vector<mitk::Image::ConstPointer> m_OtherModalPaths;
 
    mitk::Image::ConstPointer m_InputBuffer;
 
    void RenderOutputBuffer();
 
    MultiLabelSegmentation::Pointer GetOutputBuffer();
 
    void ClearOutputBuffer();
 
    mitk::DataStorage *GetDataStorage();
 
    mitk::DataNode *GetRefNode();
 
    void SetOutputBuffer(MultiLabelSegmentation::Pointer);
 
  protected:
    nnUNetTool() = default;
 
    ~nnUNetTool();
 
    void DoUpdatePreview(const Image* inputAtTimeStep, const Image* oldSegAtTimeStep, MultiLabelSegmentation* previewImage, TimeStepType timeStep) override;
    void UpdatePrepare() override;
 
  private:
    std::string m_MitkTempDir;
    std::string m_nnUNetDirectory;
    std::string m_ModelDirectory;
    std::string m_PythonPath;
    std::string m_PostProcessingJsonDirectory;
    // bool m_UseGPU; kept for future
    // bool m_AllInGPU;
    bool m_MixedPrecision;
    bool m_Mirror;
    bool m_NoPip;
    bool m_MultiModal;
    bool m_Ensemble = false;
    bool m_Predict;
    MultiLabelSegmentation::Pointer m_OutputBuffer;
    unsigned int m_GpuId;
    const std::string m_TEMPLATE_FILENAME = "XXXXXX_000_0000.nii.gz";
  };
} // namespace mitk
#endif