diff --git a/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.cpp b/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.cpp index 9d93d41ccf..83d50e0fdd 100644 --- a/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.cpp +++ b/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.cpp @@ -1,349 +1,359 @@ /*============================================================================ 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. ============================================================================*/ // MITK #include "mitkTotalSegmentatorTool.h" #include #include #include #include #include #include // us #include #include #include #include #include namespace mitk { MITK_TOOL_MACRO(MITKSEGMENTATION_EXPORT, TotalSegmentatorTool, "Total Segmentator"); } mitk::TotalSegmentatorTool::~TotalSegmentatorTool() { std::filesystem::remove_all(this->GetMitkTempDir()); } mitk::TotalSegmentatorTool::TotalSegmentatorTool() : SegWithPreviewTool(true) // prevents auto-compute across all timesteps { this->IsTimePointChangeAwareOff(); this->RequestDeactivationConfirmationOn(); } void mitk::TotalSegmentatorTool::Activated() { Superclass::Activated(); this->SetLabelTransferScope(LabelTransferScope::AllLabels); this->SetLabelTransferMode(LabelTransferMode::AddLabel); } const char **mitk::TotalSegmentatorTool::GetXPM() const { return nullptr; } us::ModuleResource mitk::TotalSegmentatorTool::GetIconResource() const { us::Module *module = us::GetModuleContext()->GetModule(); us::ModuleResource resource = module->GetResource("AI.svg"); return resource; } const char *mitk::TotalSegmentatorTool::GetName() const { return "TotalSegmentator"; } void mitk::TotalSegmentatorTool::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::TotalSegmentatorTool::DoUpdatePreview(const Image *inputAtTimeStep, const Image * /*oldSegAtTimeStep*/, LabelSetImage *previewImage, TimeStepType timeStep) { if (this->m_MitkTempDir.empty()) { this->SetMitkTempDir(IOUtil::CreateTemporaryDirectory("mitk-XXXXXX")); } ProcessExecutor::Pointer spExec = ProcessExecutor::New(); itk::CStyleCommand::Pointer spCommand = itk::CStyleCommand::New(); spCommand->SetCallback(&onPythonProcessEvent); spExec->AddObserver(ExternalProcessOutputEvent(), spCommand); m_ProgressCommand->SetProgress(5); std::string inDir, outDir, inputImagePath, outputImagePath, scriptPath; inDir = IOUtil::CreateTemporaryDirectory("totalseg-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("totalseg-out-XXXXXX", this->GetMitkTempDir()); LabelSetImage::Pointer outputBuffer; m_ProgressCommand->SetProgress(20); IOUtil::Save(inputAtTimeStep, inputImagePath); m_ProgressCommand->SetProgress(50); outputImagePath = outDir + IOUtil::GetDirectorySeparator() + token + "_000.nii.gz"; - const bool isSubTask = (this->GetSubTask() != DEFAULT_TOTAL_TASK); + const bool isSubTask = (this->GetSubTask() != DEFAULT_TOTAL_TASK) && (this->GetSubTask() != DEFAULT_TOTAL_TASK_MRI); if (isSubTask) { outputImagePath = outDir; this->run_totalsegmentator( spExec, inputImagePath, outputImagePath, !isSubTask, !isSubTask, this->GetGpuId(), this->GetSubTask()); // Construct Label Id map std::vector files = SUBTASKS_MAP.at(this->GetSubTask()); // Agglomerate individual mask files into one multi-label image. std::for_each(files.begin(), files.end(), [&](std::string &fileName) { fileName = (outDir + IOUtil::GetDirectorySeparator() + fileName); }); outputBuffer = AgglomerateLabelFiles(files, inputAtTimeStep->GetDimensions(), inputAtTimeStep->GetGeometry()); } else { this->run_totalsegmentator( - spExec, inputImagePath, outputImagePath, this->GetFast(), !isSubTask, this->GetGpuId(), DEFAULT_TOTAL_TASK); + spExec, inputImagePath, outputImagePath, this->GetFast(), !isSubTask, this->GetGpuId(), this->GetSubTask()); Image::Pointer outputImage = IOUtil::Load(outputImagePath); outputBuffer = mitk::LabelSetImage::New(); outputBuffer->InitializeByLabeledImage(outputImage); outputBuffer->SetGeometry(inputAtTimeStep->GetGeometry()); } m_ProgressCommand->SetProgress(180); mitk::ImageReadAccessor newMitkImgAcc(outputBuffer.GetPointer()); this->MapLabelsToSegmentation(outputBuffer, previewImage, m_LabelMapTotal); previewImage->SetVolume(newMitkImgAcc.GetData(), timeStep); } void mitk::TotalSegmentatorTool::UpdatePrepare() { Superclass::UpdatePrepare(); auto preview = this->GetPreviewSegmentation(); preview->RemoveLabels(preview->GetAllLabelValues()); if (m_LabelMapTotal.empty()) { this->ParseLabelMapTotalDefault(); } - const bool isSubTask = (this->GetSubTask() != DEFAULT_TOTAL_TASK); + const bool isSubTask = (this->GetSubTask() != DEFAULT_TOTAL_TASK) && (this->GetSubTask() != DEFAULT_TOTAL_TASK_MRI); if (isSubTask) { std::vector files = SUBTASKS_MAP.at(this->GetSubTask()); m_LabelMapTotal.clear(); mitk::Label::PixelType labelId = 1; for (auto const &file : files) { std::string labelName = file.substr(0, file.find('.')); m_LabelMapTotal[labelId] = labelName; labelId++; } } } mitk::LabelSetImage::Pointer mitk::TotalSegmentatorTool::AgglomerateLabelFiles(std::vector &filePaths, const unsigned int *dimensions, mitk::BaseGeometry *geometry) { Label::PixelType labelId = 1; auto aggloLabelImage = mitk::LabelSetImage::New(); auto initImage = mitk::Image::New(); initImage->Initialize(mitk::MakeScalarPixelType(), 3, dimensions); aggloLabelImage->Initialize(initImage); aggloLabelImage->SetGeometry(geometry); const auto layerIndex = aggloLabelImage->AddLayer(); aggloLabelImage->SetActiveLayer(layerIndex); for (auto const &outputImagePath : filePaths) { double rgba[4]; aggloLabelImage->GetLookupTable()->GetTableValue(labelId, rgba); mitk::Color color; color.SetRed(rgba[0]); color.SetGreen(rgba[1]); color.SetBlue(rgba[2]); auto label = mitk::Label::New(); label->SetName("object-" + std::to_string(labelId)); label->SetValue(labelId); label->SetColor(color); label->SetOpacity(rgba[3]); aggloLabelImage->AddLabel(label, layerIndex, false, false); Image::Pointer outputImage = IOUtil::Load(outputImagePath); auto source = mitk::LabelSetImage::New(); source->InitializeByLabeledImage(outputImage); source->SetGeometry(geometry); mitk::TransferLabelContent(source, aggloLabelImage, aggloLabelImage->GetConstLabelsByValue(aggloLabelImage->GetLabelValuesByGroup(layerIndex)), 0, 0, false, {{1, labelId}}); labelId++; } return aggloLabelImage; } void mitk::TotalSegmentatorTool::run_totalsegmentator(ProcessExecutor* spExec, const std::string &inputImagePath, const std::string &outputImagePath, bool isFast, bool isMultiLabel, unsigned int gpuId, const std::string &subTask) { ProcessExecutor::ArgumentListType args; std::string command = "TotalSegmentator"; #ifdef _WIN32 command += ".exe"; #endif args.clear(); args.push_back("-i"); args.push_back(inputImagePath); args.push_back("-o"); args.push_back(outputImagePath); - if (subTask != DEFAULT_TOTAL_TASK) + if (subTask == DEFAULT_TOTAL_TASK_MRI) + { + args.push_back("--task"); + args.push_back(subTask); + } + else if (subTask != DEFAULT_TOTAL_TASK) { args.push_back("-ta"); args.push_back(subTask); } if (isMultiLabel) { args.push_back("--ml"); } if (isFast) { args.push_back("--fast"); } 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::TotalSegmentatorTool::ParseLabelMapTotalDefault() { if (!this->GetLabelMapPath().empty()) { + int start_line = 0, end_line = 0; + if (this->GetSubTask() == DEFAULT_TOTAL_TASK) + start_line = 111, end_line = 229; + else if (this->GetSubTask() == DEFAULT_TOTAL_TASK_MRI) + start_line = 231, end_line = 288; std::regex sanitizer(R"([^A-Za-z0-9_])"); std::fstream newfile; newfile.open(this->GetLabelMapPath(), ios::in); std::stringstream buffer; if (newfile.is_open()) { int line = 0; std::string temp; while (std::getline(newfile, temp)) { - if (line > 111 && line < 229) + if (line > start_line && line < end_line) { buffer << temp; } ++line; } } std::string key, val; while (std::getline(std::getline(buffer, key, ':'), val, ',')) { std::string sanitized = std::regex_replace(val, sanitizer, ""); m_LabelMapTotal[std::stoi(key)] = sanitized; } } } void mitk::TotalSegmentatorTool::MapLabelsToSegmentation(const mitk::LabelSetImage* source, mitk::LabelSetImage* dest, std::map &labelMap) { auto lookupTable = mitk::LookupTable::New(); lookupTable->SetType(mitk::LookupTable::LookupTableType::MULTILABEL); for (auto const &[key, val] : labelMap) { if (source->ExistLabel(key, source->GetActiveLayer())) { Label::Pointer label = Label::New(key, val); std::array lookupTableColor; lookupTable->GetColor(key, lookupTableColor.data()); Color color; color.SetRed(lookupTableColor[0]); color.SetGreen(lookupTableColor[1]); color.SetBlue(lookupTableColor[2]); label->SetColor(color); dest->AddLabel(label, 0,false); } } } std::string mitk::TotalSegmentatorTool::GetLabelMapPath() { std::string pythonFileName; std::filesystem::path pathToLabelMap(this->GetPythonPath()); pathToLabelMap = pathToLabelMap.parent_path(); #ifdef _WIN32 pythonFileName = pathToLabelMap.string() + "/Lib/site-packages/totalsegmentator/map_to_binary.py"; #else pathToLabelMap.append("lib"); for (auto const &dir_entry : std::filesystem::directory_iterator{pathToLabelMap}) { if (dir_entry.is_directory()) { auto dirName = dir_entry.path().filename().string(); if (dirName.rfind("python", 0) == 0) { pathToLabelMap.append(dir_entry.path().filename().string()); break; } } } pythonFileName = pathToLabelMap.string() + "/site-packages/totalsegmentator/map_to_binary.py"; #endif return pythonFileName; } diff --git a/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.h b/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.h index ca4b68fe21..3e86e95220 100644 --- a/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.h +++ b/Modules/Segmentation/Interactions/mitkTotalSegmentatorTool.h @@ -1,150 +1,151 @@ /*============================================================================ 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 MITKTOTALSEGMENTATORTOOL_H #define MITKTOTALSEGMENTATORTOOL_H #include "mitkSegWithPreviewTool.h" #include #include "mitkProcessExecutor.h" namespace us { class ModuleResource; } namespace mitk { /** \brief TotalSegmentator segmentation tool. \ingroup Interaction \ingroup ToolManagerEtAl \warning Only to be instantiated by mitk::ToolManager. */ class MITKSEGMENTATION_EXPORT TotalSegmentatorTool : public SegWithPreviewTool { public: mitkClassMacro(TotalSegmentatorTool, SegWithPreviewTool); itkFactorylessNewMacro(Self); itkCloneMacro(Self); const char *GetName() const override; const char **GetXPM() const override; us::ModuleResource GetIconResource() const override; void Activated() override; itkSetMacro(MitkTempDir, std::string); itkGetConstMacro(MitkTempDir, std::string); itkSetMacro(SubTask, std::string); itkGetConstMacro(SubTask, std::string); itkSetMacro(PythonPath, std::string); itkGetConstMacro(PythonPath, std::string); itkSetMacro(GpuId, unsigned int); itkGetConstMacro(GpuId, unsigned int); itkSetMacro(Fast, bool); itkGetConstMacro(Fast, bool); itkBooleanMacro(Fast); /** * @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 *); protected: TotalSegmentatorTool(); ~TotalSegmentatorTool(); /** * @brief Overridden method from the tool manager to execute the segmentation * Implementation: * 1. Creates temp directory, if not done already. * 2. Parses Label names from map_to_binary.py for using later on. * 3. Calls "run_totalsegmentator" method. * 4. Expects an output image to be saved in the temporary directory by the python process. Loads it as * LabelSetImage and sets to previewImage. * * @param inputAtTimeStep * @param oldSegAtTimeStep * @param previewImage * @param timeStep */ void DoUpdatePreview(const Image* inputAtTimeStep, const Image* oldSegAtTimeStep, LabelSetImage* previewImage, TimeStepType timeStep) override; void UpdatePrepare() override; private: /** * @brief Runs Totalsegmentator python process with desired arguments * */ void run_totalsegmentator(ProcessExecutor*, const std::string&, const std::string&, bool, bool, unsigned int, const std::string&); /** * @brief Applies the m_LabelMapTotal lookup table on the output segmentation LabelSetImage. * */ void MapLabelsToSegmentation(const mitk::LabelSetImage*, mitk::LabelSetImage*, std::map&); /** * @brief Parses map_to_binary.py file to extract label ids and names * and stores as a map for reference in m_LabelMapTotal * */ void ParseLabelMapTotalDefault(); /** * @brief Get the Label Map Path from the virtual environment location * * @return std::string */ std::string GetLabelMapPath(); /** * @brief Agglomerate many individual mask image files into one multi-label LabelSetImage in the * given filePath order. * * @param filePaths * @param dimension * @param geometry * @return LabelSetImage::Pointer */ LabelSetImage::Pointer AgglomerateLabelFiles(std::vector& filePaths, const unsigned int* dimension, mitk::BaseGeometry* geometry); std::string m_MitkTempDir; std::string m_PythonPath; std::string m_SubTask = "total"; unsigned int m_GpuId = 0; std::map m_LabelMapTotal; bool m_Fast = true; const std::string TEMPLATE_FILENAME = "XXXXXX_000_0000.nii.gz"; const std::string DEFAULT_TOTAL_TASK = "total"; + const std::string DEFAULT_TOTAL_TASK_MRI = "total_mr"; const std::unordered_map> SUBTASKS_MAP = { {"body", { "body.nii.gz", "body_trunc.nii.gz", "body_extremities.nii.gz", "skin.nii.gz"}}, {"hip_implant", {"hip_implant.nii.gz"}}, {"cerebral_bleed", {"intracerebral_hemorrhage.nii.gz"}}, {"coronary_arteries", {"coronary_arteries.nii.gz"}}, {"lung_vessels", {"lung_vessels.nii.gz", "lung_trachea_bronchia.nii.gz"}}, {"pleural_pericard_effusion", {"pleural_effusion.nii.gz", "pericardial_effusion.nii.gz"}} }; }; // class } // namespace #endif diff --git a/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp b/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp index 4fa200eeb7..a907216f77 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.cpp @@ -1,257 +1,282 @@ /*============================================================================ 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 "mitkLog.h" #include #include #include #include #include #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==0.7.5"); spExec->Execute(workingDir, "python", args); PipInstall("torch>=2.0.0", workingDir, callback, "ltt"); PipInstall("torchvision>=0.15.0", workingDir, callback, "ltt"); } void QmitkSetupVirtualEnvUtil::InstallPytorch() { this->InstallPytorch(GetPythonPath().toStdString(), &PrintProcessEvent); } bool QmitkSetupVirtualEnvUtil::IsVenvInstalled(const QString &pythonPath) { bool isVenvInstalled = false; QProcess pyProcess; #ifdef _WIN32 const QString PYTHON_EXE = "python.exe"; #else const QString PYTHON_EXE = "python3"; #endif pyProcess.start(pythonPath + QDir::separator() + PYTHON_EXE, QStringList() << "-m" << "venv", QProcess::ReadOnly); //insuffient args to provoke stderr out if (pyProcess.waitForFinished()) { auto venvCaptured = QString(QStringDecoder(QStringDecoder::Utf8)(pyProcess.readAllStandardError())); isVenvInstalled = venvCaptured.startsWith(QString("usage")); // if venv found, shows correct usage instructions } return isVenvInstalled; } 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; } std::pair QmitkSetupVirtualEnvUtil::GetExactPythonPath(const QString &pyEnv) { QString fullPath = pyEnv; bool pythonDoesExist = false; bool isSupportedVersion = false; #ifdef _WIN32 const std::string PYTHON_EXE = "python.exe"; // check if python exist in given folder. pythonDoesExist = QFile::exists(fullPath + QDir::separator() + QString::fromStdString(PYTHON_EXE)); if (!pythonDoesExist && // check if in Scripts already, if not go there !(fullPath.endsWith("Scripts", Qt::CaseInsensitive) || fullPath.endsWith("Scripts/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("Scripts"); pythonDoesExist = QFile::exists(fullPath + QDir::separator() + QString("python.exe")); } #else const std::string PYTHON_EXE = "python3"; pythonDoesExist = QFile::exists(fullPath + QDir::separator() + QString::fromStdString(PYTHON_EXE)); if (!pythonDoesExist && !(fullPath.endsWith("bin", Qt::CaseInsensitive) || fullPath.endsWith("bin/", Qt::CaseInsensitive))) { fullPath += QDir::separator() + QString("bin"); pythonDoesExist = QFile::exists(fullPath + QDir::separator() + QString("python3")); } #endif std::pair pythonPath; if (pythonDoesExist) { std::regex sanitizer(R"(3\.(\d+))"); QProcess pyProcess; pyProcess.start(fullPath + QDir::separator() + QString::fromStdString(PYTHON_EXE), QStringList() << "--version", QProcess::ReadOnly); if (pyProcess.waitForFinished()) { auto pyVersionCaptured = QString(QStringDecoder(QStringDecoder::Utf8)(pyProcess.readAllStandardOutput())).toStdString(); std::smatch match; // Expecting "Python 3.xx.xx" or "Python 3.xx" if (std::regex_search(pyVersionCaptured, match, sanitizer) && !match.empty()) { std::string pyVersionNumber = match[0]; int pySubVersion = std::stoi(match[1]); isSupportedVersion = (pySubVersion > 8) ? (pySubVersion < 13) : false; pythonPath.second = QString::fromStdString(pyVersionNumber); } } } pythonPath.first = pythonDoesExist &&isSupportedVersion ? fullPath : ""; return pythonPath; } + +QString QmitkSetupVirtualEnvUtil::GetPipPackageVersion(const QString &pythonPath, const QString &packageName) +{ +#ifdef _WIN32 + const QString PIP_EXE = "pip.exe"; +#else + const QString PIP_EXE = "pip3"; +#endif + QString version; + QProcess pipProcess; + pipProcess.start(pythonPath + QDir::separator() + PIP_EXE, QStringList() << "show" << packageName, QProcess::ReadOnly); + if (pipProcess.waitForFinished()) + { + while (pipProcess.canReadLine()) + { + QString line = pipProcess.readLine(); + if (line.startsWith("Version")) + { + version = (line.split(":")[1]).trimmed(); + break; + } + } + } + return version; +} diff --git a/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.h b/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.h index c743854bf7..db12623ad7 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.h +++ b/Modules/SegmentationUI/Qmitk/QmitkSetupVirtualEnvUtil.h @@ -1,206 +1,212 @@ /*============================================================================ The Medical Imaging Interaction Toolkit (MITK) Copyright (c) German Cancer Research Center (DKFZ) All rights reserved. Use of this source code is governed by a 3-clause BSD license that can be found in the LICENSE file.s ============================================================================*/ #ifndef QmitkSetupVirtualEnvUtil_h_Included #define QmitkSetupVirtualEnvUtil_h_Included #include "mitkLog.h" #include "mitkProcessExecutor.h" #include #include /** * @brief Abstract Class to Setup a python virtual environment and pip install required packages. * Derive this class for creating installer for the respective tool. */ class MITKSEGMENTATIONUI_EXPORT QmitkSetupVirtualEnvUtil { public: QmitkSetupVirtualEnvUtil(const QString& baseDir); QmitkSetupVirtualEnvUtil(); /** * @brief Implement the method in child class * to setup the virtual environment. */ virtual bool SetupVirtualEnv(const QString& venvName) = 0; /** * @brief Get the Virtual Env Path object. Override this method in the respective * tool installer class. * * @return QString */ virtual QString GetVirtualEnvPath(); /** * @brief Function to Pip install a library package given the location of * pip3 executable. * Any callback function can be passed to process the output. * * @param library * @param workingDir * @param callback * @param command */ void PipInstall(const std::string &library, const std::string &workingDir, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command = "pip3"); /** * @brief Overloaded function to Pip install a library function. * * @param library * @param callback * @param command */ void PipInstall(const std::string &library, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command = "pip3"); /** * @brief Function to execute any python code given a python path. * Any callback function can be passed to process the output. * * @param args * @param pythonPath * @param callback * @param command */ void ExecutePython(const std::string &args, const std::string &pythonPath, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command = "python"); /** * @brief Overloaded function to Execute Python code. * Any callback function can be passed to process the output. * * @param args * @param callback * @param command */ void ExecutePython(const std::string &args, void (*callback)(itk::Object *, const itk::EventObject &, void *), const std::string &command = "python"); /** * @brief Installs pytorch using light-the-torch package, correctly identifying cuda version. * Requires location of pip3 executable. * Any callback function can be passed to process the output. * * @param workingDir * @param callback */ void InstallPytorch(const std::string &workingDir, void (*callback)(itk::Object *, const itk::EventObject &, void *)); /** * @brief Overloaded function to install pytorch using light-the-torch package, correctly * identifying cuda version. */ void InstallPytorch(); /** * @brief Get the Base Dir object * * @return QString& */ QString& GetBaseDir(); /** * @brief Get the System Python Path object * * @return QString& */ QString& GetSystemPythonPath(); /** * @brief Get the Python Path object * * @return QString& */ QString& GetPythonPath(); /** * @brief Get the Pip Path object * * @return QString& */ QString& GetPipPath(); /** * @brief Set the System Python Path object * * @param path */ void SetSystemPythonPath(const QString& path); /** * @brief Set the Python Path object * * @param path */ void SetPythonPath(const QString& path); /** * @brief Set the Pip Path object * * @param path */ void SetPipPath(const QString& path); /** * @brief Set the Virtual Env Path object * * @param path */ void SetVirtualEnvPath(const QString &path); /** * @brief Check if the path provide has python executable or not. * * @param pythonPath * @return true * @return false */ bool IsPythonPath(const QString &pythonPath); /** * @brief Checks if venv module is available for the python. * */ static bool IsVenvInstalled(const QString &pythonPath); + /** + * @brief Returns version of the pip installed package + * + */ + static QString GetPipPackageVersion(const QString &pythonPath, const QString &packageName); + /** * @brief Function can be used as callback to simply print out all the process execution output * parsed out from itk::EventObject. * */ static void PrintProcessEvent(itk::Object *, const itk::EventObject &e, void *); /** * @brief Get the exact Python path and version for any OS from the virtual environment path. * @return A pair of the exact python path and its Python version or empty, if an supported * version of Python could not be found. */ static std::pair GetExactPythonPath(const QString &pyEnv); private: QString m_PythonPath; QString m_PipPath; QString m_BaseDir; QString m_venvPath; QString m_SysPythonPath; }; #endif diff --git a/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorGUIControls.ui b/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorGUIControls.ui index b4a1b287b6..05aae03ebb 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorGUIControls.ui +++ b/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorGUIControls.ui @@ -1,357 +1,357 @@ QmitkTotalSegmentatorToolGUIControls 0 0 699 352 0 0 100 0 100000 100000 QmitkTotalSegmentatorToolWidget 0 0 0 0 0 0 - <html><head/><body><p>Welcome to TotalSegmentator v2 tool in MITK. [Experimental]</p><p>Please note that this is only an interface to TotalSegmentator. MITK does not ship with TotalSegmentator. Make sure to have a working internet connection to install TotalSegmentator. Or, choose an python environment with TotalSegmentator in it before inferencing. - Currently, the tool is pretrained on CT images only. </p><p>Refer to <a href="https://github.com/wasserth/TotalSegmentator"><span style=" text-decoration: underline; color:#0000ff;">https://github.com/wasserth/TotalSegmentator</span></a> to learn everything about the TotalSegmentator.</p><p><br/></p></body></html> + <html><head/><body><p>Welcome to TotalSegmentator v2 tool in MITK. [Experimental]</p><p>Please note that this is only an interface to TotalSegmentator. MITK does not ship with TotalSegmentator. Make sure to have a working internet connection to install TotalSegmentator. Or, choose an python environment with TotalSegmentator in it before inferencing. + The tool is predominantly for CT images only. For MR image choose &quot;total_mr&quot; task specifically. </p><p>Refer to <a href="https://github.com/wasserth/TotalSegmentator"><span style=" text-decoration: underline; color:#0000ff;">https://github.com/wasserth/TotalSegmentator</span></a> to learn everything about the TotalSegmentator.</p><p><br/></p></body></html> Qt::RichText true 0 0 100000 16777215 Install TotalSegmentator 0 0 Install Options true 5 true 0 0 0 0 6 0 0 Use Custom Installation: false 0 0 100000 16777215 Clear Install 0 0 Custom Env. Path: 0 0 System Python (>=3.9): 0 0 Task: 0 0 Advanced true 5 true 0 0 0 0 6 0 0 Fast: true 0 0 GPU Id: 0 0 100000 16777215 Run TotalSegmentator 0 0 true - + ctkComboBox QComboBox
ctkComboBox.h
1 ctkCollapsibleGroupBox QWidget
ctkCollapsibleGroupBox.h
 ctkCheckBox QWidget
