Medical Imaging Interaction Toolkit  2024.06.00
Medical Imaging Interaction Toolkit
mitknnUnetTool.h
Go to the documentation of this file.
1 /*============================================================================
2 
3 The Medical Imaging Interaction Toolkit (MITK)
4 
5 Copyright (c) German Cancer Research Center (DKFZ)
6 All rights reserved.
7 
8 Use of this source code is governed by a 3-clause BSD license that can be
9 found in the LICENSE file.
10 
11 ============================================================================*/
12 
13 #ifndef mitknnUnetTool_h
14 #define mitknnUnetTool_h
15 
16 #include "mitkSegWithPreviewTool.h"
17 #include "mitkCommon.h"
18 #include "mitkToolManager.h"
21 #include <numeric>
22 #include <utility>
23 
24 namespace us
25 {
26  class ModuleResource;
27 }
28 
29 namespace mitk
30 {
35  struct ModelParams
36  {
37  std::string task;
38  std::vector<std::string> folds;
39  std::string model;
40  std::string trainer;
41  std::string planId;
42  std::string outputDir;
43  std::string inputName;
44  std::string timeStamp;
45 
46  size_t generateHash() const
47  {
48  std::string toHash;
49  std::string foldsConcatenated = std::accumulate(folds.begin(), folds.end(), std::string(""));
50  toHash += this->task;
51  toHash += this->model;
52  toHash += this->inputName;
53  toHash += foldsConcatenated;
54  toHash += this->timeStamp;
55  size_t hashVal = std::hash<std::string>{}(toHash);
56  return hashVal;
57  }
58  };
59 
69  {
70  public:
72  itkFactorylessNewMacro(Self);
73  itkCloneMacro(Self);
74 
75  const char **GetXPM() const override;
76  const char *GetName() const override;
77  us::ModuleResource GetIconResource() const override;
78 
79  void Activated() override;
80 
81  itkSetMacro(nnUNetDirectory, std::string);
82  itkGetConstMacro(nnUNetDirectory, std::string);
83 
84  itkSetMacro(ModelDirectory, std::string);
85  itkGetConstMacro(ModelDirectory, std::string);
86 
87  itkSetMacro(PythonPath, std::string);
88  itkGetConstMacro(PythonPath, std::string);
89 
90  itkSetMacro(MitkTempDir, std::string);
91  itkGetConstMacro(MitkTempDir, std::string);
92 
93  itkSetMacro(PostProcessingJsonDirectory, std::string);
94  itkGetConstMacro(PostProcessingJsonDirectory, std::string);
95 
96  itkSetMacro(MixedPrecision, bool);
97  itkGetConstMacro(MixedPrecision, bool);
98  itkBooleanMacro(MixedPrecision);
99 
100  itkSetMacro(Mirror, bool);
101  itkGetConstMacro(Mirror, bool);
102  itkBooleanMacro(Mirror);
103 
104  itkSetMacro(MultiModal, bool);
105  itkGetConstMacro(MultiModal, bool);
106  itkBooleanMacro(MultiModal);
107 
108  itkSetMacro(NoPip, bool);
109  itkGetConstMacro(NoPip, bool);
110  itkBooleanMacro(NoPip);
111 
112  itkSetMacro(Ensemble, bool);
113  itkGetConstMacro(Ensemble, bool);
114  itkBooleanMacro(Ensemble);
115 
116  itkSetMacro(Predict, bool);
117  itkGetConstMacro(Predict, bool);
118  itkBooleanMacro(Predict);
119 
120  itkSetMacro(GpuId, unsigned int);
121  itkGetConstMacro(GpuId, unsigned int);
122 
127  std::vector<ModelParams> m_ParamQ;
128 
133  std::vector<mitk::Image::ConstPointer> m_OtherModalPaths;
134 
136 
141  void RenderOutputBuffer();
142 
148  LabelSetImage::Pointer GetOutputBuffer();
149 
154  void ClearOutputBuffer();
155 
159  mitk::DataStorage *GetDataStorage();
160 
161  mitk::DataNode *GetRefNode();
162 
163  void SetOutputBuffer(LabelSetImage::Pointer);
164 
165  protected:
170  nnUNetTool() = default;
171 
176  ~nnUNetTool();
177 
193  void DoUpdatePreview(const Image* inputAtTimeStep, const Image* oldSegAtTimeStep, LabelSetImage* previewImage, TimeStepType timeStep) override;
194  void UpdatePrepare() override;
195 
196  private:
197  std::string m_MitkTempDir;
198  std::string m_nnUNetDirectory;
199  std::string m_ModelDirectory;
200  std::string m_PythonPath;
201  std::string m_PostProcessingJsonDirectory;
202  // bool m_UseGPU; kept for future
203  // bool m_AllInGPU;
204  bool m_MixedPrecision;
205  bool m_Mirror;
206  bool m_NoPip;
207  bool m_MultiModal;
208  bool m_Ensemble = false;
209  bool m_Predict;
210  LabelSetImage::Pointer m_OutputBuffer;
211  unsigned int m_GpuId;
212  const std::string m_TEMPLATE_FILENAME = "XXXXXX_000_0000.nii.gz";
213  };
214 } // namespace mitk
215 #endif
mitk::SegWithPreviewTool
Base class for any auto segmentation tool that provides a preview of the new segmentation.
Definition: mitkSegWithPreviewTool.h:38
mitk::ModelParams::generateHash
size_t generateHash() const
Definition: mitknnUnetTool.h:46
mitk::ModelParams::inputName
std::string inputName
Definition: mitknnUnetTool.h:43
mitk::ModelParams
nnUNet parameter request object holding all model parameters for input. Also holds output temporary d...
Definition: mitknnUnetTool.h:35
mitk::nnUNetTool::m_InputBuffer
mitk::Image::ConstPointer m_InputBuffer
Definition: mitknnUnetTool.h:135
us
Definition: mitkAbstractFileReader.h:29
mitk::ModelParams::planId
std::string planId
Definition: mitknnUnetTool.h:41
mitkSegWithPreviewTool.h
mitk::Image
Image class for storing images.
Definition: mitkImage.h:69
mitk::ModelParams::outputDir
std::string outputDir
Definition: mitknnUnetTool.h:42
itk::SmartPointer< const Self >
mitkStandardFileLocations.h
mitk::ModelParams::trainer
std::string trainer
Definition: mitknnUnetTool.h:40
mitk::ModelParams::task
std::string task
Definition: mitknnUnetTool.h:37
mitk::nnUNetTool
nnUNet segmentation tool.
Definition: mitknnUnetTool.h:68
mitk
Find image slices visible on a given plane.
Definition: RenderingTests.dox:1
mitk::EventStateMachine
‍**
Definition: mitkEventStateMachine.h:111
mitk::nnUNetTool::m_ParamQ
std::vector< ModelParams > m_ParamQ
vector of ModelParams. Size > 1 only for ensemble prediction.
Definition: mitknnUnetTool.h:121
mitk::ModelParams::model
std::string model
Definition: mitknnUnetTool.h:39
mitkCommon.h
MITKSEGMENTATION_EXPORT
#define MITKSEGMENTATION_EXPORT
Definition: MitkSegmentationExports.h:15
mitk::DataStorage
Data management class that handles 'was created by' relations.
Definition: mitkDataStorage.h:43
us::ModuleResource
Definition: usModuleResource.h:55
mitk::ModelParams::folds
std::vector< std::string > folds
Definition: mitknnUnetTool.h:38
mitk::nnUNetTool::m_OtherModalPaths
std::vector< mitk::Image::ConstPointer > m_OtherModalPaths
Holds paths to other input image modalities.
Definition: mitknnUnetTool.h:133
mitk::TimeStepType
std::size_t TimeStepType
Definition: mitkTimeGeometry.h:27
mitkClassMacro
#define mitkClassMacro(className, SuperClassName)
Definition: mitkCommon.h:36
mitk::ModelParams::timeStamp
std::string timeStamp
Definition: mitknnUnetTool.h:44
MitkSegmentationExports.h
mitk::DataNode
Class for nodes of the DataTree.
Definition: mitkDataNode.h:63
mitkToolManager.h
mitk::LabelSetImage
LabelSetImage class for handling labels and layers in a segmentation session.
Definition: mitkLabelSetImage.h:43