diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp index 58b9b3068b..f0227f20e5 100644 --- a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp +++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp @@ -1,261 +1,263 @@ /*============================================================================ 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 "mitkSegmentAnythingPythonService.h" #include "mitkIOUtil.h" #include #include #include #include #include #include #include "mitkImageAccessByItk.h" using namespace std::chrono_literals; using sys_clock = std::chrono::system_clock; namespace mitk { const std::string SIGNALCONSTANTS::READY = "READY"; const std::string SIGNALCONSTANTS::KILL = "KILL"; const std::string SIGNALCONSTANTS::OFF = "OFF"; const std::string SIGNALCONSTANTS::CUDA_OUT_OF_MEMORY_ERROR = "CudaOutOfMemoryError"; + const std::string SIGNALCONSTANTS::TIMEOUT_ERROR = "TimeOut"; SegmentAnythingPythonService::Status SegmentAnythingPythonService::CurrentStatus = SegmentAnythingPythonService::Status::OFF; } mitk::SegmentAnythingPythonService::SegmentAnythingPythonService( std::string workingDir, std::string modelType, std::string checkPointPath, unsigned int gpuId) : m_PythonPath(workingDir), m_ModelType(modelType), m_CheckpointPath(checkPointPath), m_GpuId(gpuId) { this->CreateTempDirs(PARENT_TEMP_DIR_PATTERN); } mitk::SegmentAnythingPythonService::~SegmentAnythingPythonService() { if (CurrentStatus == Status::READY) { this->StopAsyncProcess(); } CurrentStatus = Status::OFF; std::filesystem::remove_all(this->GetMitkTempDir()); } void mitk::SegmentAnythingPythonService::onPythonProcessEvent(itk::Object*, const itk::EventObject &e, void*) { std::string testCOUT,testCERR; const auto *pEvent = dynamic_cast(&e); if (pEvent) { testCOUT = testCOUT + pEvent->GetOutput(); testCOUT.erase(std::find_if(testCOUT.rbegin(), testCOUT.rend(), [](unsigned char ch) { return !std::isspace(ch);}).base(), testCOUT.end()); // remove trailing whitespaces, if any if (SIGNALCONSTANTS::READY == testCOUT) { CurrentStatus = Status::READY; } if (SIGNALCONSTANTS::KILL == testCOUT) { CurrentStatus = Status::KILLED; } if (SIGNALCONSTANTS::CUDA_OUT_OF_MEMORY_ERROR == testCOUT) { CurrentStatus = Status::CUDAError; } MITK_INFO << testCOUT; } const auto *pErrEvent = dynamic_cast(&e); if (pErrEvent) { testCERR = testCERR + pErrEvent->GetOutput(); MITK_ERROR << testCERR; } } void mitk::SegmentAnythingPythonService::StopAsyncProcess() { std::stringstream controlStream; controlStream << SIGNALCONSTANTS::KILL; this->WriteControlFile(controlStream); m_DaemonExec->SetStop(true); m_Future.get(); } void mitk::SegmentAnythingPythonService::StartAsyncProcess() { if (nullptr != m_DaemonExec) { this->StopAsyncProcess(); } if (this->GetMitkTempDir().empty()) { this->CreateTempDirs(PARENT_TEMP_DIR_PATTERN); } std::stringstream controlStream; controlStream << SIGNALCONSTANTS::READY; this->WriteControlFile(controlStream); double timeout = 1; m_DaemonExec = SegmentAnythingProcessExecutor::New(timeout); itk::CStyleCommand::Pointer spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(&mitk::SegmentAnythingPythonService::onPythonProcessEvent); m_DaemonExec->AddObserver(ExternalProcessOutputEvent(), spCommand); m_Future = std::async(std::launch::async, &mitk::SegmentAnythingPythonService::start_python_daemon, this); } void mitk::SegmentAnythingPythonService::TransferPointsToProcess(std::stringstream &triggerCSV) { this->CheckStatus(); std::string triggerFilePath = m_InDir + IOUtil::GetDirectorySeparator() + TRIGGER_FILENAME; std::ofstream csvfile; csvfile.open(triggerFilePath, std::ofstream::out | std::ofstream::trunc); csvfile << triggerCSV.rdbuf(); csvfile.close(); } void mitk::SegmentAnythingPythonService::WriteControlFile(std::stringstream &statusStream) { std::string controlFilePath = m_InDir + IOUtil::GetDirectorySeparator() + "control.txt"; std::ofstream controlFile; controlFile.open(controlFilePath, std::ofstream::out | std::ofstream::trunc); controlFile << statusStream.rdbuf(); controlFile.close(); } void mitk::SegmentAnythingPythonService::start_python_daemon() { ProcessExecutor::ArgumentListType args; std::string command = "python"; args.push_back("-u"); args.push_back(SAM_PYTHON_FILE_NAME); args.push_back("--input-folder"); args.push_back(m_InDir); args.push_back("--output-folder"); args.push_back(m_OutDir); args.push_back("--trigger-file"); args.push_back(TRIGGER_FILENAME); args.push_back("--model-type"); args.push_back(m_ModelType); args.push_back("--checkpoint"); args.push_back(m_CheckpointPath); args.push_back("--device"); if (m_GpuId == -1) { args.push_back("cpu"); } else { args.push_back("cuda"); std::string cudaEnv = "CUDA_VISIBLE_DEVICES=" + std::to_string(m_GpuId); itksys::SystemTools::PutEnv(cudaEnv.c_str()); } try { std::stringstream logStream; for (const auto &arg : args) logStream << arg << " "; logStream << m_PythonPath; MITK_INFO << logStream.str(); m_DaemonExec->Execute(m_PythonPath, command, args); } catch (const mitk::Exception &e) { MITK_ERROR << e.GetDescription(); return; } MITK_INFO << "Python process ended."; } bool mitk::SegmentAnythingPythonService::CheckStatus() { switch (CurrentStatus) { case mitk::SegmentAnythingPythonService::Status::READY: return true; case mitk::SegmentAnythingPythonService::Status::CUDAError: mitkThrow() << "Error: Cuda Out of Memory. Change your model type in Preferences and Activate Segment Anything tool again."; case mitk::SegmentAnythingPythonService::Status::KILLED: mitkThrow() << "Error: Python process is already terminated. Cannot load requested segmentation. Activate Segment Anything tool again."; default: return false; } } void mitk::SegmentAnythingPythonService::CreateTempDirs(const std::string &dirPattern) { this->SetMitkTempDir(IOUtil::CreateTemporaryDirectory(dirPattern)); m_InDir = IOUtil::CreateTemporaryDirectory("sam-in-XXXXXX", m_MitkTempDir); m_OutDir = IOUtil::CreateTemporaryDirectory("sam-out-XXXXXX", m_MitkTempDir); } mitk::LabelSetImage::Pointer mitk::SegmentAnythingPythonService::RetrieveImageFromProcess(long timeOut) { std::string outputImagePath = m_OutDir + IOUtil::GetDirectorySeparator() + m_CurrentUId + ".nrrd"; auto start = sys_clock::now(); while (!std::filesystem::exists(outputImagePath)) { this->CheckStatus(); std::this_thread::sleep_for(100ms); if (timeOut != -1 && std::chrono::duration_cast(sys_clock::now() - start).count() > timeOut) { CurrentStatus = Status::OFF; m_DaemonExec->SetStop(true); - mitkThrow() << "Error: Operation Timed Out!"; + mitkThrow() << SIGNALCONSTANTS::TIMEOUT_ERROR; + } } LabelSetImage::Pointer outputBuffer = mitk::IOUtil::Load(outputImagePath); return outputBuffer; } void mitk::SegmentAnythingPythonService::TransferImageToProcess(const Image *inputAtTimeStep, std::string &UId) { std::string inputImagePath = m_InDir + IOUtil::GetDirectorySeparator() + UId + ".nrrd"; if (inputAtTimeStep->GetPixelType().GetNumberOfComponents() < 2) { AccessByItk_n(inputAtTimeStep, ITKWriter, (inputImagePath)); } else { mitk::IOUtil::Save(inputAtTimeStep, inputImagePath); } m_CurrentUId = UId; } template void mitk::SegmentAnythingPythonService::ITKWriter(const itk::Image *image, std::string& outputFilename) { typedef itk::Image ImageType; typedef itk::ImageFileWriter WriterType; typename WriterType::Pointer writer = WriterType::New(); writer->SetFileName(outputFilename); writer->SetInput(image); try { writer->Update(); } catch (const itk::ExceptionObject &error) { MITK_ERROR << "Error: " << error << std::endl; mitkThrow() << "Error: " << error; } } diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h index 69172e36b1..9df2b48605 100644 --- a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h +++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h @@ -1,151 +1,152 @@ /*============================================================================ 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 mitkSegmentAnythingPythonService_h #define mitkSegmentAnythingPythonService_h #include #include #include #include #include #include #include namespace mitk { /** * @brief Segment Anything Model Python process handler class. * */ class MITKSEGMENTATION_EXPORT SegmentAnythingPythonService : public itk::Object { public: enum Status { READY, OFF, KILLED, CUDAError }; SegmentAnythingPythonService(std::string, std::string, std::string, unsigned int); ~SegmentAnythingPythonService(); itkSetMacro(MitkTempDir, std::string); itkGetConstMacro(MitkTempDir, std::string); /** * @brief Static function to print out everything from itk::EventObject. * Used as callback in mitk::ProcessExecutor object. * */ static void onPythonProcessEvent(itk::Object*, const itk::EventObject&, void*); /** * @brief Checks CurrentStatus enum variable and returns * true if daemon is READY (to read files) state, false is OFF state or * throws exception if daemon is found KILL or Cuda error state. * * @return bool */ static bool CheckStatus() /*throw(mitk::Exception)*/; /** * @brief Creates temp directories and calls start_python_daemon * function async. * */ void StartAsyncProcess(); /** * @brief Writes KILL to the control file to stop the daemon process. * */ void StopAsyncProcess(); /** * @brief Writes image as nifity file with unique id (UId) as file name. * */ void TransferImageToProcess(const Image*, std::string &UId); /** * @brief Writes csv stringstream of points to a csv file for * python daemon to read. * */ void TransferPointsToProcess(std::stringstream&); /** * @brief Waits for output nifity file from the daemon to appear and * reads it as a mitk::Image * * @return Image::Pointer */ LabelSetImage::Pointer RetrieveImageFromProcess(long timeOut= -1); static Status CurrentStatus; private: /** * @brief Runs SAM python daemon using mitk::ProcessExecutor * */ void start_python_daemon(); /** * @brief Writes stringstream content into control file. * */ void WriteControlFile(std::stringstream&); /** * @brief Create a Temp Dirs * */ void CreateTempDirs(const std::string&); /** * @brief ITK-based file writer for dumping inputs into python daemon * */ template void ITKWriter(const itk::Image *image, std::string& outputFilename); std::string m_MitkTempDir; std::string m_PythonPath; std::string m_ModelType; std::string m_CheckpointPath; std::string m_InDir, m_OutDir; std::string m_CurrentUId; int m_GpuId = 0; const std::string PARENT_TEMP_DIR_PATTERN = "mitk-sam-XXXXXX"; const std::string TRIGGER_FILENAME = "trigger.csv"; const std::string SAM_PYTHON_FILE_NAME = "run_inference_daemon.py"; std::future m_Future; //ProcessExecutor::Pointer m_DaemonExec; SegmentAnythingProcessExecutor::Pointer m_DaemonExec; }; struct SIGNALCONSTANTS { static const std::string READY; static const std::string KILL; static const std::string OFF; static const std::string CUDA_OUT_OF_MEMORY_ERROR; + static const std::string TIMEOUT_ERROR; }; } // namespace #endif diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp index 9c3a649b0c..59bcbfab63 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp @@ -1,287 +1,300 @@ /*============================================================================ 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 #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) { m_EnableConfirmSegBtnFnc = [this](bool enabled) { bool result = false; auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { result = enabled && tool->HasPicks(); } return result; }; m_Preferences = GetPreferences(); m_Preferences->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) { QString modelType = QString::fromStdString(m_Preferences->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_Preferences->Get("sam python path", "")); bool isInstalled = QmitkSegmentAnythingToolGUI::IsSAMInstalled(storageDir); std::string modelType = m_Preferences->Get("sam modeltype", ""); std::string path = m_Preferences->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) + if (message.rfind("Error", 0) == 0) { this->EnableAll(true); this->WriteErrorMessage(QString::fromStdString(message)); } + else if (message == "TimeOut") + { // trying to re init the daemon + this->WriteErrorMessage(QString("STATUS: Sorry, operation timed out. Reactivating SAM tool...")); + if (this->ActivateSAMDaemon()) + { + this->WriteStatusMessage(QString("STATUS: Segment Anything tool re-initialized.")); + } + else + { + this->WriteErrorMessage(QString("STATUS: Couldn't init tool backend.")); + this->EnableAll(true); + } + } else { this->WriteStatusMessage(QString::fromStdString(message)); } } void QmitkSegmentAnythingToolGUI::OnActivateBtnClicked() { auto tool = this->GetConnectedToolAs(); if (nullptr == tool) { return; } try { this->EnableAll(false); qApp->processEvents(); QString pythonPath = QString::fromStdString(m_Preferences->Get("sam python path", "")); if (!QmitkSegmentAnythingToolGUI::IsSAMInstalled(pythonPath)) { throw std::runtime_error(WARNING_SAM_NOT_FOUND); } tool->SetPythonPath(pythonPath.toStdString()); tool->SetGpuId(m_Preferences->GetInt("sam gpuid", -1)); const QString modelType = QString::fromStdString(m_Preferences->Get("sam modeltype", "")); tool->SetModelType(modelType.toStdString()); tool->SetTimeOutLimit(m_Preferences->GetInt("sam timeout", 300)); tool->SetCheckpointPath(m_Preferences->Get("sam parent path", "")); this->WriteStatusMessage( QString("STATUS: Initializing Segment Anything Model...")); tool->SAMStatusMessageEvent += mitk::MessageDelegate1( this, &QmitkSegmentAnythingToolGUI::StatusMessageListener); if (this->ActivateSAMDaemon()) { this->WriteStatusMessage(QString("STATUS: Segment Anything tool initialized.")); } else { this->WriteErrorMessage(QString("STATUS: Couldn't init tool backend.")); this->EnableAll(true); } } 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(); } 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&) { this->EnableAll(true); this->WriteStatusMessage("A Preference change was detected. Please initialize the tool again."); auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { tool->IsReadyOff(); } } diff --git a/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.cpp b/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.cpp index 9dd1dd1065..3229b83a62 100644 --- a/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.cpp +++ b/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.cpp @@ -1,367 +1,368 @@ /*============================================================================ 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 "QmitkSegmentAnythingPreferencePage.h" #include #include #include #include #include #include #include #include #include #include namespace { mitk::IPreferences* GetPreferences() { auto* preferencesService = mitk::CoreServices::GetPreferencesService(); return preferencesService->GetSystemPreferences()->Node("org.mitk.views.segmentation"); } } QmitkSegmentAnythingPreferencePage::QmitkSegmentAnythingPreferencePage() : m_Ui(new Ui::QmitkSegmentAnythingPreferencePage), m_Control(nullptr){} QmitkSegmentAnythingPreferencePage::~QmitkSegmentAnythingPreferencePage(){} void QmitkSegmentAnythingPreferencePage::Init(berry::IWorkbench::Pointer){} void QmitkSegmentAnythingPreferencePage::CreateQtControl(QWidget* parent) { m_Control = new QWidget(parent); m_Ui->setupUi(m_Control); #ifndef _WIN32 m_Ui->sysPythonComboBox->addItem("/usr/bin"); #endif this->AutoParsePythonPaths(); m_Ui->timeoutEdit->setValidator(new QIntValidator(0, 1000, this)); m_Ui->sysPythonComboBox->addItem("Select..."); m_Ui->sysPythonComboBox->setCurrentIndex(0); connect(m_Ui->installSAMButton, SIGNAL(clicked()), this, SLOT(OnInstallBtnClicked())); connect(m_Ui->clearSAMButton, SIGNAL(clicked()), this, SLOT(OnClearInstall())); connect(m_Ui->sysPythonComboBox, QOverload::of(&QComboBox::activated), [=](int index) { OnSystemPythonChanged(m_Ui->sysPythonComboBox->itemText(index)); }); QIcon deleteIcon = QmitkStyleManager::ThemeIcon(QStringLiteral(":/org_mitk_icons/icons/awesome/scalable/actions/edit-delete.svg")); m_Ui->clearSAMButton->setIcon(deleteIcon); const QString storageDir = m_Installer.GetVirtualEnvPath(); bool isInstalled = QmitkSegmentAnythingToolGUI::IsSAMInstalled(storageDir); QString welcomeText; if (isInstalled) { m_PythonPath = GetExactPythonPath(storageDir); m_Installer.SetVirtualEnvPath(m_PythonPath); welcomeText += " Segment Anything tool is already found installed."; m_Ui->installSAMButton->setEnabled(false); } else { welcomeText += " Segment Anything tool not installed. Please click on \"Install SAM\" above. \ The installation will create a new virtual environment using the System Python selected above."; m_Ui->installSAMButton->setEnabled(true); } this->WriteStatusMessage(welcomeText); m_Ui->samModelTypeComboBox->addItems(VALID_MODELS); m_Ui->gpuComboBox->addItem(CPU_ID); this->SetGPUInfo(); this->Update(); } QWidget* QmitkSegmentAnythingPreferencePage::GetQtControl() const { return m_Control; } bool QmitkSegmentAnythingPreferencePage::PerformOk() { auto* prefs = GetPreferences(); prefs->Put("sam parent path", m_Installer.STORAGE_DIR.toStdString()); prefs->Put("sam python path", m_PythonPath.toStdString()); prefs->Put("sam modeltype", m_Ui->samModelTypeComboBox->currentText().toStdString()); prefs->PutInt("sam gpuid", FetchSelectedGPUFromUI()); prefs->PutInt("sam timeout", std::stoi(m_Ui->timeoutEdit->text().toStdString())); return true; } void QmitkSegmentAnythingPreferencePage::PerformCancel(){} void QmitkSegmentAnythingPreferencePage::Update() { auto* prefs = GetPreferences(); m_Ui->samModelTypeComboBox->setCurrentText(QString::fromStdString(prefs->Get("sam modeltype", "vit_b"))); + m_Ui->timeoutEdit->setText(QString::number(prefs->GetInt("sam timeout", 300))); int gpuId = prefs->GetInt("sam gpuid", -1); if (gpuId == -1) { m_Ui->gpuComboBox->setCurrentText(CPU_ID); } else if (m_GpuLoader.GetGPUCount() == 0) { m_Ui->gpuComboBox->setCurrentText(QString::number(gpuId)); } else { std::vector specs = m_GpuLoader.GetAllGPUSpecs(); QmitkGPUSpec gpuSpec = specs[gpuId]; m_Ui->gpuComboBox->setCurrentText(QString::number(gpuSpec.id) + ": " + gpuSpec.name + " (" + gpuSpec.memory + ")"); } } QString QmitkSegmentAnythingPreferencePage::OnSystemPythonChanged(const QString &pyEnv) { QString pyPath; if (pyEnv == QString("Select...")) { QString path = QFileDialog::getExistingDirectory(m_Ui->sysPythonComboBox->parentWidget(), "Python Path", "dir"); if (!path.isEmpty()) { this->OnSystemPythonChanged(path); // recall same function for new path validation bool oldState = m_Ui->sysPythonComboBox->blockSignals(true); // block signal firing while inserting item m_Ui->sysPythonComboBox->insertItem(0, path); m_Ui->sysPythonComboBox->setCurrentIndex(0); m_Ui->sysPythonComboBox->blockSignals(oldState); // unblock signal firing after inserting item. Remove this after Qt6 migration } } else { QString uiPyPath = this->GetPythonPathFromUI(pyEnv); pyPath = this->GetExactPythonPath(uiPyPath); } return pyPath; } QString QmitkSegmentAnythingPreferencePage::GetPythonPathFromUI(const QString &pyUI) const { QString fullPath = pyUI; if (-1 != fullPath.indexOf(")")) { fullPath = fullPath.mid(fullPath.indexOf(")") + 2); } return fullPath.simplified(); } QString QmitkSegmentAnythingPreferencePage::GetExactPythonPath(const QString &pyEnv) const { QString fullPath = pyEnv; bool isPythonExists = false; #ifdef _WIN32 isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python.exe")); if (!isPythonExists && !(fullPath.endsWith("Scripts", Qt::CaseInsensitive) || fullPath.endsWith("Scripts/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("Scripts"); isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python.exe")); } #else isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python3")); if (!isPythonExists && !(fullPath.endsWith("bin", Qt::CaseInsensitive) || fullPath.endsWith("bin/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("bin"); isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python3")); } #endif if (!isPythonExists) { fullPath.clear(); } return fullPath; } void QmitkSegmentAnythingPreferencePage::AutoParsePythonPaths() { QString homeDir = QDir::homePath(); std::vector searchDirs; #ifdef _WIN32 searchDirs.push_back(QString("C:") + QDir::separator() + QString("ProgramData") + QDir::separator() + QString("anaconda3")); #else // Add search locations for possible standard python paths here searchDirs.push_back(homeDir + QDir::separator() + "anaconda3"); searchDirs.push_back(homeDir + QDir::separator() + "miniconda3"); searchDirs.push_back(homeDir + QDir::separator() + "opt" + QDir::separator() + "miniconda3"); searchDirs.push_back(homeDir + QDir::separator() + "opt" + QDir::separator() + "anaconda3"); #endif for (QString searchDir : searchDirs) { if (searchDir.endsWith("anaconda3", Qt::CaseInsensitive)) { if (QDir(searchDir).exists()) { m_Ui->sysPythonComboBox->addItem("(base): " + searchDir); searchDir.append((QDir::separator() + QString("envs"))); } } for (QDirIterator subIt(searchDir, QDir::AllDirs, QDirIterator::NoIteratorFlags); subIt.hasNext();) { subIt.next(); QString envName = subIt.fileName(); if (!envName.startsWith('.')) // Filter out irrelevent hidden folders, if any. { m_Ui->sysPythonComboBox->addItem("(" + envName + "): " + subIt.filePath()); } } } } void QmitkSegmentAnythingPreferencePage::SetGPUInfo() { std::vector specs = m_GpuLoader.GetAllGPUSpecs(); for (const QmitkGPUSpec &gpuSpec : specs) { m_Ui->gpuComboBox->addItem(QString::number(gpuSpec.id) + ": " + gpuSpec.name + " (" + gpuSpec.memory + ")"); } if (specs.empty()) { m_Ui->gpuComboBox->setCurrentIndex(m_Ui->gpuComboBox->findText("cpu")); } else { m_Ui->gpuComboBox->setCurrentIndex(m_Ui->gpuComboBox->count()-1); } } int QmitkSegmentAnythingPreferencePage::FetchSelectedGPUFromUI() const { QString gpuInfo = m_Ui->gpuComboBox->currentText(); if ("cpu" == gpuInfo) { return -1; } else if(m_GpuLoader.GetGPUCount() == 0) { return static_cast(gpuInfo.toInt()); } else { QString gpuId = gpuInfo.split(":", QString::SplitBehavior::SkipEmptyParts).first(); return static_cast(gpuId.toInt()); } } void QmitkSegmentAnythingPreferencePage::OnInstallBtnClicked() { QString systemPython = OnSystemPythonChanged(m_Ui->sysPythonComboBox->currentText()); if (!systemPython.isEmpty()) { this->WriteStatusMessage("STATUS: Installing SAM..."); m_Ui->installSAMButton->setEnabled(false); m_Installer.SetSystemPythonPath(systemPython); bool isInstalled = false; bool isFinished = m_Installer.SetupVirtualEnv(m_Installer.VENV_NAME); if (isFinished) { isInstalled = QmitkSegmentAnythingToolGUI::IsSAMInstalled(m_Installer.GetVirtualEnvPath()); } if (isInstalled) { m_PythonPath = this->GetExactPythonPath(m_Installer.GetVirtualEnvPath()); this->WriteStatusMessage("STATUS: Successfully installed SAM."); } else { this->WriteErrorMessage("ERROR: Couldn't install SAM."); m_Ui->installSAMButton->setEnabled(true); } } } void QmitkSegmentAnythingPreferencePage::OnClearInstall() { QDir folderPath(m_Installer.GetVirtualEnvPath()); bool isDeleted = folderPath.removeRecursively(); if (isDeleted) { this->WriteStatusMessage("Deleted SAM installation."); m_Ui->installSAMButton->setEnabled(true); m_PythonPath.clear(); } else { MITK_ERROR << "The virtual environment couldn't be removed. Please check if you have the required access " "privileges or, some other process is accessing the folders."; } } void QmitkSegmentAnythingPreferencePage::WriteStatusMessage(const QString &message) { m_Ui->samInstallStatusLabel->setText(message); m_Ui->samInstallStatusLabel->setStyleSheet("font-weight: bold; color: white"); qApp->processEvents(); } void QmitkSegmentAnythingPreferencePage::WriteErrorMessage(const QString &message) { m_Ui->samInstallStatusLabel->setText(message); m_Ui->samInstallStatusLabel->setStyleSheet("font-weight: bold; color: red"); qApp->processEvents(); } QString QmitkSAMInstaller::GetVirtualEnvPath() { return STORAGE_DIR + VENV_NAME; } bool QmitkSAMInstaller::SetupVirtualEnv(const QString &venvName) { if (GetSystemPythonPath().isEmpty()) { return false; } QDir folderPath(GetBaseDir()); folderPath.mkdir(venvName); if (!folderPath.cd(venvName)) { return false; // Check if directory creation was successful. } mitk::ProcessExecutor::ArgumentListType args; auto spExec = mitk::ProcessExecutor::New(); auto spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(&PrintProcessEvent); spExec->AddObserver(mitk::ExternalProcessOutputEvent(), spCommand); args.push_back("-m"); args.push_back("venv"); args.push_back(venvName.toStdString()); #ifdef _WIN32 QString pythonFile = GetSystemPythonPath() + QDir::separator() + "python.exe"; QString pythonExeFolder = "Scripts"; #else QString pythonFile = GetSystemPythonPath() + QDir::separator() + "python3"; QString pythonExeFolder = "bin"; #endif spExec->Execute(GetBaseDir().toStdString(), pythonFile.toStdString(), args); // Setup local virtual environment if (folderPath.cd(pythonExeFolder)) { this->SetPythonPath(folderPath.absolutePath()); this->SetPipPath(folderPath.absolutePath()); this->InstallPytorch(); for (auto &package : PACKAGES) { this->PipInstall(package.toStdString(), &PrintProcessEvent); } std::string pythonCode; // python syntax to check if torch is installed with CUDA. pythonCode.append("import torch;"); pythonCode.append("print('Pytorch was installed with CUDA') if torch.cuda.is_available() else print('PyTorch was " "installed WITHOUT CUDA');"); this->ExecutePython(pythonCode, &PrintProcessEvent); return true; } return false; }