diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp index c7fffb89ed..ed0038e61d 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp @@ -1,346 +1,360 @@ /*============================================================================ 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 "QmitkSegmentAnythingToolGUI.h" #include "mitkSegmentAnythingTool.h" #include "mitkProcessExecutor.h" #include #include #include #include #include #include #include #include #include MITK_TOOL_GUI_MACRO(MITKSEGMENTATIONUI_EXPORT, QmitkSegmentAnythingToolGUI, "") namespace { mitk::IPreferences *GetPreferences() { auto *preferencesService = mitk::CoreServices::GetPreferencesService(); return preferencesService->GetSystemPreferences()->Node("org.mitk.views.segmentation"); } } QmitkSegmentAnythingToolGUI::QmitkSegmentAnythingToolGUI() : QmitkSegWithPreviewToolGUIBase(true) { // Nvidia-smi command returning zero doesn't always imply 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 SegmentAnything tool can be very slow."; this->ShowErrorMessage(warning); } m_EnableConfirmSegBtnFnc = [this](bool enabled) { bool result = false; auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { result = enabled && tool->HasPicks(); } return result; }; m_Prefences = GetPreferences(); + m_Prefences->OnPropertyChanged += + mitk::MessageDelegate1( + this, &QmitkSegmentAnythingToolGUI::OnPreferenceChangedEvent); } QmitkSegmentAnythingToolGUI::~QmitkSegmentAnythingToolGUI() { auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { tool->SAMStatusMessageEvent -= mitk::MessageDelegate1( this, &QmitkSegmentAnythingToolGUI::StatusMessageListener); } } void QmitkSegmentAnythingToolGUI::InitializeUI(QBoxLayout *mainLayout) { m_Controls.setupUi(this); m_Controls.statusLabel->setTextFormat(Qt::RichText); QString welcomeText; if (m_GpuLoader.GetGPUCount() != 0) { welcomeText = "STATUS: Welcome to Segment Anything tool. You're in luck: " + QString::number(m_GpuLoader.GetGPUCount()) + " GPU(s) were detected."; } else { welcomeText = "STATUS: Welcome to Segment Anything tool. Sorry, " + QString::number(m_GpuLoader.GetGPUCount()) + " GPUs were detected."; } connect(m_Controls.activateButton, SIGNAL(clicked()), this, SLOT(OnActivateBtnClicked())); connect(m_Controls.resetButton, SIGNAL(clicked()), this, SLOT(OnResetPicksClicked())); QIcon arrowIcon = QmitkStyleManager::ThemeIcon( QStringLiteral(":/org_mitk_icons/icons/tango/scalable/actions/media-playback-start.svg")); m_Controls.activateButton->setIcon(arrowIcon); bool isInstalled = this->ValidatePrefences(); if (isInstalled) { m_PythonPath = QString::fromStdString(m_Prefences->Get("sam python path", "")); QString modelType = QString::fromStdString(m_Prefences->Get("sam modeltype", "")); welcomeText += " SAM is already found installed. Model type '" + modelType + "' selected in Preferences."; } else { welcomeText += " SAM tool is not configured correctly. Please go to Preferences (Cntl+P) > Segment Anything to configure and/or install SAM."; } this->EnableAll(isInstalled); this->WriteStatusMessage(welcomeText); this->ShowProgressBar(false); m_Controls.samProgressBar->setMaximum(0); mainLayout->addLayout(m_Controls.verticalLayout); Superclass::InitializeUI(mainLayout); } bool QmitkSegmentAnythingToolGUI::ValidatePrefences() { const QString storageDir = QString::fromStdString(m_Prefences->Get("sam python path", "")); bool isInstalled = QmitkSegmentAnythingToolGUI::IsSAMInstalled(storageDir); std::string modelType = m_Prefences->Get("sam modeltype", ""); std::string path = m_Prefences->Get("sam parent path", ""); return (isInstalled && !modelType.empty() && !path.empty()); } void QmitkSegmentAnythingToolGUI::EnableAll(bool isEnable) { m_Controls.activateButton->setEnabled(isEnable); } void QmitkSegmentAnythingToolGUI::WriteStatusMessage(const QString &message) { m_Controls.statusLabel->setText(message); m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: white"); qApp->processEvents(); } void QmitkSegmentAnythingToolGUI::WriteErrorMessage(const QString &message) { m_Controls.statusLabel->setText(message); m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: red"); qApp->processEvents(); } void QmitkSegmentAnythingToolGUI::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 QmitkSegmentAnythingToolGUI::StatusMessageListener(const std::string &message) { if (message.rfind("Error", 0) == 0) { this->EnableAll(true); this->WriteErrorMessage(QString::fromStdString(message)); } else { this->WriteStatusMessage(QString::fromStdString(message)); } } void QmitkSegmentAnythingToolGUI::OnActivateBtnClicked() { auto tool = this->GetConnectedToolAs(); if (nullptr == tool) { return; } try { this->EnableAll(false); qApp->processEvents(); if (!QmitkSegmentAnythingToolGUI::IsSAMInstalled(m_PythonPath)) { throw std::runtime_error(WARNING_SAM_NOT_FOUND); } tool->SetPythonPath(m_PythonPath.toStdString()); tool->SetGpuId(m_Prefences->GetInt("sam gpuid", 0)); const QString modelType = QString::fromStdString(m_Prefences->Get("sam modeltype", "")); tool->SetModelType(modelType.toStdString()); this->WriteStatusMessage( QString("STATUS: Checking if model is already downloaded... This might take a while.")); tool->SAMStatusMessageEvent += mitk::MessageDelegate1( this, &QmitkSegmentAnythingToolGUI::StatusMessageListener); if (this->DownloadModel(modelType)) { this->WriteStatusMessage(QString("STATUS: Model found. Activating Segment Anything tool...")); if (this->ActivateSAMDaemon()) { this->WriteStatusMessage(QString("STATUS: Model found. Segment Anything tool Activated.")); } else { this->WriteErrorMessage(QString("STATUS: Model found. Couldn't init tool backend.")); this->EnableAll(true); } } else { tool->IsReadyOff(); this->WriteStatusMessage(QString("STATUS: Model type not found. Starting download...")); } } catch (const std::exception &e) { std::stringstream errorMsg; errorMsg << "STATUS: Error while processing parameters for Segment Anything segmentation. Reason: " << e.what(); this->ShowErrorMessage(errorMsg.str()); this->WriteErrorMessage(QString::fromStdString(errorMsg.str())); this->EnableAll(true); return; } catch (...) { std::string errorMsg = "Unkown error occured while generation Segment Anything segmentation."; this->ShowErrorMessage(errorMsg); this->EnableAll(true); return; } } bool QmitkSegmentAnythingToolGUI::ActivateSAMDaemon() { auto tool = this->GetConnectedToolAs(); if (nullptr == tool) { return false; } this->ShowProgressBar(true); qApp->processEvents(); try { tool->InitSAMPythonProcess(); while (!tool->IsPythonReady()) { qApp->processEvents(); } tool->IsReadyOn(); } catch (...) { tool->IsReadyOff(); } this->ShowProgressBar(false); return tool->GetIsReady(); } bool QmitkSegmentAnythingToolGUI::DownloadModel(const QString &modelType) { QUrl url = QmitkSegmentAnythingToolGUI::VALID_MODELS_URL_MAP[modelType]; QString modelFileName = url.fileName(); const QString storageDir = QString::fromStdString(m_Prefences->Get("sam parent path", "")); QString checkPointPath = storageDir + QDir::separator() + modelFileName; if (QFile::exists(checkPointPath)) { auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { tool->SetCheckpointPath(checkPointPath.toStdString()); } return true; } connect(&m_Manager, SIGNAL(finished(QNetworkReply*)), this, SLOT(FileDownloaded(QNetworkReply*))); QNetworkRequest request(url); m_Manager.get(request); this->ShowProgressBar(true); return false; } void QmitkSegmentAnythingToolGUI::FileDownloaded(QNetworkReply *reply) { const QString storageDir = QString::fromStdString(m_Prefences->Get("sam parent path", "")); const QString &modelFileName = reply->url().fileName(); QFile file(storageDir + QDir::separator() + modelFileName); if (file.open(QIODevice::WriteOnly)) { file.write(reply->readAll()); file.close(); disconnect(&m_Manager, SIGNAL(finished(QNetworkReply *)), this, SLOT(FileDownloaded(QNetworkReply *))); this->WriteStatusMessage(QString("STATUS: Model successfully downloaded. Activating Segment Anything....")); auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { tool->SetCheckpointPath(file.fileName().toStdString()); if (this->ActivateSAMDaemon()) { this->WriteStatusMessage(QString("STATUS: Model successfully downloaded. Segment Anything activated.")); } else { this->WriteErrorMessage(QString("STATUS: Model successfully downloaded. But couldn't init tool backend.")); this->EnableAll(true); } } } else { this->WriteErrorMessage("STATUS: Model couldn't be downloaded. Segment Anything not activated."); } this->EnableAll(true); this->ShowProgressBar(false); } void QmitkSegmentAnythingToolGUI::ShowProgressBar(bool enabled) { m_Controls.samProgressBar->setEnabled(enabled); m_Controls.samProgressBar->setVisible(enabled); } bool QmitkSegmentAnythingToolGUI::IsSAMInstalled(const QString &pythonPath) { QString fullPath = pythonPath; bool isPythonExists = false; bool isSamExists = false; #ifdef _WIN32 isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python.exe")); if (!(fullPath.endsWith("Scripts", Qt::CaseInsensitive) || fullPath.endsWith("Scripts/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("Scripts"); isPythonExists = (!isPythonExists) ? QFile::exists(fullPath + QDir::separator() + QString("python.exe")) : isPythonExists; } #else isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python3")); if (!(fullPath.endsWith("bin", Qt::CaseInsensitive) || fullPath.endsWith("bin/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("bin"); isPythonExists = (!isPythonExists) ? QFile::exists(fullPath + QDir::separator() + QString("python3")) : isPythonExists; } #endif isSamExists = QFile::exists(fullPath + QDir::separator() + QString("run_inference_daemon.py")); bool isExists = isSamExists && isPythonExists; return isExists; } void QmitkSegmentAnythingToolGUI::OnResetPicksClicked() { auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { tool->ClearPicks(); } } + +void QmitkSegmentAnythingToolGUI::OnPreferenceChangedEvent(const mitk::IPreferences::ChangeEvent &e) +{ + this->EnableAll(true); + this->WriteStatusMessage("A Preference change was detected. Please activate the tool again."); + auto tool = this->GetConnectedToolAs(); + if (nullptr != tool) + { + tool->IsReadyOff(); + } +} diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h index 0df7fe9c70..69d99ce1d3 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h @@ -1,137 +1,138 @@ /*============================================================================ 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 QmitkSegmentAnythingToolGUI_h #define QmitkSegmentAnythingToolGUI_h #include "QmitkSegWithPreviewToolGUIBase.h" #include #include "ui_QmitkSegmentAnythingGUIControls.h" #include "QmitknnUNetGPU.h" #include "QmitkSetupVirtualEnvUtil.h" #include #include #include #include #include /** \ingroup org_mitk_gui_qt_interactivesegmentation_internal \brief GUI for mitk::SegmentAnythingTool. \sa mitk::PickingTool */ class MITKSEGMENTATIONUI_EXPORT QmitkSegmentAnythingToolGUI : public QmitkSegWithPreviewToolGUIBase { Q_OBJECT public: mitkClassMacro(QmitkSegmentAnythingToolGUI, QmitkSegWithPreviewToolGUIBase); itkFactorylessNewMacro(Self); itkCloneMacro(Self); inline static const QMap VALID_MODELS_URL_MAP = { {"vit_b", QUrl("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth")}, {"vit_l", QUrl("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth")}, {"vit_h", QUrl("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth")}}; /** * @brief Checks if SegmentAnything is found inside the selected python virtual environment. * @return bool */ static bool IsSAMInstalled(const QString &); protected slots: /** * @brief Qt Slot */ void OnResetPicksClicked(); /** * @brief Qt Slot */ void OnActivateBtnClicked(); /** * @brief Qt Slot */ void FileDownloaded(QNetworkReply*); protected: QmitkSegmentAnythingToolGUI(); ~QmitkSegmentAnythingToolGUI(); void InitializeUI(QBoxLayout* mainLayout) override; /** * @brief Writes any message in white on the tool pane. */ void WriteStatusMessage(const QString&); /** * @brief Writes any message in red on the tool pane. */ void WriteErrorMessage(const QString&); void StatusMessageListener(const std::string&); + void OnPreferenceChangedEvent(const mitk::IPreferences::ChangeEvent&); /** * @brief Creates a QMessage object and shows on screen. */ void ShowErrorMessage(const std::string&, QMessageBox::Icon = QMessageBox::Critical); /** * @brief Enable (or Disable) GUI elements. Currently, on the activate button * is affected. */ void EnableAll(bool); /** * @brief Start download process for the given model type. * Download URL is looked from the VALID_MODELS_URL_MAP. * * @return bool */ bool DownloadModel(const QString&); /** * @brief Enable (or Disable) progressbar on GUI * */ void ShowProgressBar(bool); /** * @brief Requests the tool class to spawn the SAM python daemon * process. Waits until the daemon is started. * * @return bool */ bool ActivateSAMDaemon(); /** * @brief Checks if the preferences are correctly set by the user. * * @return bool */ bool ValidatePrefences(); private: mitk::IPreferences* m_Prefences; QNetworkAccessManager m_Manager; Ui_QmitkSegmentAnythingGUIControls m_Controls; QString m_PythonPath; QmitkGPULoader m_GpuLoader; bool m_FirstPreviewComputation = true; const std::string WARNING_SAM_NOT_FOUND = "SAM is not detected in the selected python environment. Please reinstall SAM."; }; #endif