ctkCheckBox.h
 diff --git a/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.cpp index bd1b00879f..8505970dd0 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.cpp +++ b/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.cpp @@ -1,511 +1,516 @@ /*============================================================================ 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 "QmitkTotalSegmentatorToolGUI.h" #include "mitkProcessExecutor.h" #include "mitkTotalSegmentatorTool.h" #include #include #include #include #include #include #include MITK_TOOL_GUI_MACRO(MITKSEGMENTATIONUI_EXPORT, QmitkTotalSegmentatorToolGUI, "") QmitkTotalSegmentatorToolGUI::QmitkTotalSegmentatorToolGUI() : QmitkMultiLabelSegWithPreviewToolGUIBase(), m_SuperclassEnableConfirmSegBtnFnc(m_EnableConfirmSegBtnFnc) { // 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) { return !m_FirstPreviewComputation ? m_SuperclassEnableConfirmSegBtnFnc(enabled) : false; }; } void QmitkTotalSegmentatorToolGUI::ConnectNewTool(mitk::SegWithPreviewTool *newTool) { Superclass::ConnectNewTool(newTool); m_FirstPreviewComputation = true; } void QmitkTotalSegmentatorToolGUI::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.pythonEnvComboBox->addItem("Select"); m_Controls.pythonEnvComboBox->setDuplicatesEnabled(false); m_Controls.pythonEnvComboBox->setDisabled(true); m_Controls.previewButton->setDisabled(true); m_Controls.statusLabel->setTextFormat(Qt::RichText); m_Controls.subtaskComboBox->addItems(VALID_TASKS); QString welcomeText; this->SetGPUInfo(); if (m_GpuLoader.GetGPUCount() != 0) { 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.previewButton, SIGNAL(clicked()), this, SLOT(OnPreviewBtnClicked())); connect(m_Controls.clearButton, SIGNAL(clicked()), this, SLOT(OnClearInstall())); connect(m_Controls.installButton, SIGNAL(clicked()), this, SLOT(OnInstallBtnClicked())); connect(m_Controls.overrideBox, SIGNAL(stateChanged(int)), this, SLOT(OnOverrideChecked(int))); connect(m_Controls.pythonEnvComboBox, QOverload::of(&QComboBox::activated), [=](int index) { OnPythonPathChanged(m_Controls.pythonEnvComboBox->itemText(index)); }); connect(m_Controls.sysPythonComboBox, QOverload::of(&QComboBox::activated), [=](int index) { OnSystemPythonChanged(m_Controls.sysPythonComboBox->itemText(index)); }); QString lastSelectedPyEnv = m_Settings.value("TotalSeg/LastCustomPythonPath").toString(); if (!lastSelectedPyEnv.isEmpty() && lastSelectedPyEnv!= "Select") { m_Controls.pythonEnvComboBox->insertItem(0, lastSelectedPyEnv); } m_Controls.fastBox->setChecked(m_Settings.value("TotalSeg/LastFast").toBool()); const QString storageDir = m_Installer.GetVirtualEnvPath(); m_IsInstalled = this->IsTotalSegmentatorInstalled(storageDir); if (m_IsInstalled) { m_PythonPath = QmitkSetupVirtualEnvUtil::GetExactPythonPath(storageDir).first; m_Installer.SetVirtualEnvPath(m_PythonPath); this->EnableAll(m_IsInstalled); welcomeText += " TotalSegmentator is already found installed."; } else { welcomeText += " TotalSegmentator is not installed. Please click on \"Install TotalSegmentator\" above."; } this->WriteStatusMessage(welcomeText); QIcon deleteIcon = QmitkStyleManager::ThemeIcon(QStringLiteral(":/org_mitk_icons/icons/awesome/scalable/actions/edit-delete.svg")); QIcon arrowIcon = QmitkStyleManager::ThemeIcon(QStringLiteral(":/org_mitk_icons/icons/tango/scalable/actions/media-playback-start.svg")); m_Controls.clearButton->setIcon(deleteIcon); m_Controls.previewButton->setIcon(arrowIcon); mainLayout->addLayout(m_Controls.verticalLayout); Superclass::InitializeUI(mainLayout); } void QmitkTotalSegmentatorToolGUI::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)); } } unsigned int QmitkTotalSegmentatorToolGUI::FetchSelectedGPUFromUI() const { QString gpuInfo = m_Controls.gpuComboBox->currentText(); if (m_GpuLoader.GetGPUCount() == 0) { return static_cast(gpuInfo.toInt()); } else { QString gpuId = gpuInfo.split(":", Qt::SkipEmptyParts).first(); return static_cast(gpuId.toInt()); } } void QmitkTotalSegmentatorToolGUI::EnableAll(bool isEnable) { m_Controls.previewButton->setEnabled(isEnable); m_Controls.subtaskComboBox->setEnabled(isEnable); m_Controls.installButton->setEnabled((!isEnable)); } void QmitkTotalSegmentatorToolGUI::OnInstallBtnClicked() { bool isInstalled = false; const auto [path, version] = OnSystemPythonChanged(m_Controls.sysPythonComboBox->currentText()); if (path.isEmpty()) { this->WriteErrorMessage("ERROR: Couldn't find compatible Python."); return; } if (!QmitkSetupVirtualEnvUtil::IsVenvInstalled(path)) { this->WriteErrorMessage("venv module not found for the selected python to create a new virtual " "environment. Please install venv or select another compatibile python"); return; } // check if python 3.13 and ask for confirmation if (version.startsWith("3.13") && QMessageBox::No == QMessageBox::question( nullptr, "Installing TotalSegmentator", QString("WARNING: This is an unsupported version of Python that may not work. " "We recommend using a supported Python version between 3.9 and 3.12.\n\n" "Continue anyway?"), QMessageBox::Yes | QMessageBox::No, QMessageBox::No)) { return; } this->WriteStatusMessage("STATUS: Installing TotalSegmentator..."); m_Installer.SetSystemPythonPath(path); isInstalled = m_Installer.SetupVirtualEnv(m_Installer.VENV_NAME); if (isInstalled) { m_PythonPath = QmitkSetupVirtualEnvUtil::GetExactPythonPath(m_Installer.GetVirtualEnvPath()).first; this->WriteStatusMessage("STATUS: Successfully installed TotalSegmentator."); } else { this->WriteErrorMessage("ERROR: Couldn't install TotalSegmentator."); } this->EnableAll(isInstalled); } void QmitkTotalSegmentatorToolGUI::OnPreviewBtnClicked() { auto tool = this->GetConnectedToolAs(); if (nullptr == tool) { return; } try { m_Controls.previewButton->setEnabled(false); qApp->processEvents(); if (!this->IsTotalSegmentatorInstalled(m_PythonPath)) { throw std::runtime_error(WARNING_TOTALSEG_NOT_FOUND); } bool isFast = m_Controls.fastBox->isChecked(); QString subTask = m_Controls.subtaskComboBox->currentText(); - if (subTask != VALID_TASKS[0]) + if (subTask != VALID_TASKS[0] && subTask != VALID_TASKS[1]) { isFast = true; } tool->SetPythonPath(m_PythonPath.toStdString()); tool->SetGpuId(FetchSelectedGPUFromUI()); tool->SetFast(isFast); tool->SetSubTask(subTask.toStdString()); this->WriteStatusMessage(QString("STATUS: Starting Segmentation task... This might take a while.")); m_FirstPreviewComputation = false; tool->UpdatePreview(); m_Controls.previewButton->setEnabled(true); } catch (const std::exception &e) { std::stringstream errorMsg; errorMsg << "STATUS: Error while processing parameters for TotalSegmentator segmentation. Reason: " << e.what(); this->ShowErrorMessage(errorMsg.str()); this->WriteErrorMessage(QString::fromStdString(errorMsg.str())); m_Controls.previewButton->setEnabled(true); m_FirstPreviewComputation = true; return; } catch (...) { std::string errorMsg = "Unknown error occurred while generation TotalSegmentator segmentation."; this->ShowErrorMessage(errorMsg); m_Controls.previewButton->setEnabled(true); m_FirstPreviewComputation = true; return; } this->SetLabelSetPreview(tool->GetPreviewSegmentation()); this->ActualizePreviewLabelVisibility(); this->WriteStatusMessage("STATUS: Segmentation task finished successfully."); QString pythonPathTextItem = m_Controls.pythonEnvComboBox->currentText(); if (!pythonPathTextItem.isEmpty() && pythonPathTextItem != "Select") // only cache if the prediction ended without errors. { QString lastSelectedPyEnv = m_Settings.value("TotalSeg/LastCustomPythonPath").toString(); if (lastSelectedPyEnv != pythonPathTextItem) { m_Settings.setValue("TotalSeg/LastCustomPythonPath", pythonPathTextItem); } } m_Settings.setValue("TotalSeg/LastFast", m_Controls.fastBox->isChecked()); } void QmitkTotalSegmentatorToolGUI::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 QmitkTotalSegmentatorToolGUI::WriteStatusMessage(const QString &message) { m_Controls.statusLabel->setText(message); m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: white"); qApp->processEvents(); } void QmitkTotalSegmentatorToolGUI::WriteErrorMessage(const QString &message) { m_Controls.statusLabel->setText(message); m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: red"); qApp->processEvents(); } bool QmitkTotalSegmentatorToolGUI::IsTotalSegmentatorInstalled(const QString &pythonPath) { QString fullPath = pythonPath; bool isPythonExists = false, isExists = 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; } isExists = QFile::exists(fullPath + QDir::separator() + QString("TotalSegmentator.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; } isExists = QFile::exists(fullPath + QDir::separator() + QString("TotalSegmentator")) && isPythonExists; #endif + if (isExists && m_Installer.TOTALSEGMENTATOR_VERSION != + QmitkSetupVirtualEnvUtil::GetPipPackageVersion(fullPath, "TotalSegmentator")) + { + isExists = false; + } return isExists; } void QmitkTotalSegmentatorToolGUI::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 irrelevant hidden folders, if any. { m_Controls.pythonEnvComboBox->addItem("(" + envName + "): " + subIt.filePath()); } } } } std::pair QmitkTotalSegmentatorToolGUI::OnSystemPythonChanged(const QString &pyEnv) { std::pair pyPath; if (pyEnv == QString("Select")) { m_Controls.previewButton->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 = QmitkSetupVirtualEnvUtil::GetExactPythonPath(uiPyPath); } return pyPath; } void QmitkTotalSegmentatorToolGUI::OnPythonPathChanged(const QString &pyEnv) { if (pyEnv == QString("Select")) { m_Controls.previewButton->setDisabled(true); QString path = QFileDialog::getExistingDirectory(m_Controls.pythonEnvComboBox->parentWidget(), "Python Path", "dir"); if (!path.isEmpty()) { this->OnPythonPathChanged(path); // recall same function for new path validation bool oldState = m_Controls.pythonEnvComboBox->blockSignals(true); // block signal firing while inserting item m_Controls.pythonEnvComboBox->insertItem(0, path); m_Controls.pythonEnvComboBox->setCurrentIndex(0); m_Controls.pythonEnvComboBox->blockSignals( oldState); // unblock signal firing after inserting item. Remove this after Qt6 migration } } else if (!this->IsTotalSegmentatorInstalled(this->GetPythonPathFromUI(pyEnv))) { this->ShowErrorMessage(WARNING_TOTALSEG_NOT_FOUND); m_Controls.previewButton->setDisabled(true); } else {// Show positive status meeage m_Controls.previewButton->setDisabled(false); QString uiPyPath = this->GetPythonPathFromUI(pyEnv); m_PythonPath = QmitkSetupVirtualEnvUtil::GetExactPythonPath(uiPyPath).first; } } QString QmitkTotalSegmentatorToolGUI::GetPythonPathFromUI(const QString &pyUI) const { QString fullPath = pyUI; if (-1 != fullPath.indexOf(")")) { fullPath = fullPath.mid(fullPath.indexOf(")") + 2); } return fullPath.simplified(); } void QmitkTotalSegmentatorToolGUI::OnOverrideChecked(int state) { bool isEnabled = false; if (state == Qt::Checked) { isEnabled = true; m_Controls.previewButton->setDisabled(true); m_PythonPath.clear(); } else { m_PythonPath.clear(); m_Controls.previewButton->setDisabled(true); if (m_IsInstalled) { const QString pythonPath = m_Installer.GetVirtualEnvPath(); auto pathObject = QmitkSetupVirtualEnvUtil::GetExactPythonPath(pythonPath); m_PythonPath = pathObject.first; this->EnableAll(m_IsInstalled); } } m_Controls.pythonEnvComboBox->setEnabled(isEnabled); } void QmitkTotalSegmentatorToolGUI::OnClearInstall() { QDir folderPath(m_Installer.GetVirtualEnvPath()); if (folderPath.removeRecursively()) { m_Controls.installButton->setEnabled(true); m_IsInstalled = false; if (!m_Controls.overrideBox->isChecked()) { m_Controls.previewButton->setEnabled(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 QmitkTotalSegmentatorToolInstaller::SetupVirtualEnv(const QString& venvName) { if (GetSystemPythonPath().isEmpty()) { return false; } if (!QmitkSetupVirtualEnvUtil::IsVenvInstalled(GetSystemPythonPath())) { 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 QmitkTotalSegmentatorToolInstaller::GetVirtualEnvPath() { return STORAGE_DIR + VENV_NAME; } diff --git a/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.h b/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.h index 10ae565d51..4beb2a2e45 100644 --- a/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.h +++ b/Modules/SegmentationUI/Qmitk/QmitkTotalSegmentatorToolGUI.h @@ -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. ============================================================================*/ #ifndef QmitkTotalSegmentatorToolGUI_h_Included #define QmitkTotalSegmentatorToolGUI_h_Included #include "QmitkMultiLabelSegWithPreviewToolGUIBase.h" #include "QmitkSetupVirtualEnvUtil.h" #include "QmitknnUNetGPU.h" #include "ui_QmitkTotalSegmentatorGUIControls.h" #include #include #include #include #include /** * @brief Installer class for TotalSegmentator Tool. * Class specifies the virtual environment name, install version, packages required to pip install * and implements SetupVirtualEnv method. * */ class QmitkTotalSegmentatorToolInstaller : public QmitkSetupVirtualEnvUtil { public: const QString VENV_NAME = ".totalsegmentator_v2"; - const QString TOTALSEGMENTATOR_VERSION = "2.0.5"; + const QString TOTALSEGMENTATOR_VERSION = "2.2.1"; const std::vector PACKAGES = {QString("Totalsegmentator==") + TOTALSEGMENTATOR_VERSION, QString("setuptools")}; /* just in case */ const QString STORAGE_DIR; inline QmitkTotalSegmentatorToolInstaller( 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::TotalSegmentatorTool. \sa mitk:: */ class MITKSEGMENTATIONUI_EXPORT QmitkTotalSegmentatorToolGUI : public QmitkMultiLabelSegWithPreviewToolGUIBase { Q_OBJECT public: mitkClassMacro(QmitkTotalSegmentatorToolGUI, QmitkMultiLabelSegWithPreviewToolGUIBase); itkFactorylessNewMacro(Self); itkCloneMacro(Self); protected slots: /** * @brief Qt Slot */ void OnPreviewBtnClicked(); /** * @brief Qt Slot */ void OnPythonPathChanged(const QString &); /** * @brief Qt Slot */ std::pair OnSystemPythonChanged(const QString &); /** * @brief Qt Slot */ void OnInstallBtnClicked(); /** * @brief Qt Slot */ void OnOverrideChecked(int); /** * @brief Qt Slot */ void OnClearInstall(); protected: QmitkTotalSegmentatorToolGUI(); ~QmitkTotalSegmentatorToolGUI() = default; void ConnectNewTool(mitk::SegWithPreviewTool *newTool) override; void InitializeUI(QBoxLayout *mainLayout) override; /** * @brief Enable (or Disable) GUI elements. */ void EnableAll(bool); /** * @brief Searches and parses paths of python virtual environments * from predefined lookout locations */ void AutoParsePythonPaths(); /** * @brief Checks if TotalSegmentator command is valid in the selected python virtual environment. * * @return bool */ bool IsTotalSegmentatorInstalled(const QString &); /** * @brief Creates a QMessage object and shows on screen. */ void ShowErrorMessage(const std::string &, QMessageBox::Icon = QMessageBox::Critical); /** * @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 Adds GPU information to the gpu combo box. * In case, there aren't any GPUs available, the combo box will be * rendered editable. */ void SetGPUInfo(); /** * @brief Returns GPU id of the selected GPU from the Combo box. * * @return unsigned int */ unsigned int FetchSelectedGPUFromUI() const; /** * @brief Get the virtual env path from UI combobox removing any * extra special characters. * * @return QString */ QString GetPythonPathFromUI(const QString &) const; /** * @brief For storing values like Python path across sessions. */ QSettings m_Settings; QString m_PythonPath; QmitkGPULoader m_GpuLoader; Ui_QmitkTotalSegmentatorToolGUIControls m_Controls; bool m_FirstPreviewComputation = true; bool m_IsInstalled = false; EnableConfirmSegBtnFunctionType m_SuperclassEnableConfirmSegBtnFnc; const std::string WARNING_TOTALSEG_NOT_FOUND = "TotalSegmentator is not detected in the selected python environment.Please select a valid " "python environment or install TotalSegmentator."; const QStringList VALID_TASKS = { "total", + "total_mr", "cerebral_bleed", "hip_implant", "coronary_arteries", "body", "lung_vessels", "pleural_pericard_effusion" }; QmitkTotalSegmentatorToolInstaller m_Installer; }; #endif