QmitknnUNetToolGUI.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 QmitknnUNetToolGUI_h
#define QmitknnUNetToolGUI_h
 
#include "QmitkMultiLabelSegWithPreviewToolGUIBase.h"
#include "QmitknnUNetFolderParser.h"
#include "QmitknnUNetGPU.h"
#include "QmitknnUNetWorker.h"
#include "mitkProcessExecutor.h"
#include "mitknnUnetTool.h"
#include "ui_QmitknnUNetToolGUIControls.h"
#include <MitkSegmentationUIExports.h>
#include <QCache>
#include <QMessageBox>
#include <QSettings>
#include <QThread>
#include <QmitkSingleNodeSelectionWidget.h>
#include <QmitknnUNetEnsembleLayout.h>
#include <boost/functional/hash.hpp>
#include <unordered_map>
 
class nnUNetCache
{
public:
  mitk::LabelSetImage::ConstPointer m_SegCache;
  static size_t GetUniqueHash(std::vector<mitk::ModelParams> &requestQ)
  {
    size_t hashCode = 0;
    for (mitk::ModelParams &request : requestQ)
    {
      boost::hash_combine(hashCode, request.generateHash());
    }
    return hashCode;
  }
};
 
class MITKSEGMENTATIONUI_EXPORT QmitknnUNetToolGUI : public QmitkMultiLabelSegWithPreviewToolGUIBase
{
  Q_OBJECT
 
public:
  mitkClassMacro(QmitknnUNetToolGUI, QmitkMultiLabelSegWithPreviewToolGUIBase);
  itkFactorylessNewMacro(Self);
  itkCloneMacro(Self);
 
  QCache<size_t, nnUNetCache> m_Cache;
 
  std::unordered_map<std::string, mitk::ProcessExecutor::Pointer> m_Processes;
 
protected slots:
 
  void OnPreviewRequested();
 
  void OnDirectoryChanged(const QString&);
 
  void OnModelChanged(const QString&);
 
  void OnTaskChanged(const QString &);
 
  void OnTrainerChanged(const QString&);
 
  void OnCheckBoxChanged(int);
 
  void SegmentationProcessFailed();
 
  void SegmentationResultHandler(mitk::nnUNetTool *, bool forceRender = false);
 
  void OnModalitiesNumberChanged(int);
 
  void OnPythonPathChanged(const QString&);
 
  void OnRefreshPresssed();
 
  void OnClearCachePressed();
 
  void OnDownloadModel();
 
  void OnDownloadWorkerExit(bool, const QString);
 
  void OnStopDownload();
 
signals:
  void Operate(QString, QString, mitk::ProcessExecutor::Pointer, mitk::ProcessExecutor::ArgumentListType);
 
protected:
  QmitknnUNetToolGUI();
  ~QmitknnUNetToolGUI();
 
  void ConnectNewTool(mitk::SegWithPreviewTool* newTool) override;
  void InitializeUI(QBoxLayout* mainLayout) override;
  void EnableWidgets(bool enabled) override;
 
private:
  void SetComboBoxToNone(ctkCheckableComboBox*);
 
  void FillAvailableModelsInfoFromJSON(const QString&);
 
  void ExportAvailableModelsAsJSON(const QString&);
 
  void ClearAllModalities();
 
  void DisplayMultiModalInfoFromJSON(const QString&);
 
  void ClearAllModalLabels();
 
  void DumpJSONfromPickle(const QString&);
 
  QString FetchResultsFolderFromEnv();
 
  unsigned int FetchSelectedGPUFromUI();
 
  void SetGPUInfo();
 
  void AddToCache(size_t&, mitk::LabelSetImage::ConstPointer);
 
  void CheckAllInCheckableComboBox(ctkCheckableComboBox*);
 
  std::pair<QStringList, QStringList> ExtractTrainerPlannerFromString(QStringList);
  void ProcessEnsembleModelsParams(mitk::nnUNetTool::Pointer);
 
  void ProcessModelParams(mitk::nnUNetTool::Pointer);
 
  void ShowEnsembleLayout(bool visible = true);
 
  void ShowErrorMessage(const std::string&, QMessageBox::Icon = QMessageBox::Critical);
 
  void WriteStatusMessage(const QString&);
 
  void WriteErrorMessage(const QString&);
 
  void AutoParsePythonPaths();
 
  bool IsModelExists(const QString&, const QString&, const QString&);
 
  void ClearAllComboBoxes();
 
  void DisableEverything();
 
  bool IsNNUNetInstalled(const QString&);
 
  mitk::ModelParams MapToRequest(
    const QString&, const QString&, const QString&, const QString&, const std::vector<std::string>&);
 
  std::vector<std::string> FetchSelectedFoldsFromUI(ctkCheckableComboBox*);
 
  std::vector<mitk::Image::ConstPointer> FetchMultiModalImagesFromUI();
 
  void UpdateCacheCountOnUI();
 
  Ui_QmitknnUNetToolGUIControls m_Controls;
  QmitkGPULoader m_GpuLoader;
 
  std::vector<QmitkSingleNodeSelectionWidget*> m_Modalities;
  std::vector<QLabel*> m_ModalLabels;
 
  std::vector<std::unique_ptr<QmitknnUNetTaskParamsUITemplate>> m_EnsembleParams;
 
  mitk::NodePredicateBase::Pointer m_MultiModalPredicate;
 
  QString m_PythonPath;
 
  int m_UI_ROWS;
 
  std::shared_ptr<QmitknnUNetFolderParser> m_ParentFolder = nullptr;
 
  const QStringList m_VALID_MODELS = {"2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres", "ensembles"};
 
  const QString m_CACHE_COUNT_BASE_LABEL = "Cached Items: ";
 
  const QString m_MITK_EXPORT_JSON_FILENAME = "mitk_export.json";
 
  const QString m_AVAILABLE_MODELS_JSON_FILENAME = "available_models.json";
 
  const QString m_PICKLE_FILENAME = "plans.pkl";
 
  QSettings m_Settings;
 
  bool m_IsResultsFolderValid = false;
 
  QThread* m_nnUNetThread;
  nnUNetDownloadWorker* m_Worker;
  
  bool m_FirstPreviewComputation = true;
  EnableConfirmSegBtnFunctionType m_SuperclassEnableConfirmSegBtnFnc;
};
#endif