diff --git a/Modules/Segmentation/Interactions/mitknnUnetTool.cpp b/Modules/Segmentation/Interactions/mitknnUnetTool.cpp index 07afb75ef2..2302046894 100644 --- a/Modules/Segmentation/Interactions/mitknnUnetTool.cpp +++ b/Modules/Segmentation/Interactions/mitknnUnetTool.cpp @@ -1,345 +1,331 @@ /*============================================================================ 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. ============================================================================*/ #include "mitknnUnetTool.h" #include "mitkIOUtil.h" #include "mitkProcessExecutor.h" #include #include #include #include #include namespace mitk { MITK_TOOL_MACRO(MITKSEGMENTATION_EXPORT, nnUNetTool, "nnUNet tool"); } mitk::nnUNetTool::nnUNetTool() { this->SetMitkTempDir(IOUtil::CreateTemporaryDirectory("mitk-XXXXXX")); } mitk::nnUNetTool::~nnUNetTool() { itksys::SystemTools::RemoveADirectory(this->GetMitkTempDir()); } void mitk::nnUNetTool::Activated() { Superclass::Activated(); } void mitk::nnUNetTool::UpdateCleanUp() { // This overriden method is intentionally left out for setting later upon demand // in the `RenderOutputBuffer` method. } void mitk::nnUNetTool::RenderOutputBuffer() { - MITK_INFO << "inn RenderOutputBuffer"; if (this->m_OutputBuffer != nullptr) { Superclass::SetNodeProperties(this->m_OutputBuffer); try { if (nullptr != this->GetPreviewSegmentationNode()) { this->GetPreviewSegmentationNode()->SetVisibility(!this->GetSelectedLabels().empty()); } if (this->GetSelectedLabels().empty()) { this->ResetPreviewNode(); } } catch (const mitk::Exception &e) { MITK_INFO << e.GetDescription(); } } } void mitk::nnUNetTool::SetNodeProperties(LabelSetImage::Pointer segmentation) { // This overriden method doesn't set node properties. Intentionally left out for setting later upon demand // in the `RenderOutputBuffer` method. this->m_OutputBuffer = segmentation; - MITK_INFO << "inn derived SetNodeProperties " << this->m_OutputBuffer.GetPointer(); } mitk::LabelSetImage::Pointer mitk::nnUNetTool::GetOutputBuffer() { - MITK_INFO << "inn GetOutputBuffer " << this->m_OutputBuffer.GetPointer(); return this->m_OutputBuffer; } void mitk::nnUNetTool::ClearOutputBuffer() { - MITK_INFO << "inn ClearOutputBuffer " << this->m_OutputBuffer.GetPointer(); this->m_OutputBuffer = nullptr; } us::ModuleResource mitk::nnUNetTool::GetIconResource() const { us::Module *module = us::GetModuleContext()->GetModule(); us::ModuleResource resource = module->GetResource("Watershed_48x48.png"); return resource; } const char **mitk::nnUNetTool::GetXPM() const { return nullptr; } const char *mitk::nnUNetTool::GetName() const { return "nnUNet"; } mitk::DataStorage *mitk::nnUNetTool::GetDataStorage() { return this->GetToolManager()->GetDataStorage(); } mitk::DataNode *mitk::nnUNetTool::GetRefNode() { return this->GetToolManager()->GetReferenceData(0); } namespace { void onPythonProcessEvent(itk::Object * /*pCaller*/, const itk::EventObject &e, void *) { std::string testCOUT; std::string testCERR; const auto *pEvent = dynamic_cast(&e); if (pEvent) { testCOUT = testCOUT + pEvent->GetOutput(); MITK_INFO << testCOUT; } const auto *pErrEvent = dynamic_cast(&e); if (pErrEvent) { testCERR = testCERR + pErrEvent->GetOutput(); MITK_ERROR << testCERR; } } } // namespace mitk::LabelSetImage::Pointer mitk::nnUNetTool::ComputeMLPreview(const Image *inputAtTimeStep, TimeStepType timeStep) { - MITK_INFO << "IN ComputeMLPreview " << timeStep; if (m_InputBuffer == inputAtTimeStep) { return m_OutputBuffer; } std::string inDir, outDir, inputImagePath, outputImagePath, scriptPath; ProcessExecutor::Pointer spExec = ProcessExecutor::New(); itk::CStyleCommand::Pointer spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(&onPythonProcessEvent); spExec->AddObserver(ExternalProcessOutputEvent(), spCommand); ProcessExecutor::ArgumentListType args; inDir = IOUtil::CreateTemporaryDirectory("nnunet-in-XXXXXX", this->GetMitkTempDir()); std::ofstream tmpStream; inputImagePath = IOUtil::CreateTemporaryFile(tmpStream, m_TEMPLATE_FILENAME, inDir + IOUtil::GetDirectorySeparator()); tmpStream.close(); std::size_t found = inputImagePath.find_last_of(IOUtil::GetDirectorySeparator()); std::string fileName = inputImagePath.substr(found + 1); std::string token = fileName.substr(0, fileName.find("_")); if (this->GetNoPip()) { scriptPath = this->GetnnUNetDirectory() + IOUtil::GetDirectorySeparator() + "nnunet" + IOUtil::GetDirectorySeparator() + "inference" + IOUtil::GetDirectorySeparator() + "predict_simple.py"; } try { IOUtil::Save(inputAtTimeStep, inputImagePath); if (this->GetMultiModal()) { std::string outModalFile; int len = inDir.length() + token.length() + 1 + 8 + 2 + 7; outModalFile.reserve(len); const std::string fileFormat(".nii.gz"); const std::string fileNamePart("_000_000"); for (size_t i = 0; i < this->m_OtherModalPaths.size(); ++i) { mitk::Image::ConstPointer modalImage = this->m_OtherModalPaths[i]; outModalFile.append(inDir); outModalFile.push_back(IOUtil::GetDirectorySeparator()); outModalFile.append(token); outModalFile.append(fileNamePart); outModalFile.append(std::to_string(i + 1)); outModalFile.append(fileFormat); IOUtil::Save(modalImage.GetPointer(), outModalFile); outModalFile.clear(); } } } catch (const mitk::Exception &e) { /* Can't throw mitk exception to the caller. Refer: T28691 */ MITK_ERROR << e.GetDescription(); return nullptr; } // Code calls external process std::string command = "nnUNet_predict"; if (this->GetNoPip()) { #ifdef _WIN32 command = "python"; #else command = "python3"; #endif } for (ModelParams &modelparam : m_ParamQ) { - MITK_INFO << "IN ComputeMLPreview Segmenting... "; outDir = IOUtil::CreateTemporaryDirectory("nnunet-out-XXXXXX", this->GetMitkTempDir()); outputImagePath = outDir + IOUtil::GetDirectorySeparator() + token + "_000.nii.gz"; modelparam.outputDir = outDir; args.clear(); if (this->GetNoPip()) { args.push_back(scriptPath); } args.push_back("-i"); args.push_back(inDir); args.push_back("-o"); args.push_back(outDir); args.push_back("-t"); args.push_back(modelparam.task); if (modelparam.model.find("cascade") != std::string::npos) { args.push_back("-ctr"); } else { args.push_back("-tr"); } args.push_back(modelparam.trainer); args.push_back("-m"); args.push_back(modelparam.model); args.push_back("-p"); args.push_back(modelparam.planId); if (!modelparam.folds.empty()) { args.push_back("-f"); for (auto fold : modelparam.folds) { args.push_back(fold); } } args.push_back("--num_threads_nifti_save"); args.push_back("1"); // fixing to 1 if (!this->GetMirror()) { args.push_back("--disable_tta"); } if (!this->GetMixedPrecision()) { args.push_back("--disable_mixed_precision"); } if (this->GetEnsemble()) { args.push_back("--save_npz"); } try { - // for (auto arg : args) - // { - // MITK_INFO << arg; - // } std::string resultsFolderEnv = "RESULTS_FOLDER=" + this->GetModelDirectory(); itksys::SystemTools::PutEnv(resultsFolderEnv.c_str()); std::string cudaEnv = "CUDA_VISIBLE_DEVICES=" + std::to_string(this->GetGpuId()); itksys::SystemTools::PutEnv(cudaEnv.c_str()); spExec->Execute(this->GetPythonPath(), command, args); } catch (const mitk::Exception &e) { /* Can't throw mitk exception to the caller. Refer: T28691 */ MITK_ERROR << e.GetDescription(); return nullptr; } } if (this->GetEnsemble() && !this->GetPostProcessingJsonDirectory().empty()) { args.clear(); command = "nnUNet_ensemble"; outDir = IOUtil::CreateTemporaryDirectory("nnunet-ensemble-out-XXXXXX", this->GetMitkTempDir()); outputImagePath = outDir + IOUtil::GetDirectorySeparator() + token + "_000.nii.gz"; args.push_back("-f"); for (ModelParams &modelparam : m_ParamQ) { args.push_back(modelparam.outputDir); } args.push_back("-o"); args.push_back(outDir); if (!this->GetPostProcessingJsonDirectory().empty()) { args.push_back("-pp"); args.push_back(this->GetPostProcessingJsonDirectory()); } - for (auto arg : args) - { - MITK_INFO << arg; - } + spExec->Execute(this->GetPythonPath(), command, args); } try { - //outputImagePath = "/Users/ashis/DKFZ/nnUNet/nnUNet/out_dir/heart_single/la_000.nii.gz"; LabelSetImage::Pointer resultImage = LabelSetImage::New(); Image::Pointer outputImage = IOUtil::Load(outputImagePath); resultImage->InitializeByLabeledImage(outputImage); resultImage->SetGeometry(inputAtTimeStep->GetGeometry()); m_InputBuffer = inputAtTimeStep; return resultImage; } catch (const mitk::Exception &e) { /* Can't throw mitk exception to the caller. Refer: T28691 */ MITK_ERROR << e.GetDescription(); return nullptr; } } diff --git a/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.cpp index 97e900534a..0736edf7a8 100644 --- a/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.cpp @@ -1,438 +1,438 @@ /*============================================================================ 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. ============================================================================*/ #include "QmitknnUNetToolGUI.h" #include "mitknnUnetTool.h" #include #include #include #include MITK_TOOL_GUI_MACRO(MITKSEGMENTATIONUI_EXPORT, QmitknnUNetToolGUI, "") QmitknnUNetToolGUI::QmitknnUNetToolGUI() : QmitkAutoMLSegmentationToolGUIBase() { // Nvidia-smi command returning zero doesn't alway mean lack of GPUs. // Pytorch uses its own libraries to communicate to the GPUs. Hence, only a warning can be given. if (m_GpuLoader.GetGPUCount() == 0) { std::string warning = "WARNING: No GPUs were detected on your machine. The nnUNet tool might not work."; ShowErrorMessage(warning); } // define predicates for multi modal data selection combobox auto imageType = mitk::TNodePredicateDataType::New(); auto labelSetImageType = mitk::NodePredicateNot::New(mitk::TNodePredicateDataType::New()); this->m_MultiModalPredicate = mitk::NodePredicateAnd::New(imageType, labelSetImageType).GetPointer(); } void QmitknnUNetToolGUI::ConnectNewTool(mitk::AutoSegmentationWithPreviewTool *newTool) { Superclass::ConnectNewTool(newTool); newTool->IsTimePointChangeAwareOff(); } void QmitknnUNetToolGUI::InitializeUI(QBoxLayout *mainLayout) { m_Controls.setupUi(this); #ifndef _WIN32 m_Controls.pythonEnvComboBox->addItem("/usr/bin"); #endif m_Controls.pythonEnvComboBox->addItem("Select"); AutoParsePythonPaths(); connect(m_Controls.previewButton, SIGNAL(clicked()), this, SLOT(OnPreviewRequested())); connect(m_Controls.modeldirectoryBox, SIGNAL(directoryChanged(const QString &)), this, SLOT(OnDirectoryChanged(const QString &))); connect( m_Controls.modelBox, SIGNAL(currentTextChanged(const QString &)), this, SLOT(OnModelChanged(const QString &))); connect(m_Controls.taskBox, SIGNAL(currentTextChanged(const QString &)), this, SLOT(OnTaskChanged(const QString &))); connect( m_Controls.plannerBox, SIGNAL(currentTextChanged(const QString &)), this, SLOT(OnTrainerChanged(const QString &))); connect(m_Controls.nopipBox, SIGNAL(stateChanged(int)), this, SLOT(OnCheckBoxChanged(int))); connect(m_Controls.multiModalBox, SIGNAL(stateChanged(int)), this, SLOT(OnCheckBoxChanged(int))); connect(m_Controls.multiModalSpinBox, SIGNAL(valueChanged(int)), this, SLOT(OnModalitiesNumberChanged(int))); connect(m_Controls.posSpinBox, SIGNAL(valueChanged(int)), this, SLOT(OnModalPositionChanged(int))); connect(m_Controls.pythonEnvComboBox, #if QT_VERSION >= 0x050F00 // 5.15 SIGNAL(textActivated(const QString &)), #elif QT_VERSION >= 0x050C00 // 5.12 SIGNAL(activated(const QString &)), #endif this, SLOT(OnPythonPathChanged(const QString &))); connect(m_Controls.refreshdirectoryBox, SIGNAL(clicked()), this, SLOT(OnRefreshPresssed())); connect(m_Controls.clearCacheButton, SIGNAL(clicked()), this, SLOT(OnClearCachePressed())); m_Controls.codedirectoryBox->setVisible(false); m_Controls.nnUnetdirLabel->setVisible(false); m_Controls.multiModalSpinBox->setVisible(false); m_Controls.multiModalSpinLabel->setVisible(false); m_Controls.posSpinBoxLabel->setVisible(false); m_Controls.posSpinBox->setVisible(false); m_Controls.previewButton->setEnabled(false); QSize sz(16, 16); m_Controls.stopButton->setVisible(false); QIcon refreshIcon; refreshIcon.addPixmap(style()->standardIcon(QStyle::SP_BrowserReload).pixmap(sz), QIcon::Normal, QIcon::Off); m_Controls.refreshdirectoryBox->setIcon(refreshIcon); m_Controls.refreshdirectoryBox->setEnabled(true); m_Controls.statusLabel->setTextFormat(Qt::RichText); ShowStatusMessage(QString("STATUS: Welcome to nnUNet. " + QString::number(m_GpuLoader.GetGPUCount()) + " GPUs were detected.")); if (m_GpuLoader.GetGPUCount() != 0) { m_Controls.gpuSpinBox->setMaximum(m_GpuLoader.GetGPUCount() - 1); } mainLayout->addLayout(m_Controls.verticalLayout); Superclass::InitializeUI(mainLayout); m_UI_ROWS = m_Controls.advancedSettingsLayout->rowCount(); // Must do. Row count is correct only here. } void QmitknnUNetToolGUI::OnPreviewRequested() { mitk::nnUNetTool::Pointer tool = this->GetConnectedToolAs(); if (nullptr != tool) { try { size_t hashKey(0); m_Controls.previewButton->setEnabled(false); // To prevent misclicked back2back prediction. qApp->processEvents(); tool->PredictOn(); // purposefully placed to make tool->GetMTime different than before. QString modelName = m_Controls.modelBox->currentText(); if (modelName.startsWith("ensemble", Qt::CaseInsensitive)) { ProcessEnsembleModelsParams(tool); } else { ProcessModelParams(tool); } QString pythonPathTextItem = m_Controls.pythonEnvComboBox->currentText(); QString pythonPath = pythonPathTextItem.mid(pythonPathTextItem.indexOf(" ") + 1); bool isNoPip = m_Controls.nopipBox->isChecked(); #ifdef _WIN32 if (!isNoPip && !(pythonPath.endsWith("Scripts", Qt::CaseInsensitive) || pythonPath.endsWith("Scripts/", Qt::CaseInsensitive))) { pythonPath += QDir::separator() + QString("Scripts"); } #else if (!(pythonPath.endsWith("bin", Qt::CaseInsensitive) || pythonPath.endsWith("bin/", Qt::CaseInsensitive))) { pythonPath += QDir::separator() + QString("bin"); } #endif std::string nnUNetDirectory; if (isNoPip) { nnUNetDirectory = m_Controls.codedirectoryBox->directory().toStdString(); } - // else if (!IsNNUNetInstalled(pythonPath)) - // { - // throw std::runtime_error("nnUNet is not detected in the selected python environment. Please select a valid " - // "python environment or install nnUNet."); - // } + else if (!IsNNUNetInstalled(pythonPath)) + { + throw std::runtime_error("nnUNet is not detected in the selected python environment. Please select a valid " + "python environment or install nnUNet."); + } tool->SetnnUNetDirectory(nnUNetDirectory); tool->SetPythonPath(pythonPath.toStdString()); tool->SetModelDirectory(m_ParentFolder->getResultsFolder().toStdString()); // checkboxes tool->SetMirror(m_Controls.mirrorBox->isChecked()); tool->SetMixedPrecision(m_Controls.mixedPrecisionBox->isChecked()); tool->SetNoPip(isNoPip); tool->SetMultiModal(m_Controls.multiModalBox->isChecked()); bool doCache = m_Controls.enableCachingCheckBox->isChecked(); // Spinboxes tool->SetGpuId(static_cast(m_Controls.gpuSpinBox->value())); // Multi-Modal tool->MultiModalOff(); if (m_Controls.multiModalBox->isChecked()) { tool->m_OtherModalPaths.clear(); tool->m_OtherModalPaths = FetchMultiModalImagesFromUI(); tool->MultiModalOn(); } if (doCache) { hashKey = nnUNetCache::GetUniqueHash(tool->m_ParamQ); if (this->m_Cache.contains(hashKey)) { tool->PredictOff(); // purposefully placed to make tool->GetMTime different than before. } } if (tool->GetPredict()) { tool->m_InputBuffer = nullptr; ShowStatusMessage(QString("STATUS: Starting Segmentation task... This might take a while.")); tool->UpdatePreview(); if (tool->GetOutputBuffer() == nullptr) { SegmentationProcessFailed(); } else { SegmentationResultHandler(tool); } if (doCache) { addToCache(hashKey, tool->GetMLPreview()); } tool->PredictOff(); // purposefully placed to make tool->GetMTime different than before. } else { MITK_INFO << "won't do segmentation. Key found: " << QString::number(hashKey).toStdString(); if (this->m_Cache.contains(hashKey)) { nnUNetCache *cacheObject = this->m_Cache[hashKey]; MITK_INFO << "fetched pointer " << cacheObject->m_SegCache.GetPointer(); tool->SetNodeProperties(const_cast(cacheObject->m_SegCache.GetPointer())); SegmentationResultHandler(tool); } } m_Controls.previewButton->setEnabled(true); } catch (const std::exception &e) { std::stringstream errorMsg; errorMsg << "Error while processing parameters for nnUNet segmentation. Reason: " << e.what(); ShowErrorMessage(errorMsg.str()); ShowStatusMessage(QString::fromStdString(errorMsg.str())); m_Controls.previewButton->setEnabled(true); tool->PredictOff(); return; } catch (...) { std::string errorMsg = "Unkown error occured while generation nnUNet segmentation."; ShowErrorMessage(errorMsg); m_Controls.previewButton->setEnabled(true); tool->PredictOff(); return; } } } std::vector QmitknnUNetToolGUI::FetchMultiModalImagesFromUI() { std::vector paths; if (m_Controls.multiModalBox->isChecked() && !m_Modalities.empty()) { for (QmitkDataStorageComboBox *modality : m_Modalities) { if (modality->objectName() != "multiModal_0") { paths.push_back(dynamic_cast(modality->GetSelectedNode()->GetData())); } } } return paths; } bool QmitknnUNetToolGUI::IsNNUNetInstalled(const QString &pythonPath) { QString fullPath = pythonPath; #ifdef _WIN32 if (!(fullPath.endsWith("Scripts", Qt::CaseInsensitive) || fullPath.endsWith("Scripts/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("Scripts"); } #else if (!(fullPath.endsWith("bin", Qt::CaseInsensitive) || fullPath.endsWith("bin/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("bin"); } #endif fullPath = fullPath.mid(fullPath.indexOf(" ") + 1); return QFile::exists(fullPath + QDir::separator() + QString("nnUNet_predict")); } void QmitknnUNetToolGUI::ShowErrorMessage(const std::string &message, QMessageBox::Icon icon) { this->setCursor(Qt::ArrowCursor); QMessageBox *messageBox = new QMessageBox(icon, nullptr, message.c_str()); messageBox->exec(); delete messageBox; MITK_WARN << message; } void QmitknnUNetToolGUI::ShowStatusMessage(const QString &message) { m_Controls.statusLabel->setText(message); } void QmitknnUNetToolGUI::ProcessEnsembleModelsParams(mitk::nnUNetTool::Pointer tool) { if (m_EnsembleParams[0]->modelBox->currentText() == m_EnsembleParams[1]->modelBox->currentText()) { throw std::runtime_error("Both models you have selected for ensembling are the same."); } QString taskName = m_Controls.taskBox->currentText(); bool isPPJson = m_Controls.postProcessingCheckBox->isChecked(); std::vector requestQ; QString ppDirFolderNamePart1 = "ensemble_"; QStringList ppDirFolderNameParts; for (std::unique_ptr &layout : m_EnsembleParams) { QStringList ppDirFolderName; QString modelName = layout->modelBox->currentText(); ppDirFolderName << modelName; ppDirFolderName << "__"; QString trainer = layout->trainerBox->currentText(); ppDirFolderName << trainer; ppDirFolderName << "__"; QString planId = layout->plannerBox->currentText(); ppDirFolderName << planId; if (!IsModelExists(modelName, taskName, QString(trainer + "__" + planId))) { std::string errorMsg = "The configuration " + modelName.toStdString() + " you have selected doesn't exist. Check your Results Folder again."; throw std::runtime_error(errorMsg); } std::vector testfold = FetchSelectedFoldsFromUI(layout->foldBox); mitk::ModelParams modelObject = MapToRequest(modelName, taskName, trainer, planId, testfold); requestQ.push_back(modelObject); ppDirFolderNameParts << ppDirFolderName.join(QString("")); } tool->EnsembleOn(); if (isPPJson) { QString ppJsonFilePossibility1 = QDir::cleanPath( m_ParentFolder->getResultsFolder() + QDir::separator() + "nnUNet" + QDir::separator() + "ensembles" + QDir::separator() + taskName + QDir::separator() + ppDirFolderNamePart1 + ppDirFolderNameParts.first() + "--" + ppDirFolderNameParts.last() + QDir::separator() + "postprocessing.json"); QString ppJsonFilePossibility2 = QDir::cleanPath( m_ParentFolder->getResultsFolder() + QDir::separator() + "nnUNet" + QDir::separator() + "ensembles" + QDir::separator() + taskName + QDir::separator() + ppDirFolderNamePart1 + ppDirFolderNameParts.last() + "--" + ppDirFolderNameParts.first() + QDir::separator() + "postprocessing.json"); if (QFile(ppJsonFilePossibility1).exists()) { tool->SetPostProcessingJsonDirectory(ppJsonFilePossibility1.toStdString()); const QString statusMsg = "Post Processing JSON file found: " + ppJsonFilePossibility1; ShowStatusMessage(statusMsg); } else if (QFile(ppJsonFilePossibility2).exists()) { tool->SetPostProcessingJsonDirectory(ppJsonFilePossibility2.toStdString()); const QString statusMsg = "Post Processing JSON file found:" + ppJsonFilePossibility2; ShowStatusMessage(statusMsg); } else { std::string errorMsg = "No post processing file was found for the selected ensemble combination. Continuing anyway..."; ShowErrorMessage(errorMsg); } } tool->m_ParamQ.clear(); tool->m_ParamQ = requestQ; } void QmitknnUNetToolGUI::ProcessModelParams(mitk::nnUNetTool::Pointer tool) { tool->EnsembleOff(); std::vector requestQ; QString modelName = m_Controls.modelBox->currentText(); QString taskName = m_Controls.taskBox->currentText(); QString trainer = m_Controls.trainerBox->currentText(); QString planId = m_Controls.plannerBox->currentText(); std::vector fetchedFolds = FetchSelectedFoldsFromUI(m_Controls.foldBox); mitk::ModelParams modelObject = MapToRequest(modelName, taskName, trainer, planId, fetchedFolds); requestQ.push_back(modelObject); tool->m_ParamQ.clear(); tool->m_ParamQ = requestQ; } bool QmitknnUNetToolGUI::IsModelExists(const QString &modelName, const QString &taskName, const QString &trainerPlanner) { QString modelSearchPath = QDir::cleanPath(m_ParentFolder->getResultsFolder() + QDir::separator() + "nnUNet" + QDir::separator() + modelName + QDir::separator() + taskName + QDir::separator() + trainerPlanner); if (QDir(modelSearchPath).exists()) { return true; } return false; } void QmitknnUNetToolGUI::CheckAllInCheckableComboBox(ctkCheckableComboBox *foldBox) { // Recalling all added items to check-mark it. const QAbstractItemModel *qaim = foldBox->checkableModel(); // for (int i = 0; i < folds.size(); ++i) auto rows = qaim->rowCount(); for (int i = 0; i < rows; ++i) { const QModelIndex mi = qaim->index(i, 0); foldBox->setCheckState(mi, Qt::Checked); } } std::pair QmitknnUNetToolGUI::ExtractTrainerPlannerFromString(QStringList trainerPlanners) { QString splitterString = "__"; QStringList trainers, planners; foreach (QString trainerPlanner, trainerPlanners) { trainers << trainerPlanner.split(splitterString, QString::SplitBehavior::SkipEmptyParts).first(); planners << trainerPlanner.split(splitterString, QString::SplitBehavior::SkipEmptyParts).last(); } trainers.removeDuplicates(); planners.removeDuplicates(); std::pair returnPair = {trainers, planners}; return returnPair; } std::vector QmitknnUNetToolGUI::FetchSelectedFoldsFromUI(ctkCheckableComboBox *foldBox) { std::vector folds; if (!(foldBox->allChecked() || foldBox->noneChecked())) { QModelIndexList foldList = foldBox->checkedIndexes(); foreach (QModelIndex index, foldList) { QString foldQString = foldBox->itemText(index.row()).split("_", QString::SplitBehavior::SkipEmptyParts).last(); folds.push_back(foldQString.toStdString()); } } return folds; } void QmitknnUNetToolGUI::OnClearCachePressed() { m_Cache.clear(); UpdateCacheCountOnUI(); } void QmitknnUNetToolGUI::UpdateCacheCountOnUI() { QString cacheText = m_CACHE_COUNT_BASE_LABEL + QString::number(m_Cache.size()); m_Controls.cacheCountLabel->setText(cacheText); } void QmitknnUNetToolGUI::addToCache(size_t &hashKey, mitk::LabelSetImage::ConstPointer mlPreview) { nnUNetCache *newCacheObj = new nnUNetCache; newCacheObj->m_SegCache = mlPreview; this->m_Cache.insert(hashKey, newCacheObj); MITK_INFO << "New hash: " << hashKey << " " << newCacheObj->m_SegCache.GetPointer(); UpdateCacheCountOnUI(); } diff --git a/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.h b/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.h index e709800b1b..349dd6e6e7 100644 --- a/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.h +++ b/Modules/SegmentationUI/Qmitk/QmitknnUNetToolGUI.h @@ -1,271 +1,276 @@ /*============================================================================ 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.s ============================================================================*/ #ifndef QmitknnUNetToolGUI_h_Included #define QmitknnUNetToolGUI_h_Included #include "QmitkAutoMLSegmentationToolGUIBase.h" #include "QmitknnUNetFolderParser.h" #include "QmitknnUNetGPU.h" #include "mitknnUnetTool.h" #include "ui_QmitknnUNetToolGUIControls.h" #include #include #include #include #include #include class nnUNetCache { public: mitk::LabelSetImage::ConstPointer m_SegCache; static size_t GetUniqueHash(std::vector &requestQ) { size_t hashCode = 0; for (mitk::ModelParams &request : requestQ) { boost::hash_combine(hashCode, request.generateHash()); } return hashCode; } }; class MITKSEGMENTATIONUI_EXPORT QmitknnUNetToolGUI : public QmitkAutoMLSegmentationToolGUIBase { Q_OBJECT public: mitkClassMacro(QmitknnUNetToolGUI, QmitkAutoMLSegmentationToolGUIBase); itkFactorylessNewMacro(Self); itkCloneMacro(Self); QCache m_Cache; protected slots: /** * @brief Qt slot * */ void OnPreviewRequested(); /** * @brief Qt slot * */ void OnDirectoryChanged(const QString &); /** * @brief Qt slot * */ void OnModelChanged(const QString &); /** * @brief Qt slot * */ void OnTaskChanged(const QString &); /** * @brief Qt slot * */ void OnTrainerChanged(const QString &); /** * @brief Qt slot * */ void OnCheckBoxChanged(int); /** * @brief Qthread slot to capture failures from thread worker and * shows error message * */ void SegmentationProcessFailed(); /** * @brief Qthread to capture sucessfull nnUNet segmentation. * Further, renders the LabelSet image */ void SegmentationResultHandler(mitk::nnUNetTool *); /** * @brief Qt Slot * */ void OnModalitiesNumberChanged(int); /** * @brief Qt Slot * */ void OnPythonPathChanged(const QString &); /** * @brief Qt Slot * */ void OnModalPositionChanged(int); /** * @brief Qt slot * */ void OnRefreshPresssed(); /** * @brief Qt slot * */ void OnClearCachePressed(); protected: QmitknnUNetToolGUI(); ~QmitknnUNetToolGUI() = default; void ConnectNewTool(mitk::AutoSegmentationWithPreviewTool *newTool) override; void InitializeUI(QBoxLayout *mainLayout) override; void EnableWidgets(bool enabled) override; private: + /** + * @brief Inserts the hash and segmentation into cache and + * updates count on UI. + */ void addToCache(size_t&, mitk::LabelSetImage::ConstPointer); + /** * @brief Checks all the entries of the ctkCheckableComboBox ui widget. * This feature is not present in ctkCheckableComboBox API. */ void CheckAllInCheckableComboBox(ctkCheckableComboBox *); /** * @brief Parses the folder names containing trainer and planner together and, * returns it as separate lists. * @return std::pair */ std::pair ExtractTrainerPlannerFromString(QStringList); /** * @brief Parses the ensemble UI elements and sets to nnUNetTool object pointer. * */ void ProcessEnsembleModelsParams(mitk::nnUNetTool::Pointer); /** * @brief Parses the UI elements and sets to nnUNetTool object pointer. * */ void ProcessModelParams(mitk::nnUNetTool::Pointer); /** * @brief Creates and renders QmitknnUNetTaskParamsUITemplate layout for ensemble input. */ void ShowEnsembleLayout(bool visible = true); /** * @brief Creates a QMessage object and shows on screen. */ void ShowErrorMessage(const std::string &, QMessageBox::Icon = QMessageBox::Critical); /** * @brief Shows any message on the tool pane. */ void ShowStatusMessage(const QString &); /** * @brief Searches and parses paths of python virtual enviroments * from predefined lookout locations */ void AutoParsePythonPaths(); /** * @brief Check if pretrained model sub folder inside RESULTS FOLDER exist. * */ bool IsModelExists(const QString &, const QString &, const QString &); /** * @brief Clears all combo boxes * Any new combo box added in the future can be featured here for clearance. * */ void ClearAllComboBoxes(); /** * @brief Checks if nnUNet_predict command is valid in the selected python virtual environment. * * @return bool */ bool IsNNUNetInstalled(const QString &); /** * @brief Mapper function to map QString entries from UI to ModelParam attributes. * * @return mitk::ModelParams */ mitk::ModelParams MapToRequest( const QString &, const QString &, const QString &, const QString &, const std::vector &); /** * @brief Returns checked fold names from the ctk-Checkable-ComboBox. * * @return std::vector */ std::vector FetchSelectedFoldsFromUI(ctkCheckableComboBox *); /** * @brief Returns all paths from the dynamically generated ctk-path-line-edit boxes. * * @return std::vector */ std::vector FetchMultiModalImagesFromUI(); void UpdateCacheCountOnUI(); Ui_QmitknnUNetToolGUIControls m_Controls; QmitkGPULoader m_GpuLoader; /** * @brief Stores all dynamically added ctk-path-line-edit UI elements. * */ std::vector m_Modalities; std::vector> m_EnsembleParams; mitk::NodePredicateBase::Pointer m_MultiModalPredicate; /** * @brief Stores row count of the "advancedSettingsLayout" layout element. This value helps dynamically add * ctk-path-line-edit UI elements at the right place. Forced to initialize in the InitializeUI method since there is * no guarantee of retrieving exact row count anywhere else. * */ int m_UI_ROWS; /** * @brief Stores path of the model director (RESULTS_FOLDER appended by "nnUNet"). * */ std::shared_ptr m_ParentFolder = nullptr; /** * @brief Valid list of models supported by nnUNet * */ const QStringList m_VALID_MODELS = {"2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres", "ensembles"}; const QString m_CACHE_COUNT_BASE_LABEL = "Cached Items: "; }; #endif