diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp index 21b95bcd85..1a7c48886a 100644 --- a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp +++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp @@ -1,325 +1,384 @@ /*============================================================================ 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 #include 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_PointSet = 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_PointSet); m_PointSetNode->SetName(std::string(this->GetName()) + "_PointSet"); 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)); this->SetLabelTransferMode(LabelTransferMode::AllLabels); } void mitk::SegmentAnythingTool::Deactivated() { this->ClearSeeds(); // remove from data storage and disable interaction GetDataStorage()->Remove(m_PointSetNode); m_PointSetNode = nullptr; m_PointSet = nullptr; Superclass::Deactivated(); } void mitk::SegmentAnythingTool::ConnectActionsAndFunctions() { CONNECT_FUNCTION("ShiftSecondaryButtonPressed", OnAddPoint); CONNECT_FUNCTION("ShiftPrimaryButtonPressed", OnAddPoint); CONNECT_FUNCTION("DeletePoint", OnDelete); } 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_PointSet.IsNotNull()) { const auto positionEvent = dynamic_cast(interactionEvent); if (positionEvent != nullptr) { m_PointSet->InsertPoint(m_PointSet->GetSize(), positionEvent->GetPositionInWorld()); this->UpdatePreview(); } } } void mitk::SegmentAnythingTool::OnDelete(StateMachineAction*, InteractionEvent* /*interactionEvent*/) { if (!this->IsUpdating() && m_PointSet.IsNotNull()) { if (this->m_PointSet->GetSize() > 0) { m_PointSet->RemovePointAtEnd(0); this->UpdatePreview(); } } } void mitk::SegmentAnythingTool::ClearPicks() { this->ClearSeeds(); - this->UpdatePreview(); + //this->UpdatePreview(); } bool mitk::SegmentAnythingTool::HasPicks() const { return this->m_PointSet.IsNotNull() && this->m_PointSet->GetSize()>0; } void mitk::SegmentAnythingTool::ClearSeeds() { if (this->m_PointSet.IsNotNull()) { this->m_PointSet = mitk::PointSet::New(); // renew pointset //ensure that the seed points are visible for all timepoints. dynamic_cast(m_PointSet->GetTimeGeometry())->SetStepDuration(std::numeric_limits::max()); this->m_PointSetNode->SetData(this->m_PointSet); } } 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_PointSet.IsNotNull()) { if (this->m_MitkTempDir.empty()) { this->SetMitkTempDir(IOUtil::CreateTemporaryDirectory("mitk-XXXXXX")); + + m_InDir = IOUtil::CreateTemporaryDirectory("sam-in-XXXXXX", this->GetMitkTempDir()); + std::ofstream tmpStream; + inputImagePath = + IOUtil::CreateTemporaryFile(tmpStream, TEMPLATE_FILENAME, m_InDir + IOUtil::GetDirectorySeparator()); + tmpStream.close(); + std::size_t found = inputImagePath.find_last_of(IOUtil::GetDirectorySeparator()); + std::string fileName = inputImagePath.substr(found + 1); + token = fileName.substr(0, fileName.find("_")); + m_OutDir = IOUtil::CreateTemporaryDirectory("sam-out-XXXXXX", this->GetMitkTempDir()); } if (this->HasPicks()) { ProcessExecutor::Pointer spExec = ProcessExecutor::New(); itk::CStyleCommand::Pointer spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(&onPythonProcessEvent); spExec->AddObserver(ExternalProcessOutputEvent(), spCommand); - std::string inDir, outDir, inputImagePath, pickleFilePath, outputImagePath, scriptPath; - inDir = IOUtil::CreateTemporaryDirectory("sam-in-XXXXXX", this->GetMitkTempDir()); - std::ofstream tmpStream; - inputImagePath = IOUtil::CreateTemporaryFile(tmpStream, 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("_")); - outDir = IOUtil::CreateTemporaryDirectory("sam-out-XXXXXX", this->GetMitkTempDir()); - pickleFilePath = outDir + IOUtil::GetDirectorySeparator() + "dump.pkl"; - outputImagePath = outDir + IOUtil::GetDirectorySeparator() + token + "_000.nii.gz"; + pickleFilePath = m_OutDir + IOUtil::GetDirectorySeparator() + "dump.pkl"; + outputImagePath = m_OutDir + IOUtil::GetDirectorySeparator() + token + "_000.nii.gz"; LabelSetImage::Pointer outputBuffer; - //IOUtil::Save(inputAtTimeStep, inputImagePath); - MITK_INFO << "No.of points: " << m_PointSet->GetSize(); - auto point = m_PointSet->GetPoint(0); - MITK_INFO << point[0] << " " << point[1] << " " << point[2]; - Point2D p2D; - p2D.SetElement(0, point[0]); - p2D.SetElement(1, point[1]); - + IOUtil::Save(inputAtTimeStep, inputImagePath); - this->GetWorkingPlaneGeometry()->WorldToIndex(p2D, p2D); this->SetPythonPath("C:\\DKFZ\\SAM_work\\sam_env\\Scripts"); this->SetModelType("vit_b"); this->SetCheckpointPath("C:\\DKFZ\\SAM_work\\sam_vit_b_01ec64.pth"); - if (false)//m_IsGenerateEmbeddings) + if (m_IsGenerateEmbeddings) { this->run_generate_embeddings( - spExec, inputImagePath, outDir, this->GetModelType(), this->GetCheckpointPath(), this->GetGpuId()); + spExec, inputImagePath, m_OutDir, this->GetModelType(), this->GetCheckpointPath(), this->GetGpuId()); } - //run_segmentation_from_points( - // spExec, pickleFilePath, outputImagePath, this->GetModelType(), this->GetCheckpointPath(), this->GetGpuId()); - - outputImagePath = "C:\\DKFZ\\SAM_work\\test_seg_3d.nii.gz"; + auto pointsVec = GetPointsAsCSVString(inputAtTimeStep->GetGeometry()); + std::string pointsCSV; + std::string labelsCSV; + pointsCSV.reserve(64); + labelsCSV.reserve(m_PointSet->GetSize() * 2); + for (const std::pair& pointPair : pointsVec) + { + auto& p2D = pointPair.first; + pointsCSV.append(std::to_string(static_cast(p2D[0]))); + pointsCSV.append(","); + pointsCSV.append(std::to_string(static_cast(p2D[1]))); + pointsCSV.append(";"); + + auto &label = pointPair.second; + labelsCSV.append(label); + labelsCSV.append(","); + } + pointsCSV.pop_back(); + labelsCSV.pop_back(); + MITK_INFO << " pointsCSV " << pointsCSV; + MITK_INFO << " labelsCSV " << labelsCSV; + + run_segmentation_from_points(spExec, + pickleFilePath, + outputImagePath, + this->GetModelType(), + this->GetCheckpointPath(), + pointsCSV, + labelsCSV, + this->GetGpuId()); + + //outputImagePath = "C:\\DKFZ\\SAM_work\\test_seg_3d.nii.gz"; Image::Pointer outputImage = IOUtil::Load(outputImagePath); previewImage->InitializeByLabeledImage(outputImage); previewImage->SetGeometry(this->GetWorkingPlaneGeometry()->Clone()); } } } void mitk::SegmentAnythingTool::run_generate_embeddings(ProcessExecutor* spExec, const std::string& inputImagePath, const std::string& outputPicklePath, const std::string& modelType, const std::string& checkpointPath, const unsigned int gpuId) { ProcessExecutor::ArgumentListType args; std::string command = "python"; args.push_back("C:\\DKFZ\\SAM_work\\sam-playground\\endpoints\\generate_embedding.py"); args.push_back("--input"); args.push_back(inputImagePath); args.push_back("--output"); args.push_back(outputPicklePath); args.push_back("--model-type"); args.push_back(modelType); args.push_back("--checkpoint"); args.push_back(checkpointPath); try { std::string cudaEnv = "CUDA_VISIBLE_DEVICES=" + std::to_string(gpuId); 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; } } void mitk::SegmentAnythingTool::run_segmentation_from_points(ProcessExecutor *spExec, const std::string &pickleFilePath, const std::string &outputImagePath, const std::string &modelType, const std::string &checkpointPath, + const std::string &points, + const std::string &labels, const unsigned int gpuId) { ProcessExecutor::ArgumentListType args; std::string command = "python"; args.push_back("C:\\DKFZ\\SAM_work\\sam-playground\\endpoints\\generate_masks.py"); args.push_back("--embedding"); args.push_back(pickleFilePath); args.push_back("--output"); args.push_back(outputImagePath); args.push_back("--model-type"); args.push_back(modelType); args.push_back("--checkpoint"); args.push_back(checkpointPath); + + args.push_back("--input-points"); + args.push_back(points); + + args.push_back("--input-labels"); + args.push_back(labels); // TODO: add more arguments here-- ashis try { std::string cudaEnv = "CUDA_VISIBLE_DEVICES=" + std::to_string(gpuId); 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; } } + +bool mitk::SegmentAnythingTool::run_download_model(std::string targetDir) +{ + MITK_INFO << "model will be downloading to " << targetDir; + return true; +} + +std::vector> mitk::SegmentAnythingTool::GetPointsAsCSVString( + const mitk::BaseGeometry *baseGeometry) +{ + MITK_INFO << "No.of points: " << m_PointSet->GetSize(); + std::vector> clickVec; + clickVec.reserve(m_PointSet->GetSize()); + for (int i = 0; i < m_PointSet->GetSize(); ++i) + { + auto point = m_PointSet->GetPoint(i); + baseGeometry->WorldToIndex(point, point); + MITK_INFO << point[0] << " " << point[1] << " " << point[2]; // remove + Point2D p2D; + p2D.SetElement(0, point[0]); + p2D.SetElement(1, point[1]); + clickVec.push_back(std::pair(p2D, "1")); + } + return clickVec; +} diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h index 1b934f0785..d4dc5231e2 100644 --- a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h +++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h @@ -1,136 +1,146 @@ /*============================================================================ 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 mitkSegmentAnythingTool_h #define mitkSegmentAnythingTool_h #include "mitkSegWithPreviewTool.h" #include "mitkPointSet.h" #include "mitkProcessExecutor.h" #include namespace us { class ModuleResource; } namespace mitk { /** CHANGE THIS --ashis \brief Extracts a single region from a segmentation image and creates a new image with same geometry of the input image. The region is extracted in 3D space. This is done by performing region growing within the desired region. Use shift click to add the seed point. \ingroup ToolManagerEtAl \sa mitk::Tool \sa QmitkInteractiveSegmentation */ class MITKSEGMENTATION_EXPORT SegmentAnythingTool : public SegWithPreviewTool { public: mitkClassMacro(SegmentAnythingTool, SegWithPreviewTool); itkFactorylessNewMacro(Self); itkCloneMacro(Self); const char **GetXPM() const override; const char *GetName() const override; us::ModuleResource GetIconResource() const override; void Activated() override; void Deactivated() override; /** * Clears all picks and updates the preview. */ void ClearPicks(); bool HasPicks() const; itkSetMacro(MitkTempDir, std::string); itkGetConstMacro(MitkTempDir, std::string); itkSetMacro(PythonPath, std::string); itkGetConstMacro(PythonPath, std::string); itkSetMacro(ModelType, std::string); itkGetConstMacro(ModelType, std::string); itkSetMacro(CheckpointPath, std::string); itkGetConstMacro(CheckpointPath, std::string); itkSetMacro(GpuId, unsigned int); itkGetConstMacro(GpuId, unsigned int); itkSetMacro(IsAuto, bool); itkGetConstMacro(IsAuto, bool); itkBooleanMacro(IsAuto); itkSetMacro(IsReady, bool); itkGetConstMacro(IsReady, bool); itkBooleanMacro(IsReady); /** * @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 &e, void *); + bool run_download_model(std::string); + protected: SegmentAnythingTool(); ~SegmentAnythingTool() override; void ConnectActionsAndFunctions() override; /// \brief Add point action of StateMachine pattern virtual void OnAddPoint(StateMachineAction*, InteractionEvent* interactionEvent); /// \brief Delete action of StateMachine pattern virtual void OnDelete(StateMachineAction*, InteractionEvent* interactionEvent); /// \brief Clear all seed points. void ClearSeeds(); void DoUpdatePreview(const Image* inputAtTimeStep, const Image* oldSegAtTimeStep, LabelSetImage* previewImage, TimeStepType timeStep) override; -private: + std::vector> GetPointsAsCSVString(const mitk::BaseGeometry*); + + private: /** * @brief Runs SAM python process with desired arguments to generate embeddings for the input image * */ void run_generate_embeddings(ProcessExecutor*, const std::string&, const std::string&, const std::string&, const std::string&, const unsigned int); - void run_segmentation_from_points(ProcessExecutor *, const std::string &, const std::string &, const std::string &, const std::string &, const unsigned int); + void run_segmentation_from_points(ProcessExecutor *, const std::string &, const std::string &, const std::string &, + const std::string &, + const std::string &, + const std::string &, + const unsigned int); std::string m_MitkTempDir; std::string m_PythonPath; std::string m_ModelType; std::string m_CheckpointPath; + std::string m_InDir, m_OutDir, inputImagePath, pickleFilePath, outputImagePath, token; //rename as per standards + unsigned int m_GpuId = 0; PointSet::Pointer m_PointSet; DataNode::Pointer m_PointSetNode; bool m_IsGenerateEmbeddings = true; bool m_IsAuto = false; bool m_IsReady = false; const std::string TEMPLATE_FILENAME = "XXXXXX_000_0000.nii.gz"; }; } // namespace #endif diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingGUIControls.ui b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingGUIControls.ui new file mode 100644 index 0000000000..011eaba5d4 --- /dev/null +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingGUIControls.ui @@ -0,0 +1,372 @@ + + + QmitkSegmentAnythingGUIControls + + + + 0 + 0 + 699 + 490 + + + + + 0 + 0 + + + + + 100 + 0 + + + + + 100000 + 100000 + + + + QmitkSegmentAnythingToolWidget + + + + 0 + + + 0 + + + 0 + + + 0 + + + + + + + + 0 + 0 + + + + <html><head/><body><p>Welcome to Segment Anything Model (SAM) tool in MITK. [Experimental]</p><p>Please note that this is only an interface to SAM. MITK does not ship with SAM. Make sure to have a working internet connection to install SAM via MITK. + </p><p>Refer to <a href="https://segment-anything.com/"><span style=" text-decoration: underline; color:#0000ff;">https://segment-anything.com</span></a> to learn everything about the SAM.</p><p><br/></p></body></html> + + + Qt::RichText + + + true + + + + + + + + + 0 + 0 + + + + + 100000 + 16777215 + + + + Install SAM + + + + + + + + + 0 + 0 + + + + Install Options + + + true + + + 5 + + + true + + + + + + 0 + + + 0 + + + 0 + + + 0 + + + 6 + + + + + + + + + 0 + 0 + + + + System Python: + + + + + + + + 0 + 0 + + + + + 100000 + 16777215 + + + + Clear Install + + + + + + + + + + + + + + 0 + 0 + + + + Model Type + + + + + + + + + + + + 0 + 0 + + + + Interactive + + + + + + + true + + + + + + + + 0 + 0 + + + + Auto + + + + + + + false + + + + + + + + + 0 + 0 + + + + Press SHIFT and click to pick region(s). +Press DEL to remove last pick. + + + + + + + + + + + 0 + 0 + + + + Advanced + + + true + + + 5 + + + true + + + + + + 0 + + + 0 + + + 0 + + + 0 + + + 6 + + + + + + 0 + 0 + + + + GPU Id: + + + + + + + + + + + + + + + + 0 + 0 + + + + + 100000 + 16777215 + + + + Reset Picks + + + + + + + + 0 + 0 + + + + + 100000 + 16777215 + + + + Activate SAM + + + + + + + + 0 + 0 + + + + true + + + + + + + + + ctkComboBox + QComboBox +
ctkComboBox.h
+ 1 + + + ctkCollapsibleGroupBox + QWidget +
ctkCollapsibleGroupBox.h
+ + + + + diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp index 00b7a287c2..ba33b07199 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.cpp @@ -1,98 +1,264 @@ /*============================================================================ 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 // Mirgrate all to QML file -#include -#include -#include +#include "mitkSegmentAnythingTool.h" +#include +#include +#include +#include +#include + MITK_TOOL_GUI_MACRO(MITKSEGMENTATIONUI_EXPORT, QmitkSegmentAnythingToolGUI, "") QmitkSegmentAnythingToolGUI::QmitkSegmentAnythingToolGUI() : QmitkSegWithPreviewToolGUIBase(true) { - auto enablePickingDelegate = [this](bool enabled) + // 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; }; - - m_EnableConfirmSegBtnFnc = enablePickingDelegate; } -void QmitkSegmentAnythingToolGUI::OnResetPicksClicked() +void QmitkSegmentAnythingToolGUI::InitializeUI(QBoxLayout *mainLayout) { - auto tool = this->GetConnectedToolAs(); - if (nullptr != tool) + 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) { - tool->ClearPicks(); + welcomeText = "STATUS: Welcome to TotalSegmentator tool. You're in luck: " + + QString::number(m_GpuLoader.GetGPUCount()) + " GPU(s) were detected."; + } + else + { + welcomeText = "STATUS: Welcome to TotalSegmentator 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, + 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); + + mainLayout->addLayout(m_Controls.verticalLayout); + Superclass::InitializeUI(mainLayout); } -void QmitkSegmentAnythingToolGUI::InitializeUI(QBoxLayout* mainLayout) +void QmitkSegmentAnythingToolGUI::SetGPUInfo() { - QLabel* welcomeLabel = new QLabel("Segment Anything Model Tool. [Experimental]", this); - mainLayout->addWidget(welcomeLabel); - - auto radioPick = new QRadioButton("Interactive", this); - radioPick->setToolTip("Click on certain parts of the image to segment."); - radioPick->setChecked(true); - mainLayout->addWidget(radioPick); - m_RadioPick = radioPick; + 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)); + } +} - auto radioAutoMode = new QRadioButton("Auto", this); - radioAutoMode->setToolTip("Automatic segmentation without clicks."); - radioAutoMode->setChecked(false); - mainLayout->addWidget(radioAutoMode); - m_RadioRelabel = radioAutoMode; +void QmitkSegmentAnythingToolGUI::WriteStatusMessage(const QString &message) +{ + m_Controls.statusLabel->setText(message); + m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: white"); + qApp->processEvents(); +} - QLabel* clickLabel = new QLabel("Press SHIFT and click to pick region(s).\nPress DEL to remove last pick.", this); - mainLayout->addWidget(clickLabel); +void QmitkSegmentAnythingToolGUI::WriteErrorMessage(const QString &message) +{ + m_Controls.statusLabel->setText(message); + m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: red"); + qApp->processEvents(); +} - auto clearButton = new QPushButton("Reset picks",this); - connect(clearButton, &QPushButton::clicked, this, &QmitkSegmentAnythingToolGUI::OnResetPicksClicked); - mainLayout->addWidget(clearButton); - m_ClearPicksBtn = clearButton; +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; +} - auto activate = new QPushButton("Activate", this); - // connect(clearButton, &QPushButton::clicked, this, &QmitkSegmentAnythingToolGUI::OnResetPicksClicked); - mainLayout->addWidget(activate); +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()); + } + } + } +} - Superclass::InitializeUI(mainLayout); +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::EnableWidgets(bool enabled) +void QmitkSegmentAnythingToolGUI::OnActivateBtnClicked() { - Superclass::EnableWidgets(enabled); + auto tool = this->GetConnectedToolAs(); + if (nullptr == tool) + { + return; + } - if (nullptr != m_ClearPicksBtn) + try { - m_ClearPicksBtn->setEnabled(enabled); + 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; } - if (nullptr != m_RadioPick) +} + +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))) { - m_RadioPick->setEnabled(enabled); + fullPath += QDir::separator() + QString("Scripts"); + isPythonExists = + (!isPythonExists) ? QFile::exists(fullPath + QDir::separator() + QString("python.exe")) : isPythonExists; } - if (nullptr != m_RadioRelabel) +#else + isPythonExists = QFile::exists(fullPath + QDir::separator() + QString("python3")); + if (!(fullPath.endsWith("bin", Qt::CaseInsensitive) || fullPath.endsWith("bin/", Qt::CaseInsensitive))) { - m_RadioRelabel->setEnabled(enabled); + 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; + return isExists; } +void QmitkSegmentAnythingToolGUI::OnResetPicksClicked() +{ + auto tool = this->GetConnectedToolAs(); + if (nullptr != tool) + { + tool->ClearPicks(); + } +} \ No newline at end of file diff --git a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h index c71e3c9351..81e85aff58 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h +++ b/Modules/SegmentationUI/Qmitk/QmitkSegmentAnythingToolGUI.h @@ -1,49 +1,103 @@ /*============================================================================ 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 /** \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 : +protected slots: + /** + * @brief Qt Slot + */ void OnResetPicksClicked(); + /** + * @brief Qt Slot + */ + void OnActivateBtnClicked(); + protected: QmitkSegmentAnythingToolGUI(); ~QmitkSegmentAnythingToolGUI() = default; - void InitializeUI(QBoxLayout* mainLayout) override; - void EnableWidgets(bool enabled) override; + 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 &); private: - QWidget* m_ClearPicksBtn = nullptr; - QWidget* m_RadioPick = nullptr; - QWidget* m_RadioRelabel = nullptr; + 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/files.cmake b/Modules/SegmentationUI/files.cmake index 4d4788f862..eb83981157 100644 --- a/Modules/SegmentationUI/files.cmake +++ b/Modules/SegmentationUI/files.cmake @@ -1,120 +1,121 @@ set(CPP_FILES Qmitk/QmitkSegWithPreviewToolGUIBase.cpp Qmitk/QmitkMultiLabelSegWithPreviewToolGUIBase.cpp Qmitk/QmitkBinaryThresholdToolGUIBase.cpp Qmitk/QmitkBinaryThresholdToolGUI.cpp Qmitk/QmitkBinaryThresholdULToolGUI.cpp Qmitk/QmitkConfirmSegmentationDialog.cpp Qmitk/QmitkCopyToClipBoardDialog.cpp Qmitk/QmitkDrawPaintbrushToolGUI.cpp Qmitk/QmitkErasePaintbrushToolGUI.cpp Qmitk/QmitkEditableContourToolGUIBase.cpp Qmitk/QmitkGrowCutToolGUI.cpp Qmitk/QmitkLiveWireTool2DGUI.cpp Qmitk/QmitkLassoToolGUI.cpp Qmitk/QmitkOtsuTool3DGUI.cpp Qmitk/QmitkPaintbrushToolGUI.cpp Qmitk/QmitkPickingToolGUI.cpp Qmitk/QmitkSlicesInterpolator.cpp Qmitk/QmitkToolGUI.cpp Qmitk/QmitkToolGUIArea.cpp Qmitk/QmitkToolSelectionBox.cpp Qmitk/QmitknnUNetFolderParser.cpp Qmitk/QmitknnUNetToolGUI.cpp Qmitk/QmitknnUNetWorker.cpp Qmitk/QmitknnUNetGPU.cpp Qmitk/QmitkSurfaceStampWidget.cpp Qmitk/QmitkMaskStampWidget.cpp Qmitk/QmitkStaticDynamicSegmentationDialog.cpp Qmitk/QmitkSimpleLabelSetListWidget.cpp Qmitk/QmitkSegmentationTaskListWidget.cpp Qmitk/QmitkTotalSegmentatorToolGUI.cpp Qmitk/QmitkSetupVirtualEnvUtil.cpp Qmitk/QmitkMultiLabelInspector.cpp Qmitk/QmitkMultiLabelManager.cpp Qmitk/QmitkMultiLabelTreeModel.cpp Qmitk/QmitkMultiLabelTreeView.cpp Qmitk/QmitkLabelColorItemDelegate.cpp Qmitk/QmitkLabelToggleItemDelegate.cpp Qmitk/QmitkSegmentAnythingToolGUI.cpp SegmentationUtilities/QmitkBooleanOperationsWidget.cpp SegmentationUtilities/QmitkContourModelToImageWidget.cpp SegmentationUtilities/QmitkImageMaskingWidget.cpp SegmentationUtilities/QmitkMorphologicalOperationsWidget.cpp SegmentationUtilities/QmitkSurfaceToImageWidget.cpp SegmentationUtilities/QmitkSegmentationUtilityWidget.cpp SegmentationUtilities/QmitkDataSelectionWidget.cpp ) set(MOC_H_FILES Qmitk/QmitkSegWithPreviewToolGUIBase.h Qmitk/QmitkMultiLabelSegWithPreviewToolGUIBase.h Qmitk/QmitkBinaryThresholdToolGUIBase.h Qmitk/QmitkBinaryThresholdToolGUI.h Qmitk/QmitkBinaryThresholdULToolGUI.h Qmitk/QmitkConfirmSegmentationDialog.h Qmitk/QmitkCopyToClipBoardDialog.h Qmitk/QmitkDrawPaintbrushToolGUI.h Qmitk/QmitkErasePaintbrushToolGUI.h Qmitk/QmitkEditableContourToolGUIBase.h Qmitk/QmitkGrowCutToolGUI.h Qmitk/QmitkLiveWireTool2DGUI.h Qmitk/QmitkLassoToolGUI.h Qmitk/QmitkOtsuTool3DGUI.h Qmitk/QmitkPaintbrushToolGUI.h Qmitk/QmitkPickingToolGUI.h Qmitk/QmitkSlicesInterpolator.h Qmitk/QmitkToolGUI.h Qmitk/QmitkToolGUIArea.h Qmitk/QmitkToolSelectionBox.h Qmitk/QmitknnUNetFolderParser.h Qmitk/QmitknnUNetToolGUI.h Qmitk/QmitknnUNetGPU.h Qmitk/QmitknnUNetWorker.h Qmitk/QmitknnUNetEnsembleLayout.h Qmitk/QmitkSurfaceStampWidget.h Qmitk/QmitkMaskStampWidget.h Qmitk/QmitkStaticDynamicSegmentationDialog.h Qmitk/QmitkSimpleLabelSetListWidget.h Qmitk/QmitkSegmentationTaskListWidget.h Qmitk/QmitkTotalSegmentatorToolGUI.h Qmitk/QmitkSetupVirtualEnvUtil.h Qmitk/QmitkMultiLabelInspector.h Qmitk/QmitkMultiLabelManager.h Qmitk/QmitkMultiLabelTreeModel.h Qmitk/QmitkMultiLabelTreeView.h Qmitk/QmitkLabelColorItemDelegate.h Qmitk/QmitkLabelToggleItemDelegate.h Qmitk/QmitkSegmentAnythingToolGUI.h SegmentationUtilities/QmitkBooleanOperationsWidget.h SegmentationUtilities/QmitkContourModelToImageWidget.h SegmentationUtilities/QmitkImageMaskingWidget.h SegmentationUtilities/QmitkMorphologicalOperationsWidget.h SegmentationUtilities/QmitkSurfaceToImageWidget.h SegmentationUtilities/QmitkSegmentationUtilityWidget.h SegmentationUtilities/QmitkDataSelectionWidget.h ) set(UI_FILES Qmitk/QmitkConfirmSegmentationDialog.ui Qmitk/QmitkGrowCutToolWidgetControls.ui Qmitk/QmitkOtsuToolWidgetControls.ui Qmitk/QmitkSurfaceStampWidgetGUIControls.ui Qmitk/QmitkMaskStampWidgetGUIControls.ui Qmitk/QmitknnUNetToolGUIControls.ui Qmitk/QmitkEditableContourToolGUIControls.ui Qmitk/QmitkSegmentationTaskListWidget.ui Qmitk/QmitkTotalSegmentatorGUIControls.ui Qmitk/QmitkMultiLabelInspectorControls.ui Qmitk/QmitkMultiLabelManagerControls.ui + Qmitk/QmitkSegmentAnythingGUIControls.ui SegmentationUtilities/QmitkBooleanOperationsWidgetControls.ui SegmentationUtilities/QmitkContourModelToImageWidgetControls.ui SegmentationUtilities/QmitkImageMaskingWidgetControls.ui SegmentationUtilities/QmitkMorphologicalOperationsWidgetControls.ui SegmentationUtilities/QmitkSurfaceToImageWidgetControls.ui SegmentationUtilities/QmitkDataSelectionWidgetControls.ui ) set(QRC_FILES resources/SegmentationUI.qrc )