diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp index 7fc7f8dd1a..11afc1b276 100644 --- a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp +++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp @@ -1,431 +1,435 @@ /*============================================================================ 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 "mitkSegmentAnythingTool.h" #include "mitkProperties.h" #include "mitkToolManager.h" #include "mitkInteractionPositionEvent.h" // us #include #include #include #include #include "mitkIOUtil.h" #include "mitkSegTool2D.h" #include #include #include +#include + using namespace std::chrono_literals; namespace mitk { MITK_TOOL_MACRO(MITKSEGMENTATION_EXPORT, SegmentAnythingTool, "SegmentAnythingTool"); } mitk::SegmentAnythingTool::SegmentAnythingTool() : SegWithPreviewTool(false, "PressMoveReleaseAndPointSetting") { this->ResetsToEmptyPreviewOn(); this->IsTimePointChangeAwareOff(); } mitk::SegmentAnythingTool::~SegmentAnythingTool() { std::filesystem::remove_all(this->GetMitkTempDir()); } const char **mitk::SegmentAnythingTool::GetXPM() const { return nullptr; } const char *mitk::SegmentAnythingTool::GetName() const { return "SAM"; } us::ModuleResource mitk::SegmentAnythingTool::GetIconResource() const { us::Module *module = us::GetModuleContext()->GetModule(); us::ModuleResource resource = module->GetResource("AI.svg"); return resource; } void mitk::SegmentAnythingTool::Activated() { Superclass::Activated(); m_PointSetPositive = mitk::PointSet::New(); //ensure that the seed points are visible for all timepoints. //dynamic_cast(m_PointSet->GetTimeGeometry())->SetStepDuration(std::numeric_limits::max()); m_PointSetNode = mitk::DataNode::New(); m_PointSetNode->SetData(m_PointSetPositive); m_PointSetNode->SetName(std::string(this->GetName()) + "_PointSetPositive"); m_PointSetNode->SetBoolProperty("helper object", true); m_PointSetNode->SetColor(0.0, 1.0, 0.0); m_PointSetNode->SetVisibility(true); this->GetDataStorage()->Add(m_PointSetNode, this->GetToolManager()->GetWorkingData(0)); m_PointSetNegative = mitk::PointSet::New(); m_PointSetNodeNegative = mitk::DataNode::New(); m_PointSetNodeNegative->SetData(m_PointSetNegative); m_PointSetNodeNegative->SetName(std::string(this->GetName()) + "_PointSetNegative"); m_PointSetNodeNegative->SetBoolProperty("helper object", true); m_PointSetNodeNegative->SetColor(1.0, 0.0, 0.0); m_PointSetNodeNegative->SetVisibility(true); this->GetDataStorage()->Add(m_PointSetNodeNegative, this->GetToolManager()->GetWorkingData(0)); this->SetLabelTransferScope(LabelTransferScope::AllLabels); this->SetLabelTransferMode(LabelTransferMode::AddLabel); } void mitk::SegmentAnythingTool::Deactivated() { this->ClearSeeds(); // remove from data storage and disable interaction GetDataStorage()->Remove(m_PointSetNode); GetDataStorage()->Remove(m_PointSetNodeNegative); m_PointSetNode = nullptr; m_PointSetNodeNegative = nullptr; m_PointSetPositive = nullptr; m_PointSetNegative = nullptr; Superclass::Deactivated(); } void mitk::SegmentAnythingTool::ConnectActionsAndFunctions() { CONNECT_FUNCTION("ShiftSecondaryButtonPressed", OnAddNegativePoint); CONNECT_FUNCTION("ShiftPrimaryButtonPressed", OnAddPoint); CONNECT_FUNCTION("DeletePoint", OnDelete); } void mitk::SegmentAnythingTool::OnAddNegativePoint(StateMachineAction *, InteractionEvent *interactionEvent) { if (!this->GetIsReady() || m_PointSetPositive->GetSize() == 0) { return; } if (!this->IsUpdating() && m_PointSetNegative.IsNotNull()) { const auto positionEvent = dynamic_cast(interactionEvent); if (positionEvent != nullptr) { m_PointSetNegative->InsertPoint(m_PointSetCount, positionEvent->GetPositionInWorld()); m_PointSetCount++; this->UpdatePreview(); } } } void mitk::SegmentAnythingTool::OnAddPoint(StateMachineAction*, InteractionEvent* interactionEvent) { if (!this->GetIsReady()) { return; } m_IsGenerateEmbeddings = false; if ((nullptr == this->GetWorkingPlaneGeometry()) || !mitk::Equal(*(interactionEvent->GetSender()->GetCurrentWorldPlaneGeometry()), *(this->GetWorkingPlaneGeometry()))) { m_IsGenerateEmbeddings = true; this->ClearSeeds(); this->SetWorkingPlaneGeometry(interactionEvent->GetSender()->GetCurrentWorldPlaneGeometry()->Clone()); } if (!this->IsUpdating() && m_PointSetPositive.IsNotNull()) { const auto positionEvent = dynamic_cast(interactionEvent); if (positionEvent != nullptr) { m_PointSetPositive->InsertPoint(m_PointSetCount, positionEvent->GetPositionInWorld()); m_PointSetCount++; this->UpdatePreview(); } } } void mitk::SegmentAnythingTool::OnDelete(StateMachineAction*, InteractionEvent* /*interactionEvent*/) { if (!this->IsUpdating() && m_PointSetPositive.IsNotNull()) { PointSet::Pointer removeSet = m_PointSetPositive; decltype(m_PointSetPositive->GetMaxId().Index()) maxId = 0; if (m_PointSetPositive->GetSize() > 0) { maxId = m_PointSetPositive->GetMaxId().Index(); } if (m_PointSetNegative->GetSize() > 0 && (maxId < m_PointSetNegative->GetMaxId().Index())) { removeSet = m_PointSetNegative; } removeSet->RemovePointAtEnd(0); --m_PointSetCount; this->UpdatePreview(); } } void mitk::SegmentAnythingTool::ClearPicks() { this->ClearSeeds(); this->UpdatePreview(); } bool mitk::SegmentAnythingTool::HasPicks() const { return this->m_PointSetPositive.IsNotNull() && this->m_PointSetPositive->GetSize()>0; } void mitk::SegmentAnythingTool::ClearSeeds() { if (this->m_PointSetPositive.IsNotNull()) { m_PointSetCount -= m_PointSetPositive->GetSize(); this->m_PointSetPositive = mitk::PointSet::New(); // renew pointset //ensure that the seed points are visible for all timepoints. dynamic_cast(m_PointSetPositive->GetTimeGeometry())->SetStepDuration(std::numeric_limits::max()); this->m_PointSetNode->SetData(this->m_PointSetPositive); } if (this->m_PointSetNegative.IsNotNull()) { m_PointSetCount -= m_PointSetNegative->GetSize(); this->m_PointSetNegative = mitk::PointSet::New(); // renew pointset // ensure that the seed points are visible for all timepoints. dynamic_cast(m_PointSetNegative->GetTimeGeometry())->SetStepDuration(std::numeric_limits::max()); this->m_PointSetNodeNegative->SetData(this->m_PointSetNegative); } } void mitk::SegmentAnythingTool::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; } } void mitk::SegmentAnythingTool::DoUpdatePreview(const Image* inputAtTimeStep, const Image* oldSegAtTimeStep, LabelSetImage* previewImage, TimeStepType timeStep) { if (nullptr != oldSegAtTimeStep && nullptr != previewImage && m_PointSetPositive.IsNotNull()) { if (this->m_MitkTempDir.empty()) { //this->SetPythonPath("C:\\DKFZ\\SAM_work\\sam_env\\Scripts"); - this->SetPythonPath("C:\\DKFZ\\SAM_work\\sam-embed"); + //this->SetPythonPath("C:\\DKFZ\\SAM_work\\sam-embed"); + MITK_INFO << "Python Path: " << this->GetPythonPath(); this->SetModelType("vit_b"); this->SetCheckpointPath("C:\\DKFZ\\SAM_work\\sam_vit_b_01ec64.pth"); this->SetMitkTempDir(IOUtil::CreateTemporaryDirectory("mitk-XXXXXX")); m_InDir = IOUtil::CreateTemporaryDirectory("sam-in-XXXXXX", this->GetMitkTempDir()); m_OutDir = IOUtil::CreateTemporaryDirectory("sam-out-XXXXXX", this->GetMitkTempDir()); m_Future = std::async(std::launch::async, &mitk::SegmentAnythingTool::start_python_daemon, this); } if (this->HasPicks()) { std::string uniquePlaneID = GetHashForCurrentPlane(); std::string inputImagePath = m_InDir + IOUtil::GetDirectorySeparator() + uniquePlaneID + ".nii.gz"; std::size_t found = inputImagePath.find_last_of(IOUtil::GetDirectorySeparator()); std::string fileName = inputImagePath.substr(found + 1); token = fileName.substr(0, fileName.find("_")); outputImagePath = m_OutDir + IOUtil::GetDirectorySeparator() + uniquePlaneID + ".nii.gz"; IOUtil::Save(inputAtTimeStep, inputImagePath); auto csvStream = this->GetPointsAsCSVString(inputAtTimeStep->GetGeometry()); std::string triggerFilePath = m_InDir + IOUtil::GetDirectorySeparator() + m_TRIGGER_FILENAME; std::ofstream csvfile; csvfile.open(triggerFilePath, std::ofstream::out | std::ofstream::trunc); csvfile << csvStream.rdbuf(); csvfile.close(); auto status = m_Future.wait_for(0ms); // check if python daemon has stopped for some reason if (status == std::future_status::ready) { MITK_INFO << "Thread finished... restarting.."; m_Future = std::async(std::launch::async, &mitk::SegmentAnythingTool::start_python_daemon, this); } //outputImagePath = "C:\\DKFZ\\SAM_work\\test_seg_3d.nii.gz"; + std::this_thread::sleep_for(10ms); while (!std::filesystem::exists(outputImagePath)); Image::Pointer outputImage = IOUtil::Load(outputImagePath); //auto endloading = std::chrono::system_clock::now(); //MITK_INFO << "Loaded image in MITK. Elapsed: " // << std::chrono::duration_cast(endloading- endPython).count(); //mitk::SegTool2D::WriteSliceToVolume(previewImage, this->GetWorkingPlaneGeometry(), outputImage, timeStep, true); previewImage->InitializeByLabeledImage(outputImage); previewImage->SetGeometry(this->GetWorkingPlaneGeometry()->Clone()); std::filesystem::remove(outputImagePath); } else { previewImage->SetGeometry(this->GetWorkingPlaneGeometry()->Clone()); this->ResetPreviewContentAtTimeStep(timeStep); } } } std::string mitk::SegmentAnythingTool::GetHashForCurrentPlane() { mitk::Vector3D normal = this->GetWorkingPlaneGeometry()->GetNormal(); std::stringstream hashstream; hashstream << normal[0] << normal[1] << normal[2]; mitk::Point3D point = m_PointSetPositive->GetPoint(0); for (int i = 0; i < 3; ++i) { if (normal[i]!=0) { hashstream << point[i]; } } size_t hashVal = std::hash{}(hashstream.str()); return std::to_string(hashVal); } void mitk::SegmentAnythingTool::start_python_daemon() { ProcessExecutor::ArgumentListType args; std::string command = "python"; args.push_back("-u"); args.push_back("C:\\DKFZ\\SAM_work\\sam-playground\\endpoints\\run_inference_daemon.py"); 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(m_TRIGGER_FILENAME); args.push_back("--model-type"); args.push_back(this->GetModelType()); args.push_back("--checkpoint"); args.push_back(this->GetCheckpointPath()); ProcessExecutor::Pointer spExec = ProcessExecutor::New(); itk::CStyleCommand::Pointer spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(&onPythonProcessEvent); spExec->AddObserver(ExternalProcessOutputEvent(), spCommand); try { std::string cudaEnv = "CUDA_VISIBLE_DEVICES=" + std::to_string(this->GetGpuId()); itksys::SystemTools::PutEnv(cudaEnv.c_str()); std::stringstream logStream; for (const auto &arg : args) logStream << arg << " "; logStream << this->GetPythonPath(); MITK_INFO << logStream.str(); spExec->Execute(this->GetPythonPath(), command, args); } catch (const mitk::Exception &e) { MITK_ERROR << e.GetDescription(); return; } MITK_INFO << "ending python process....."; } bool mitk::SegmentAnythingTool::run_download_model(std::string targetDir) { MITK_INFO << "model will be downloading to " << targetDir; return true; } std::stringstream mitk::SegmentAnythingTool::GetPointsAsCSVString(const mitk::BaseGeometry* baseGeometry) { MITK_INFO << "No.of points: " << m_PointSetPositive->GetSize(); std::stringstream pointsAndLabels; pointsAndLabels << "Point,Label\n"; mitk::PointSet::PointsIterator pointSetItPos = m_PointSetPositive->Begin(); mitk::PointSet::PointsIterator pointSetItNeg = m_PointSetNegative->Begin(); const char SPACE = ' '; const char newLine = ' '; while (pointSetItPos != m_PointSetPositive->End() || pointSetItNeg != m_PointSetNegative->End()) { if (pointSetItPos != m_PointSetPositive->End()) { mitk::Point3D point = pointSetItPos.Value(); if (baseGeometry->IsInside(point)) { Point2D p2D = Get2DIndicesfrom3DWorld(baseGeometry, point); pointsAndLabels << (int)p2D[0] << SPACE << (int)p2D[1] << ",1" << std::endl; } ++pointSetItPos; } if (pointSetItNeg != m_PointSetNegative->End()) { mitk::Point3D point = pointSetItNeg.Value(); if (baseGeometry->IsInside(point)) { Point2D p2D = Get2DIndicesfrom3DWorld(baseGeometry, point); pointsAndLabels << (int)p2D[0] << SPACE << (int)p2D[1] << ",0" << std::endl; } ++pointSetItNeg; } } return pointsAndLabels; } std::vector> mitk::SegmentAnythingTool::GetPointsAsVector(const mitk::BaseGeometry *baseGeometry) { std::vector> clickVec; clickVec.reserve(m_PointSetPositive->GetSize()+m_PointSetNegative->GetSize()); mitk::PointSet::PointsIterator pointSetItPos = m_PointSetPositive->Begin(); mitk::PointSet::PointsIterator pointSetItNeg = m_PointSetNegative->Begin(); while (pointSetItPos != m_PointSetPositive->End() || pointSetItNeg != m_PointSetNegative->End()) { if (pointSetItPos != m_PointSetPositive->End()) { mitk::Point3D point = pointSetItPos.Value(); if (baseGeometry->IsInside(point)) { Point2D p2D = Get2DIndicesfrom3DWorld(baseGeometry, point); clickVec.push_back(std::pair(p2D, "1")); } ++pointSetItPos; } if (pointSetItNeg != m_PointSetNegative->End()) { mitk::Point3D point = pointSetItNeg.Value(); if (baseGeometry->IsInside(point)) { Point2D p2D = Get2DIndicesfrom3DWorld(baseGeometry, point); clickVec.push_back(std::pair(p2D, "0")); } ++pointSetItNeg; } } return clickVec; } mitk::Point2D mitk::SegmentAnythingTool::Get2DIndicesfrom3DWorld(const mitk::BaseGeometry* baseGeometry, mitk::Point3D& point3d) { baseGeometry->WorldToIndex(point3d, point3d); MITK_INFO << point3d[0] << " " << point3d[1] << " " << point3d[2]; // remove Point2D point2D; point2D.SetElement(0, point3d[0]); point2D.SetElement(1, point3d[1]); return point2D; } diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp index ba33b07199..62f44822b6 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp @@ -1,264 +1,445 @@ /*============================================================================ 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 MITK_TOOL_GUI_MACRO(MITKSEGMENTATIONUI_EXPORT, QmitkSegmentAnythingToolGUI, "") 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 TotalSegmentator 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; }; } void QmitkSegmentAnythingToolGUI::InitializeUI(QBoxLayout *mainLayout) { m_Controls.setupUi(this); #ifndef _WIN32 m_Controls.sysPythonComboBox->addItem("/usr/bin"); #endif this->AutoParsePythonPaths(); m_Controls.sysPythonComboBox->addItem("Select"); m_Controls.sysPythonComboBox->setCurrentIndex(0); m_Controls.statusLabel->setTextFormat(Qt::RichText); m_Controls.modelTypeComboBox->addItems(VALID_MODELS); QString welcomeText; this->SetGPUInfo(); if (m_GpuLoader.GetGPUCount() != 0) { - welcomeText = "STATUS: Welcome to TotalSegmentator tool. You're in luck: " + + 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 TotalSegmentator tool. Sorry, " + + 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())); - /* connect(m_Controls.clearButton, SIGNAL(clicked()), this, SLOT(OnClearInstall())); connect(m_Controls.installButton, SIGNAL(clicked()), this, SLOT(OnInstallBtnClicked())); - connect(m_Controls.sysPythonComboBox, + connect(m_Controls.clearButton, SIGNAL(clicked()), this, SLOT(OnClearInstall())); + /*connect(m_Controls.sysPythonComboBox, QOverload::of(&QComboBox::activated), [=](int index) { OnSystemPythonChanged(m_Controls.sysPythonComboBox->itemText(index)); }); */ QIcon deleteIcon = QmitkStyleManager::ThemeIcon(QStringLiteral(":/org_mitk_icons/icons/awesome/scalable/actions/edit-delete.svg")); QIcon arrowIcon = QmitkStyleManager::ThemeIcon(QStringLiteral(":/org_mitk_icons/icons/awesome/scalable/actions/go-next.svg")); m_Controls.clearButton->setIcon(deleteIcon); m_Controls.activateButton->setIcon(arrowIcon); + const QString storageDir = m_Installer.GetVirtualEnvPath(); + m_IsInstalled = this->IsSAMInstalled(storageDir); + if (m_IsInstalled) + { + m_PythonPath = GetExactPythonPath(storageDir); + m_Installer.SetVirtualEnvPath(m_PythonPath); + this->EnableAll(m_IsInstalled); + welcomeText += " SAM is already found installed."; + } + else + { + welcomeText += " SAM is not installed. Please click on \"Install SAM\" above."; + } + mainLayout->addLayout(m_Controls.verticalLayout); Superclass::InitializeUI(mainLayout); } +QString QmitkSegmentAnythingToolGUI::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 QmitkSegmentAnythingToolGUI::EnableAll(bool isEnable) +{ + m_Controls.activateButton->setEnabled(isEnable); + m_Controls.installButton->setEnabled((!isEnable)); +} + void QmitkSegmentAnythingToolGUI::SetGPUInfo() { std::vector specs = m_GpuLoader.GetAllGPUSpecs(); for (const QmitkGPUSpec &gpuSpec : specs) { m_Controls.gpuComboBox->addItem(QString::number(gpuSpec.id) + ": " + gpuSpec.name + " (" + gpuSpec.memory + ")"); } if (specs.empty()) { m_Controls.gpuComboBox->setEditable(true); m_Controls.gpuComboBox->addItem(QString::number(0)); m_Controls.gpuComboBox->setValidator(new QIntValidator(0, 999, this)); } } 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::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() + "environments"); 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_Controls.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_Controls.sysPythonComboBox->addItem("(" + envName + "): " + subIt.filePath()); } } } } unsigned int QmitkSegmentAnythingToolGUI::FetchSelectedGPUFromUI() const { QString gpuInfo = m_Controls.gpuComboBox->currentText(); 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 QmitkSegmentAnythingToolGUI::OnActivateBtnClicked() { auto tool = this->GetConnectedToolAs(); if (nullptr == tool) { return; } try { m_Controls.activateButton->setEnabled(false); qApp->processEvents(); /* if (!this->IsSAMInstalled(m_PythonPath)) { throw std::runtime_error(WARNING_SAM_NOT_FOUND); }*/ tool->SetIsAuto(m_Controls.autoRButton->isChecked()); tool->SetPythonPath(m_PythonPath.toStdString()); tool->SetGpuId(FetchSelectedGPUFromUI()); QString modelType = m_Controls.modelTypeComboBox->currentText(); std::string path = "c:\\Data"; tool->SetModelType(modelType.toStdString()); this->WriteStatusMessage( QString("STATUS: Checking if model is already downloaded... This might take a while.")); if (tool->run_download_model(path)) { this->WriteStatusMessage(QString("STATUS: Model is sucessfully found.")); tool->IsReadyOn(); } else { tool->IsReadyOff(); this->WriteStatusMessage(QString("STATUS: Model couldn't be downloaded. Please try again.")); } m_Controls.activateButton->setEnabled(true); } catch (const std::exception &e) { std::stringstream errorMsg; errorMsg << "STATUS: Error while processing parameters for SAM segmentation. Reason: " << e.what(); this->ShowErrorMessage(errorMsg.str()); this->WriteErrorMessage(QString::fromStdString(errorMsg.str())); m_Controls.activateButton->setEnabled(true); return; } catch (...) { std::string errorMsg = "Unkown error occured while generation SAM segmentation."; this->ShowErrorMessage(errorMsg); m_Controls.activateButton->setEnabled(true); return; } } +QString QmitkSegmentAnythingToolGUI::GetPythonPathFromUI(const QString &pyUI) const +{ + QString fullPath = pyUI; + if (-1 != fullPath.indexOf(")")) + { + fullPath = fullPath.mid(fullPath.indexOf(")") + 2); + } + return fullPath.simplified(); +} + +QString QmitkSegmentAnythingToolGUI::OnSystemPythonChanged(const QString &pyEnv) +{ + QString pyPath; + if (pyEnv == QString("Select")) + { + m_Controls.activateButton->setDisabled(true); + QString path = + QFileDialog::getExistingDirectory(m_Controls.sysPythonComboBox->parentWidget(), "Python Path", "dir"); + if (!path.isEmpty()) + { + this->OnSystemPythonChanged(path); // recall same function for new path validation + bool oldState = m_Controls.sysPythonComboBox->blockSignals(true); // block signal firing while inserting item + m_Controls.sysPythonComboBox->insertItem(0, path); + m_Controls.sysPythonComboBox->setCurrentIndex(0); + m_Controls.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; +} + bool QmitkSegmentAnythingToolGUI::IsSAMInstalled(const QString &pythonPath) { QString fullPath = pythonPath; bool isPythonExists = 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 - bool isExists = QFile::exists(fullPath + QDir::separator() + QString("MITK_SAM")) && isPythonExists; + bool isExists = /*QFile::exists(fullPath + QDir::separator() + QString("MITK_SAM"))*/ true && isPythonExists; return isExists; } void QmitkSegmentAnythingToolGUI::OnResetPicksClicked() { auto tool = this->GetConnectedToolAs(); if (nullptr != tool) { tool->ClearPicks(); } -} \ No newline at end of file +} + +void QmitkSegmentAnythingToolGUI::OnInstallBtnClicked() +{ + bool isInstalled = false; + //QNetworkRequest request(url); + //QNetworkReply *reply = manager.get(request); + QString systemPython = OnSystemPythonChanged(m_Controls.sysPythonComboBox->currentText()); + if (systemPython.isEmpty()) + { + this->WriteErrorMessage("ERROR: Couldn't find Python."); + } + else + { + this->WriteStatusMessage("STATUS: Installing SAM..."); + m_Installer.SetSystemPythonPath(systemPython); + isInstalled = m_Installer.SetupVirtualEnv(m_Installer.VENV_NAME); + if (isInstalled) + { + m_PythonPath = this->GetExactPythonPath(m_Installer.GetVirtualEnvPath()); + this->WriteStatusMessage("STATUS: Successfully installed SAM."); + } + else + { + this->WriteErrorMessage("ERROR: Couldn't install TotalSegmentator."); + } + } + this->EnableAll(isInstalled); +} + +void QmitkSegmentAnythingToolGUI::OnClearInstall() +{ + QDir folderPath(m_Installer.GetVirtualEnvPath()); + if (folderPath.removeRecursively()) + { + m_Controls.installButton->setEnabled(true); + m_IsInstalled = false; + } + 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."; + } +} + +bool QmitkSegmentAnythingToolInstaller::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; +} + +QString QmitkSegmentAnythingToolInstaller::GetVirtualEnvPath() +{ + return STORAGE_DIR + VENV_NAME; +} diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h index 81e85aff58..54e32b6cb6 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h @@ -1,103 +1,167 @@ /*============================================================================ 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 + + +/** + * @brief Installer class for SegmentAnythingModel Tool. + * Class specifies the virtual environment name, install version, packages required to pip install + * and implements SetupVirtualEnv method. + * + */ +class QmitkSegmentAnythingToolInstaller : public QmitkSetupVirtualEnvUtil +{ +public: + const QString VENV_NAME = ".sam"; + const QString SAM_VERSION = "1.0"; //currently, unused + const std::vector PACKAGES = {QString("numpy"), + QString("opencv-python"), + QString("git+https://github.com/facebookresearch/segment-anything.git"), + QString("SimpleITK")}; + const QString STORAGE_DIR; + inline QmitkSegmentAnythingToolInstaller( + const QString baseDir = QStandardPaths::writableLocation(QStandardPaths::GenericDataLocation) + QDir::separator() + + qApp->organizationName() + QDir::separator()) + : QmitkSetupVirtualEnvUtil(baseDir), STORAGE_DIR(baseDir){}; + bool SetupVirtualEnv(const QString &) override; + QString GetVirtualEnvPath() override; +}; + /** \ingroup org_mitk_gui_qt_interactivesegmentation_internal \brief GUI for mitk::PickingTool. \sa mitk::PickingTool */ class MITKSEGMENTATIONUI_EXPORT QmitkSegmentAnythingToolGUI : public QmitkSegWithPreviewToolGUIBase { Q_OBJECT public: mitkClassMacro(QmitkSegmentAnythingToolGUI, QmitkSegWithPreviewToolGUIBase); itkFactorylessNewMacro(Self); itkCloneMacro(Self); protected slots: /** * @brief Qt Slot */ void OnResetPicksClicked(); /** * @brief Qt Slot */ void OnActivateBtnClicked(); + /** + * @brief Qt Slot + */ + void OnInstallBtnClicked(); + + /** + * @brief Qt Slot + */ + QString OnSystemPythonChanged(const QString&); + + /** + * @brief Qt Slot + */ + void OnClearInstall(); + protected: QmitkSegmentAnythingToolGUI(); ~QmitkSegmentAnythingToolGUI() = default; 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 &); /** * @brief Creates a QMessage object and shows on screen. */ void ShowErrorMessage(const std::string &, QMessageBox::Icon = QMessageBox::Critical); /** * @brief Adds GPU information to the gpu combo box. * In case, there aren't any GPUs avaialble, the combo box will be * rendered editable. */ void SetGPUInfo(); /** * @brief Searches and parses paths of python virtual enviroments * from predefined lookout locations */ void AutoParsePythonPaths(); /** * @brief Returns GPU id of the selected GPU from the Combo box. * @return unsigned int */ unsigned int FetchSelectedGPUFromUI() const; /** * @brief Checks if TotalSegmentator command is valid in the selected python virtual environment. * @return bool */ bool IsSAMInstalled(const QString &); + /** + * @brief Get the Exact Python Path for any OS + * from the virtual environment path. + * @return QString + */ + QString GetExactPythonPath(const QString &) const; + + /** + * @brief Enable (or Disable) GUI elements. + */ + void EnableAll(bool); + + /** + * @brief Get the virtual env path from UI combobox removing any + * extra special characters. + * + * @return QString + */ + QString GetPythonPathFromUI(const QString &) const; + private: + QmitkSegmentAnythingToolInstaller m_Installer; Ui_QmitkSegmentAnythingGUIControls m_Controls; QString m_PythonPath; QmitkGPULoader m_GpuLoader; bool m_FirstPreviewComputation = true; bool m_IsInstalled = false; const std::string WARNING_SAM_NOT_FOUND = "SAM is not detected in the selected python environment.Please reinstall SAM."; const QStringList VALID_MODELS = {"vit_h (default)", "vit_l", "vit_b"}; }; #endif diff --git a/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp b/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp index 8cc7ecb0e2..0c81c41141 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp @@ -1,182 +1,183 @@ /*============================================================================ 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 ============================================================================*/ #include "QmitkSetupVirtualEnvUtil.h" #include "mitkLogMacros.h" #include #include QmitkSetupVirtualEnvUtil::QmitkSetupVirtualEnvUtil() { m_BaseDir = QStandardPaths::writableLocation(QStandardPaths::GenericDataLocation) + QDir::separator() + qApp->organizationName() + QDir::separator(); } QmitkSetupVirtualEnvUtil::QmitkSetupVirtualEnvUtil(const QString &baseDir) { m_BaseDir = baseDir; } QString& QmitkSetupVirtualEnvUtil::GetBaseDir() { return m_BaseDir; } QString QmitkSetupVirtualEnvUtil::GetVirtualEnvPath() { return m_venvPath; } QString& QmitkSetupVirtualEnvUtil::GetSystemPythonPath() { return m_SysPythonPath; } QString& QmitkSetupVirtualEnvUtil::GetPythonPath() { return m_PythonPath; } QString& QmitkSetupVirtualEnvUtil::GetPipPath() { return m_PipPath; } void QmitkSetupVirtualEnvUtil::SetVirtualEnvPath(const QString &path) { m_venvPath = path; } void QmitkSetupVirtualEnvUtil::SetPipPath(const QString &path) { m_PipPath = path; } void QmitkSetupVirtualEnvUtil::SetPythonPath(const QString &path) { if (this->IsPythonPath(path)) { m_PythonPath = path; } else { MITK_INFO << "Python was not detected in " + path.toStdString(); } } void QmitkSetupVirtualEnvUtil::SetSystemPythonPath(const QString &path) { if (this->IsPythonPath(path)) { m_SysPythonPath = path; } else { MITK_INFO << "Python was not detected in " + path.toStdString(); } } void QmitkSetupVirtualEnvUtil::PrintProcessEvent(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; } } void QmitkSetupVirtualEnvUtil::InstallPytorch(const std::string &workingDir, void (*callback)(itk::Object *, const itk::EventObject &, void *)) { mitk::ProcessExecutor::ArgumentListType args; auto spExec = mitk::ProcessExecutor::New(); auto spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(callback); spExec->AddObserver(mitk::ExternalProcessOutputEvent(), spCommand); args.push_back("-m"); args.push_back("pip"); args.push_back("install"); args.push_back("light-the-torch"); spExec->Execute(workingDir, "python", args); PipInstall("torch", workingDir, callback, "ltt"); + PipInstall("torchvision", workingDir, callback, "ltt"); } void QmitkSetupVirtualEnvUtil::InstallPytorch() { this->InstallPytorch(GetPythonPath().toStdString(), &PrintProcessEvent); } void QmitkSetupVirtualEnvUtil::PipInstall(const std::string &library, const std::string &workingDir, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command) { mitk::ProcessExecutor::ArgumentListType args; auto spExec = mitk::ProcessExecutor::New(); auto spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(callback); spExec->AddObserver(mitk::ExternalProcessOutputEvent(), spCommand); args.push_back("install"); args.push_back(library); spExec->Execute(workingDir, command, args); } void QmitkSetupVirtualEnvUtil::PipInstall(const std::string &library, void (*callback)(itk::Object*, const itk::EventObject&, void*), const std::string& command) { this->PipInstall(library, this->GetPipPath().toStdString(), callback, command); } void QmitkSetupVirtualEnvUtil::ExecutePython(const std::string &pythonCode, const std::string &workingDir, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command) { mitk::ProcessExecutor::ArgumentListType args; auto spExec = mitk::ProcessExecutor::New(); auto spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(callback); spExec->AddObserver(mitk::ExternalProcessOutputEvent(), spCommand); args.push_back("-c"); args.push_back(pythonCode); spExec->Execute(workingDir, command, args); } void QmitkSetupVirtualEnvUtil::ExecutePython(const std::string &args, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command) { this->ExecutePython(args, this->GetPythonPath().toStdString(), callback, command); } bool QmitkSetupVirtualEnvUtil::IsPythonPath(const QString &pythonPath) { QString fullPath = pythonPath; bool isExists = #ifdef _WIN32 QFile::exists(fullPath + QDir::separator() + QString("python.exe")); #else QFile::exists(fullPath + QDir::separator() + QString("python3")); #endif return isExists; }