diff --git a/Modules/Segmentation/Interactions/mitkMedSAMTool.cpp b/Modules/Segmentation/Interactions/mitkMedSAMTool.cpp
index 27f254cd40..e16a0b183d 100644
--- a/Modules/Segmentation/Interactions/mitkMedSAMTool.cpp
+++ b/Modules/Segmentation/Interactions/mitkMedSAMTool.cpp
@@ -1,147 +1,147 @@
/*============================================================================
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 "mitkMedSAMTool.h"
#include "mitkToolManager.h"
#include "mitkGeometryData.h"
#include "mitkInteractionPositionEvent.h"
// us
#include <usGetModuleContext.h>
#include <usModule.h>
#include <usModuleContext.h>
#include <usModuleResource.h>
namespace mitk
{
MITK_TOOL_MACRO(MITKSEGMENTATION_EXPORT, MedSAMTool, "MedSAMTool");
}
const char *mitk::MedSAMTool::GetName() const
{
return "MedSAM";
}
void mitk::MedSAMTool::Activated()
{
SegWithPreviewTool::Activated();
this->SetLabelTransferScope(LabelTransferScope::ActiveLabel);
this->SetLabelTransferMode(LabelTransferMode::MapLabel);
}
void mitk::MedSAMTool::Deactivated()
{
SegWithPreviewTool::Deactivated();
GetDataStorage()->Remove(m_BoundingBoxNode);
m_BoundingBoxNode = nullptr;
m_PythonService.reset();
}
void mitk::MedSAMTool::ConnectActionsAndFunctions()
{
CONNECT_FUNCTION("ShiftPrimaryButtonPressed", OnRenderWindowClicked);
CONNECT_FUNCTION("DeletePoint", OnDelete);
}
void mitk::MedSAMTool::OnRenderWindowClicked(StateMachineAction *, InteractionEvent *interactionEvent)
{
if (!this->GetIsReady())
{
return;
}
if ((nullptr == this->GetWorkingPlaneGeometry()) ||
!mitk::Equal(*(interactionEvent->GetSender()->GetCurrentWorldPlaneGeometry()),
*(this->GetWorkingPlaneGeometry())))
{
this->ClearPicks();
this->SetWorkingPlaneGeometry(interactionEvent->GetSender()->GetCurrentWorldPlaneGeometry()->Clone());
auto boundingBox = mitk::GeometryData::New();
boundingBox->SetGeometry(static_cast<mitk::Geometry3D *>(
this->InitializeWithImageGeometry(interactionEvent->GetSender()->GetCurrentWorldPlaneGeometry()->Clone())));
m_BoundingBoxNode = mitk::DataNode::New();
m_BoundingBoxNode->SetData(boundingBox);
m_BoundingBoxNode->SetVisibility(true);
m_BoundingBoxNode->SetName(std::string(this->GetName()) + "_Boundingbox");
m_BoundingBoxNode->SetProperty("layer", mitk::IntProperty::New(99));
m_BoundingBoxNode->AddProperty("Bounding Shape.Handle Size Factor", mitk::DoubleProperty::New(0.02));
m_BoundingBoxNode->SetBoolProperty("pickable", true);
this->GetDataStorage()->Add(m_BoundingBoxNode, this->GetToolManager()->GetWorkingData(0));
this->CreateBoundingShapeInteractor(false);
m_BoundingShapeInteractor->EnableInteraction(true);
m_BoundingShapeInteractor->SetDataNode(m_BoundingBoxNode);
mitk::RenderingManager::GetInstance()->InitializeViews();
}
}
void mitk::MedSAMTool::CreateBoundingShapeInteractor(bool rotationEnabled)
{
if (m_BoundingShapeInteractor.IsNull())
{
m_BoundingShapeInteractor = mitk::BoundingShapeInteractor::New();
m_BoundingShapeInteractor->LoadStateMachine("BoundingShapeInteraction.xml",
us::ModuleRegistry::GetModule("MitkBoundingShape"));
m_BoundingShapeInteractor->SetEventConfig("BoundingShapeMouseConfig.xml",
us::ModuleRegistry::GetModule("MitkBoundingShape"));
}
m_BoundingShapeInteractor->SetRotationEnabled(rotationEnabled);
}
mitk::Geometry3D::Pointer mitk::MedSAMTool::InitializeWithImageGeometry(const mitk::BaseGeometry *geometry) const
{
if (geometry == nullptr)
mitkThrow() << "Geometry is not valid.";
auto boundingGeometry = mitk::Geometry3D::New();
boundingGeometry->SetBounds(geometry->GetBounds());
boundingGeometry->SetImageGeometry(geometry->GetImageGeometry());
boundingGeometry->SetOrigin(geometry->GetOrigin());
boundingGeometry->SetSpacing(geometry->GetSpacing());
boundingGeometry->SetIndexToWorldTransform(geometry->GetIndexToWorldTransform()->Clone());
boundingGeometry->Modified();
return boundingGeometry;
}
bool mitk::MedSAMTool::HasPicks() const
{
return m_BoundingBoxNode.IsNotNull();
}
void mitk::MedSAMTool::OnDelete(StateMachineAction*, InteractionEvent*)
{
this->ClearPicks();
}
void mitk::MedSAMTool::ClearPicks()
{
this->GetDataStorage()->Remove(m_BoundingBoxNode);
m_BoundingBoxNode = nullptr;
this->UpdatePreview();
this->SetWorkingPlaneGeometry(nullptr);
}
-std::stringstream mitk::MedSAMTool::GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry)
+std::string mitk::MedSAMTool::GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry) const
{
auto geometry = m_BoundingBoxNode->GetData()->GetGeometry();
mitk::BoundingBox::ConstPointer boundingBox = geometry->GetBoundingBox();
mitk::Point3D BBmin = boundingBox->GetMinimum();
mitk::Point3D BBmax = boundingBox->GetMaximum();
std::stringstream pointsAndLabels;
pointsAndLabels << "Coordinates\n";
const char SPACE = ' ';
Point2D p2D_min =
this->Get2DIndicesfrom3DWorld(baseGeometry, geometry->GetIndexToWorldTransform()->TransformPoint(BBmin));
Point2D p2D_max =
this->Get2DIndicesfrom3DWorld(baseGeometry, geometry->GetIndexToWorldTransform()->TransformPoint(BBmax));
pointsAndLabels << abs(static_cast<int>(p2D_min[0])) << SPACE << abs(static_cast<int>(p2D_min[1])) << SPACE
<< abs(static_cast<int>(p2D_max[0])) << SPACE << abs(static_cast<int>(p2D_max[1]));
- return pointsAndLabels;
+ return pointsAndLabels.str();
}
diff --git a/Modules/Segmentation/Interactions/mitkMedSAMTool.h b/Modules/Segmentation/Interactions/mitkMedSAMTool.h
index 40dc2170a9..9771eb99c2 100644
--- a/Modules/Segmentation/Interactions/mitkMedSAMTool.h
+++ b/Modules/Segmentation/Interactions/mitkMedSAMTool.h
@@ -1,84 +1,84 @@
/*============================================================================
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 mitkMedSAMTool_h
#define mitkMedSAMTool_h
#include "mitkSegmentAnythingTool.h"
#include "mitkBoundingShapeInteractor.h"
#include <MitkSegmentationExports.h>
#include <mitkRenderingManager.h>
namespace us
{
class ModuleResource;
}
namespace mitk
{
/**
\brief Medical Segment Anything Model interactive 2D tool class.
\ingroup ToolManagerEtAl
\sa mitk::Tool
\sa QmitkInteractiveSegmentation
*/
class MITKSEGMENTATION_EXPORT MedSAMTool : public SegmentAnythingTool
{
public:
mitkClassMacro(MedSAMTool, SegmentAnythingTool);
itkFactorylessNewMacro(Self);
itkCloneMacro(Self);
const char *GetName() const override;
void Activated() override;
void Deactivated() override;
bool HasPicks() const override;
void ClearPicks() override;
void ConnectActionsAndFunctions() override;
- std::stringstream GetPointsAsCSVString(const mitk::BaseGeometry *) override;
+ std::string GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry) const override;
/**
* @brief Adds bounding box in the render window when clicked.
*
*/
- void OnRenderWindowClicked(StateMachineAction *, InteractionEvent *);
+ void OnRenderWindowClicked(StateMachineAction *, InteractionEvent *interactionEvent);
/**
* @brief Deletes bounding box from the render window.
*
*/
void OnDelete(StateMachineAction *, InteractionEvent *);
protected:
MedSAMTool() = default;
~MedSAMTool() = default;
private:
/**
* @brief Initializes the Bounding Shape Interactor object
*
*/
- void CreateBoundingShapeInteractor(bool);
+ void CreateBoundingShapeInteractor(bool rotationEnabled);
/**
* @brief initializes a new bounding shape using the selected image geometry.
*
*/
- mitk::Geometry3D::Pointer InitializeWithImageGeometry(const mitk::BaseGeometry *) const;
+ mitk::Geometry3D::Pointer InitializeWithImageGeometry(const mitk::BaseGeometry *geometry) const;
DataNode::Pointer m_BoundingBoxNode;
BoundingShapeInteractor::Pointer m_BoundingShapeInteractor;
};
}
#endif
diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp
index d88d00d4c8..7d00578b3c 100644
--- a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp
+++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.cpp
@@ -1,269 +1,269 @@
/*============================================================================
The Medical Imaging Interaction Toolkit (MITK)
Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.
Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.
============================================================================*/
#include "mitkSegmentAnythingPythonService.h"
#include "mitkIOUtil.h"
#include <mitkSegmentAnythingProcessExecutor.h>
#include <itksys/SystemTools.hxx>
#include <chrono>
#include <thread>
#include <filesystem>
#include <itkImageFileWriter.h>
#include "mitkImageAccessByItk.h"
#include <mitkLocaleSwitch.h>
using namespace std::chrono_literals;
using sys_clock = std::chrono::system_clock;
namespace mitk
{
const std::string SIGNALCONSTANTS::READY = "READY";
const std::string SIGNALCONSTANTS::KILL = "KILL";
const std::string SIGNALCONSTANTS::OFF = "OFF";
const std::string SIGNALCONSTANTS::CUDA_OUT_OF_MEMORY_ERROR = "CudaOutOfMemoryError";
const std::string SIGNALCONSTANTS::TIMEOUT_ERROR = "TimeOut";
SegmentAnythingPythonService::Status SegmentAnythingPythonService::CurrentStatus =
SegmentAnythingPythonService::Status::OFF;
}
mitk::SegmentAnythingPythonService::SegmentAnythingPythonService(
std::string workingDir, std::string modelType, std::string checkPointPath, unsigned int gpuId, std::string backend)
: m_PythonPath(workingDir),
m_ModelType(modelType),
m_CheckpointPath(checkPointPath),
m_Backend(backend),
m_GpuId(gpuId)
{
this->CreateTempDirs(PARENT_TEMP_DIR_PATTERN);
}
mitk::SegmentAnythingPythonService::~SegmentAnythingPythonService()
{
if (CurrentStatus == Status::READY)
{
this->StopAsyncProcess();
}
CurrentStatus = Status::OFF;
std::filesystem::remove_all(this->GetMitkTempDir());
}
void mitk::SegmentAnythingPythonService::onPythonProcessEvent(itk::Object*, const itk::EventObject &e, void*)
{
std::string testCOUT,testCERR;
const auto *pEvent = dynamic_cast<const mitk::ExternalProcessStdOutEvent *>(&e);
if (pEvent)
{
testCOUT = testCOUT + pEvent->GetOutput();
testCOUT.erase(std::find_if(testCOUT.rbegin(), testCOUT.rend(), [](unsigned char ch) {
return !std::isspace(ch);}).base(), testCOUT.end()); // remove trailing whitespaces, if any
if (SIGNALCONSTANTS::READY == testCOUT)
{
CurrentStatus = Status::READY;
}
if (SIGNALCONSTANTS::KILL == testCOUT)
{
CurrentStatus = Status::KILLED;
}
if (SIGNALCONSTANTS::CUDA_OUT_OF_MEMORY_ERROR == testCOUT)
{
CurrentStatus = Status::CUDAError;
}
MITK_INFO << testCOUT;
}
const auto *pErrEvent = dynamic_cast<const mitk::ExternalProcessStdErrEvent *>(&e);
if (pErrEvent)
{
testCERR = testCERR + pErrEvent->GetOutput();
MITK_ERROR << testCERR;
}
}
void mitk::SegmentAnythingPythonService::StopAsyncProcess()
{
std::stringstream controlStream;
controlStream << SIGNALCONSTANTS::KILL;
- this->WriteControlFile(controlStream);
+ this->WriteControlFile(controlStream.str());
m_DaemonExec->SetStop(true);
m_Future.get();
}
void mitk::SegmentAnythingPythonService::StartAsyncProcess()
{
if (nullptr != m_DaemonExec)
{
this->StopAsyncProcess();
}
if (this->GetMitkTempDir().empty())
{
this->CreateTempDirs(PARENT_TEMP_DIR_PATTERN);
}
std::stringstream controlStream;
controlStream << SIGNALCONSTANTS::READY;
- this->WriteControlFile(controlStream);
+ this->WriteControlFile(controlStream.str());
double timeout = 1;
m_DaemonExec = SegmentAnythingProcessExecutor::New(timeout);
itk::CStyleCommand::Pointer spCommand = itk::CStyleCommand::New();
spCommand->SetCallback(&mitk::SegmentAnythingPythonService::onPythonProcessEvent);
m_DaemonExec->AddObserver(ExternalProcessOutputEvent(), spCommand);
m_Future = std::async(std::launch::async, &mitk::SegmentAnythingPythonService::start_python_daemon, this);
}
-void mitk::SegmentAnythingPythonService::TransferPointsToProcess(std::stringstream &triggerCSV)
+void mitk::SegmentAnythingPythonService::TransferPointsToProcess(const std::string &triggerCSV) const
{
this->CheckStatus();
std::string triggerFilePath = m_InDir + IOUtil::GetDirectorySeparator() + TRIGGER_FILENAME;
std::ofstream csvfile;
csvfile.open(triggerFilePath, std::ofstream::out | std::ofstream::trunc);
- csvfile << triggerCSV.rdbuf();
+ csvfile << triggerCSV;
csvfile.close();
}
-void mitk::SegmentAnythingPythonService::WriteControlFile(std::stringstream &statusStream)
+void mitk::SegmentAnythingPythonService::WriteControlFile(std::string &statusString) const
{
std::string controlFilePath = m_InDir + IOUtil::GetDirectorySeparator() + "control.txt";
std::ofstream controlFile;
controlFile.open(controlFilePath, std::ofstream::out | std::ofstream::trunc);
- controlFile << statusStream.rdbuf();
+ controlFile << statusString;
controlFile.close();
}
-void mitk::SegmentAnythingPythonService::start_python_daemon()
+void mitk::SegmentAnythingPythonService::start_python_daemon() const
{
ProcessExecutor::ArgumentListType args;
std::string command = "python";
args.push_back("-u");
args.push_back(SAM_PYTHON_FILE_NAME);
args.push_back("--input-folder");
args.push_back(m_InDir);
args.push_back("--output-folder");
args.push_back(m_OutDir);
args.push_back("--trigger-file");
args.push_back(TRIGGER_FILENAME);
args.push_back("--model-type");
args.push_back(m_ModelType);
args.push_back("--checkpoint");
args.push_back(m_CheckpointPath);
args.push_back("--backend");
args.push_back(m_Backend);
args.push_back("--device");
if (m_GpuId == -1)
{
args.push_back("cpu");
}
else
{
args.push_back("cuda");
std::string cudaEnv = "CUDA_VISIBLE_DEVICES=" + std::to_string(m_GpuId);
itksys::SystemTools::PutEnv(cudaEnv.c_str());
}
try
{
std::stringstream logStream;
for (const auto &arg : args)
logStream << arg << " ";
logStream << m_PythonPath;
MITK_INFO << logStream.str();
m_DaemonExec->Execute(m_PythonPath, command, args);
}
catch (const mitk::Exception &e)
{
MITK_ERROR << e.GetDescription();
return;
}
MITK_INFO << "Python process ended.";
}
bool mitk::SegmentAnythingPythonService::CheckStatus()
{
switch (CurrentStatus)
{
case mitk::SegmentAnythingPythonService::Status::READY:
return true;
case mitk::SegmentAnythingPythonService::Status::CUDAError:
mitkThrow() << "Error: Cuda Out of Memory. Change your model type in Preferences and Activate Segment Anything tool again.";
case mitk::SegmentAnythingPythonService::Status::KILLED:
mitkThrow() << "Error: Python process is already terminated. Cannot load requested segmentation. Activate Segment Anything tool again.";
default:
return false;
}
}
void mitk::SegmentAnythingPythonService::CreateTempDirs(const std::string &dirPattern)
{
this->SetMitkTempDir(IOUtil::CreateTemporaryDirectory(dirPattern));
m_InDir = IOUtil::CreateTemporaryDirectory("sam-in-XXXXXX", m_MitkTempDir);
m_OutDir = IOUtil::CreateTemporaryDirectory("sam-out-XXXXXX", m_MitkTempDir);
}
-mitk::LabelSetImage::Pointer mitk::SegmentAnythingPythonService::RetrieveImageFromProcess(long timeOut)
+mitk::LabelSetImage::Pointer mitk::SegmentAnythingPythonService::RetrieveImageFromProcess(long timeOut) const
{
std::string outputImagePath = m_OutDir + IOUtil::GetDirectorySeparator() + m_CurrentUId + ".nrrd";
auto start = sys_clock::now();
while (!std::filesystem::exists(outputImagePath))
{
this->CheckStatus();
std::this_thread::sleep_for(100ms);
if (timeOut != -1 && std::chrono::duration_cast<std::chrono::seconds>(sys_clock::now() - start).count() > timeOut)
{
CurrentStatus = Status::OFF;
m_DaemonExec->SetStop(true);
mitkThrow() << SIGNALCONSTANTS::TIMEOUT_ERROR;
}
}
LabelSetImage::Pointer outputBuffer = mitk::IOUtil::Load<LabelSetImage>(outputImagePath);
return outputBuffer;
}
void mitk::SegmentAnythingPythonService::TransferImageToProcess(const Image *inputAtTimeStep, std::string &UId)
{
std::string inputImagePath = m_InDir + IOUtil::GetDirectorySeparator() + UId + ".nrrd";
if (inputAtTimeStep->GetPixelType().GetNumberOfComponents() < 2)
{
AccessByItk_n(inputAtTimeStep, ITKWriter, (inputImagePath));
}
else
{
mitk::IOUtil::Save(inputAtTimeStep, inputImagePath);
}
m_CurrentUId = UId;
}
template <typename TPixel, unsigned int VImageDimension>
-void mitk::SegmentAnythingPythonService::ITKWriter(const itk::Image<TPixel, VImageDimension> *image, std::string& outputFilename)
+void mitk::SegmentAnythingPythonService::ITKWriter(const itk::Image<TPixel, VImageDimension> *image, std::string& outputFilename) const
{
typedef itk::Image<TPixel, VImageDimension> ImageType;
typedef itk::ImageFileWriter<ImageType> WriterType;
typename WriterType::Pointer writer = WriterType::New();
mitk::LocaleSwitch localeSwitch("C");
writer->SetFileName(outputFilename);
writer->SetInput(image);
try
{
writer->Update();
}
catch (const itk::ExceptionObject &error)
{
MITK_ERROR << "Error: " << error << std::endl;
mitkThrow() << "Error: " << error;
}
}
diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h
index 7f15e9eda6..dacd6766fe 100644
--- a/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h
+++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingPythonService.h
@@ -1,152 +1,168 @@
/*============================================================================
The Medical Imaging Interaction Toolkit (MITK)
Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.
Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.
============================================================================*/
#ifndef mitkSegmentAnythingPythonService_h
#define mitkSegmentAnythingPythonService_h
#include <mitkSegmentAnythingProcessExecutor.h>
#include <MitkSegmentationExports.h>
#include <thread>
#include <future>
#include <mitkImage.h>
#include <mitkLabelSetImage.h>
#include <itkImage.h>
namespace mitk
{
/**
* @brief Segment Anything Model Python process handler class.
*
*/
class MITKSEGMENTATION_EXPORT SegmentAnythingPythonService : public itk::Object
{
public:
enum Status
{
READY,
OFF,
KILLED,
CUDAError
};
- SegmentAnythingPythonService(std::string, std::string, std::string, unsigned int, std::string);
+ /**
+ * @brief Construct a new Segment Anything Python Service object. Specify working directory,
+ * ViT model type, checkpoint path, gpu id and backend
+ *
+ * @param workingDir of python process
+ * @param modelType of ViT
+ * @param checkPointPath of specified model type
+ * @param gpuId
+ * @param backend SAM or MedSAM
+ */
+ SegmentAnythingPythonService(std::string workingDir, std::string modelType,
+ std::string checkPointPath, unsigned int gpuId, std::string backend);
+
+ /**
+ * @brief Destroy the Segment Anything Python Service object. Stop the async python process
+ * and deletes temporary directories
+ */
~SegmentAnythingPythonService();
itkSetMacro(MitkTempDir, std::string);
itkGetConstMacro(MitkTempDir, std::string);
/**
* @brief Static function to print out everything from itk::EventObject.
* Used as callback in mitk::ProcessExecutor object.
*
*/
static void onPythonProcessEvent(itk::Object*, const itk::EventObject&, void*);
/**
* @brief Checks CurrentStatus enum variable and returns
* true if daemon is READY (to read files) state, false is OFF state or
* throws exception if daemon is found KILL or Cuda error state.
*
* @return bool
*/
static bool CheckStatus() /*throw(mitk::Exception)*/;
/**
* @brief Creates temp directories and calls start_python_daemon
* function async.
*
*/
void StartAsyncProcess();
/**
* @brief Writes KILL to the control file to stop the daemon process.
*
*/
void StopAsyncProcess();
/**
* @brief Writes image as nifity file with unique id (UId) as file name.
*
*/
- void TransferImageToProcess(const Image*, std::string &UId);
+ void TransferImageToProcess(const Image *inputAtTimeStep, std::string &UId);
/**
* @brief Writes csv stringstream of points to a csv file for
* python daemon to read.
*
*/
- void TransferPointsToProcess(std::stringstream&);
+ void TransferPointsToProcess(const std::string &triggerCSV) const;
/**
* @brief Waits for output nifity file from the daemon to appear and
* reads it as a mitk::Image
*
- * @return Image::Pointer
+ * @return LabelSetImage::Pointer
*/
- LabelSetImage::Pointer RetrieveImageFromProcess(long timeOut= -1);
+ LabelSetImage::Pointer RetrieveImageFromProcess(long timeOut= -1) const;
static Status CurrentStatus;
private:
/**
* @brief Runs SAM python daemon using mitk::ProcessExecutor
*
*/
- void start_python_daemon();
+ void start_python_daemon() const;
/**
* @brief Writes stringstream content into control file.
*
*/
- void WriteControlFile(std::stringstream&);
+ void WriteControlFile(std::string &statusString) const;
/**
* @brief Create a Temp Dirs
*
*/
- void CreateTempDirs(const std::string&);
+ void CreateTempDirs(const std::string &dirPattern);
/**
* @brief ITK-based file writer for dumping inputs into python daemon
*
*/
template <typename TPixel, unsigned int VImageDimension>
- void ITKWriter(const itk::Image<TPixel, VImageDimension> *image, std::string& outputFilename);
+ void ITKWriter(const itk::Image<TPixel, VImageDimension> *image, std::string& outputFilename) const;
std::string m_MitkTempDir;
std::string m_PythonPath;
std::string m_ModelType;
std::string m_CheckpointPath;
std::string m_InDir, m_OutDir;
std::string m_Backend;
std::string m_CurrentUId;
int m_GpuId = 0;
const std::string PARENT_TEMP_DIR_PATTERN = "mitk-sam-XXXXXX";
const std::string TRIGGER_FILENAME = "trigger.csv";
const std::string SAM_PYTHON_FILE_NAME = "run_inference_daemon.py";
std::future<void> m_Future;
SegmentAnythingProcessExecutor::Pointer m_DaemonExec;
};
struct SIGNALCONSTANTS
{
static const std::string READY;
static const std::string KILL;
static const std::string OFF;
static const std::string CUDA_OUT_OF_MEMORY_ERROR;
static const std::string TIMEOUT_ERROR;
};
} // namespace
#endif
diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp
index 73d4a54ab4..d90b14acaa 100644
--- a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp
+++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.cpp
@@ -1,425 +1,425 @@
/*============================================================================
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 <chrono>
#include <thread>
#include <iomanip>
#include "mitkInteractionPositionEvent.h"
#include "mitkPointSetShapeProperty.h"
#include "mitkProperties.h"
#include "mitkToolManager.h"
#include <mitkSegTool2D.h>
// us
#include <usGetModuleContext.h>
#include <usModule.h>
#include <usModuleContext.h>
#include <usModuleResource.h>
#include <itkIntensityWindowingImageFilter.h>
#include "mitkImageAccessByItk.h"
using namespace std::chrono_literals;
namespace mitk
{
MITK_TOOL_MACRO(MITKSEGMENTATION_EXPORT, SegmentAnythingTool, "SegmentAnythingTool");
}
mitk::SegmentAnythingTool::SegmentAnythingTool() : SegWithPreviewTool(true, "PressMoveReleaseAndPointSetting")
{
this->ResetsToEmptyPreviewOn();
this->IsTimePointChangeAwareOff();
this->KeepActiveAfterAcceptOn();
}
const char **mitk::SegmentAnythingTool::GetXPM() const
{
return nullptr;
}
const char *mitk::SegmentAnythingTool::GetName() const
{
return "Segment Anything";
}
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();
m_PointSetNodePositive = mitk::DataNode::New();
m_PointSetNodePositive->SetData(m_PointSetPositive);
m_PointSetNodePositive->SetName(std::string(this->GetName()) + "_PointSetPositive");
m_PointSetNodePositive->SetBoolProperty("helper object", true);
m_PointSetNodePositive->SetColor(0.0, 1.0, 0.0);
m_PointSetNodePositive->SetVisibility(true);
m_PointSetNodePositive->SetProperty("Pointset.2D.shape",
mitk::PointSetShapeProperty::New(mitk::PointSetShapeProperty::CIRCLE));
m_PointSetNodePositive->SetProperty("Pointset.2D.fill shape", mitk::BoolProperty::New(true));
this->GetDataStorage()->Add(m_PointSetNodePositive, 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);
m_PointSetNodeNegative->SetProperty("Pointset.2D.shape",
mitk::PointSetShapeProperty::New(mitk::PointSetShapeProperty::CIRCLE));
m_PointSetNodeNegative->SetProperty("Pointset.2D.fill shape", mitk::BoolProperty::New(true));
this->GetDataStorage()->Add(m_PointSetNodeNegative, this->GetToolManager()->GetWorkingData(0));
this->SetLabelTransferScope(LabelTransferScope::ActiveLabel);
this->SetLabelTransferMode(LabelTransferMode::MapLabel);
}
void mitk::SegmentAnythingTool::Deactivated()
{
this->ClearSeeds();
GetDataStorage()->Remove(m_PointSetNodePositive);
GetDataStorage()->Remove(m_PointSetNodeNegative);
m_PointSetNodePositive = nullptr;
m_PointSetNodeNegative = nullptr;
m_PointSetPositive = nullptr;
m_PointSetNegative = nullptr;
m_PythonService.reset();
Superclass::Deactivated();
}
void mitk::SegmentAnythingTool::ConnectActionsAndFunctions()
{
CONNECT_FUNCTION("ShiftSecondaryButtonPressed", OnAddNegativePoint);
CONNECT_FUNCTION("ShiftPrimaryButtonPressed", OnAddPositivePoint);
CONNECT_FUNCTION("DeletePoint", OnDelete);
}
void mitk::SegmentAnythingTool::InitSAMPythonProcess()
{
if (nullptr != m_PythonService)
{
m_PythonService.reset();
}
this->ClearPicks();
m_PythonService = std::make_unique<mitk::SegmentAnythingPythonService>(
this->GetPythonPath(), this->GetModelType(), this->GetCheckpointPath(), this->GetGpuId(), this->GetBackend());
m_PythonService->StartAsyncProcess();
}
bool mitk::SegmentAnythingTool::IsPythonReady() const
{
return m_PythonService->CheckStatus();
}
void mitk::SegmentAnythingTool::OnAddNegativePoint(StateMachineAction *, InteractionEvent *interactionEvent)
{
if (!this->GetIsReady() || m_PointSetPositive->GetSize() == 0 || nullptr == this->GetWorkingPlaneGeometry() ||
!mitk::Equal(*(interactionEvent->GetSender()->GetCurrentWorldPlaneGeometry()), *(this->GetWorkingPlaneGeometry())))
{
return;
}
if (!this->IsUpdating() && m_PointSetNegative.IsNotNull())
{
const auto positionEvent = dynamic_cast<mitk::InteractionPositionEvent *>(interactionEvent);
if (positionEvent != nullptr)
{
m_PointSetNegative->InsertPoint(m_PointSetCount, positionEvent->GetPositionInWorld());
m_PointSetCount++;
this->UpdatePreview();
}
}
}
void mitk::SegmentAnythingTool::OnAddPositivePoint(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<mitk::InteractionPositionEvent *>(interactionEvent);
if (positionEvent != nullptr)
{
m_PointSetPositive->InsertPoint(m_PointSetCount, positionEvent->GetPositionInWorld());
++m_PointSetCount;
this->UpdatePreview();
}
}
}
void mitk::SegmentAnythingTool::OnDelete(StateMachineAction *, InteractionEvent *)
{
if (!this->IsUpdating() && m_PointSetPositive.IsNotNull() && m_PointSetNegative.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
this->m_PointSetNodePositive->SetData(this->m_PointSetPositive);
}
if (this->m_PointSetNegative.IsNotNull())
{
m_PointSetCount -= m_PointSetNegative->GetSize();
this->m_PointSetNegative = mitk::PointSet::New(); // renew pointset
this->m_PointSetNodeNegative->SetData(this->m_PointSetNegative);
}
}
void mitk::SegmentAnythingTool::ConfirmCleanUp()
{
auto previewImage = this->GetPreviewSegmentation();
for (unsigned int timeStep = 0; timeStep < previewImage->GetTimeSteps(); ++timeStep)
{
this->ResetPreviewContentAtTimeStep(timeStep);
}
this->ClearSeeds();
RenderingManager::GetInstance()->RequestUpdateAll();
}
template <typename TPixel, unsigned int VImageDimension>
void mitk::SegmentAnythingTool::ITKWindowing(const itk::Image<TPixel, VImageDimension> *inputImage,
Image *mitkImage,
ScalarType min,
ScalarType max)
{
typedef itk::Image<TPixel, VImageDimension> ImageType;
typedef itk::IntensityWindowingImageFilter<ImageType, ImageType> IntensityFilterType;
typename IntensityFilterType::Pointer filter = IntensityFilterType::New();
filter->SetInput(inputImage);
filter->SetWindowMinimum(min);
filter->SetWindowMaximum(max);
filter->SetOutputMinimum(min);
filter->SetOutputMaximum(max);
filter->Update();
mitkImage->SetImportVolume((void *)(filter->GetOutput()->GetPixelContainer()->GetBufferPointer()), 0, 0, Image::ManageMemory);
filter->GetOutput()->GetPixelContainer()->ContainerManageMemoryOff();
}
namespace
{
// Checks if the image has valid size across each dimension. The check is
// critical for 2D images since 2D image are not valid in Saggital and Coronal views.
bool IsImageAtTimeStepValid(const mitk::Image *inputAtTimeStep)
{
int total = 0;
total += (inputAtTimeStep->GetDimension(0) > 1);
total += (inputAtTimeStep->GetDimension(1) > 1);
total += (inputAtTimeStep->GetDimension(2) > 1);
return (total > 1);
}
}
void mitk::SegmentAnythingTool::DoUpdatePreview(const Image *inputAtTimeStep,
const Image * /*oldSegAtTimeStep*/,
LabelSetImage *previewImage,
TimeStepType timeStep)
{
if (nullptr != previewImage && ::IsImageAtTimeStepValid(inputAtTimeStep))
{
if (this->HasPicks() && nullptr != m_PythonService)
{
mitk::LevelWindow levelWindow;
this->GetToolManager()->GetReferenceData(0)->GetLevelWindow(levelWindow);
std::string uniquePlaneID = GetHashForCurrentPlane(levelWindow);
m_ProgressCommand->SetProgress(20);
try
{
- std::stringstream csvStream;
+ std::string csvString;
this->EmitSAMStatusMessageEvent("Prompting Segment Anything Model...");
m_ProgressCommand->SetProgress(50);
if (inputAtTimeStep->GetPixelType().GetNumberOfComponents() < 2)
{
auto filteredImage = mitk::Image::New();
filteredImage->Initialize(inputAtTimeStep);
AccessByItk_n(inputAtTimeStep, ITKWindowing, // apply level window filter
(filteredImage, levelWindow.GetLowerWindowBound(), levelWindow.GetUpperWindowBound()));
m_PythonService->TransferImageToProcess(filteredImage, uniquePlaneID);
- csvStream = this->GetPointsAsCSVString(filteredImage->GetGeometry());
+ csvString = this->GetPointsAsCSVString(filteredImage->GetGeometry());
}
else
{
m_PythonService->TransferImageToProcess(inputAtTimeStep, uniquePlaneID);
- csvStream = this->GetPointsAsCSVString(inputAtTimeStep->GetGeometry());
+ csvString = this->GetPointsAsCSVString(inputAtTimeStep->GetGeometry());
}
m_ProgressCommand->SetProgress(100);
- m_PythonService->TransferPointsToProcess(csvStream);
+ m_PythonService->TransferPointsToProcess(csvString);
m_ProgressCommand->SetProgress(150);
std::this_thread::sleep_for(100ms);
mitk::LabelSetImage::Pointer outputBuffer = m_PythonService->RetrieveImageFromProcess(this->GetTimeOutLimit());
m_ProgressCommand->SetProgress(180);
mitk::SegTool2D::WriteSliceToVolume(previewImage, this->GetWorkingPlaneGeometry(), outputBuffer.GetPointer(), timeStep, false);
this->SetSelectedLabels({MASK_VALUE});
this->EmitSAMStatusMessageEvent("Successfully generated segmentation.");
}
catch (const mitk::Exception &e)
{
this->ClearPicks();
this->EmitSAMStatusMessageEvent(e.GetDescription());
mitkThrow() << e.GetDescription();
}
}
else if (nullptr != this->GetWorkingPlaneGeometry())
{
this->ResetPreviewContentAtTimeStep(timeStep);
RenderingManager::GetInstance()->ForceImmediateUpdateAll();
}
}
}
-std::string mitk::SegmentAnythingTool::GetHashForCurrentPlane(const mitk::LevelWindow &levelWindow)
+std::string mitk::SegmentAnythingTool::GetHashForCurrentPlane(const mitk::LevelWindow &levelWindow) const
{
mitk::Vector3D normal = this->GetWorkingPlaneGeometry()->GetNormal();
mitk::Point3D center = this->GetWorkingPlaneGeometry()->GetCenter();
std::stringstream hashstream;
hashstream << std::setprecision(3) << std::fixed << normal[0]; //Consider only 3 digits after decimal
hashstream << std::setprecision(3) << std::fixed << normal[1];
hashstream << std::setprecision(3) << std::fixed << normal[2];
hashstream << std::setprecision(3) << std::fixed << center[0];
hashstream << std::setprecision(3) << std::fixed << center[1];
hashstream << std::setprecision(3) << std::fixed << center[2];
hashstream << levelWindow.GetLowerWindowBound();
hashstream << levelWindow.GetUpperWindowBound();
size_t hashVal = std::hash<std::string>{}(hashstream.str());
return std::to_string(hashVal);
}
-std::stringstream mitk::SegmentAnythingTool::GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry)
+std::string mitk::SegmentAnythingTool::GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry) const
{
MITK_INFO << "No.of points: " << m_PointSetPositive->GetSize();
std::stringstream pointsAndLabels;
pointsAndLabels << "Point,Label\n";
mitk::PointSet::PointsConstIterator pointSetItPos = m_PointSetPositive->Begin();
mitk::PointSet::PointsConstIterator pointSetItNeg = m_PointSetNegative->Begin();
const char SPACE = ' ';
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 << static_cast<int>(p2D[0]) << SPACE << static_cast<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 << static_cast<int>(p2D[0]) << SPACE << static_cast<int>(p2D[1]) << ",0" << std::endl;
}
++pointSetItNeg;
}
}
- return pointsAndLabels;
+ return pointsAndLabels.str();
}
std::vector<std::pair<mitk::Point2D, std::string>> mitk::SegmentAnythingTool::GetPointsAsVector(
- const mitk::BaseGeometry *baseGeometry)
+ const mitk::BaseGeometry *baseGeometry) const
{
std::vector<std::pair<mitk::Point2D, std::string>> clickVec;
clickVec.reserve(m_PointSetPositive->GetSize() + m_PointSetNegative->GetSize());
mitk::PointSet::PointsConstIterator pointSetItPos = m_PointSetPositive->Begin();
mitk::PointSet::PointsConstIterator 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,
const mitk::Point3D &point3d)
{
mitk::Point3D index3D;
baseGeometry->WorldToIndex(point3d, index3D);
MITK_INFO << index3D[0] << " " << index3D[1] << " " << index3D[2]; // remove
Point2D point2D;
point2D.SetElement(0, index3D[0]);
point2D.SetElement(1, index3D[1]);
return point2D;
}
void mitk::SegmentAnythingTool::EmitSAMStatusMessageEvent(const std::string& status)
{
SAMStatusMessageEvent.Send(status);
}
diff --git a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h
index 99c2e20f65..652f414dfc 100644
--- a/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h
+++ b/Modules/Segmentation/Interactions/mitkSegmentAnythingTool.h
@@ -1,221 +1,221 @@
/*============================================================================
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 "mitkSegmentAnythingPythonService.h"
#include <MitkSegmentationExports.h>
#include <itkImage.h>
#include <mitkLevelWindow.h>
namespace us
{
class ModuleResource;
}
namespace mitk
{
/**
\brief Segment Anything Model interactive 2D tool class.
\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;
/**
* @brief Clears all picks and updates the preview.
*/
virtual void ClearPicks();
/**
* @brief Checks if any point exists in the either
* of the pointsets
*
* @return bool
*/
virtual 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(Backend, std::string);
itkGetConstMacro(Backend, std::string);
itkSetMacro(GpuId, int);
itkGetConstMacro(GpuId, int);
itkSetMacro(TimeOutLimit, long);
itkGetConstMacro(TimeOutLimit, long);
itkSetMacro(IsReady, bool);
itkGetConstMacro(IsReady, bool);
itkBooleanMacro(IsReady);
/**
* @brief Initializes python service and
* starts async python daemon of SegmentAnything model.
*
*/
void InitSAMPythonProcess();
/**
* @brief Checks if Python daemon is ready to accept inputs.
*
* @return bool
*/
bool IsPythonReady() const;
Message1<const std::string&> SAMStatusMessageEvent;
protected:
SegmentAnythingTool();
~SegmentAnythingTool() = default;
void ConnectActionsAndFunctions() override;
/*
* @brief Add positive seed point action of StateMachine pattern
*/
virtual void OnAddPositivePoint(StateMachineAction*, InteractionEvent *interactionEvent);
/*
* @brief Add negative seed point action of StateMachine pattern
*/
virtual void OnAddNegativePoint(StateMachineAction*, InteractionEvent *interactionEvent);
/*
* @brief Delete action of StateMachine pattern. The function deletes positive or negative points in
the reverse order of creation. This is done by finding & deleting the Point having the highest
PointIdentifier value from either of the PointSets m_PointSetPositive & m_PointSetNegative.
*/
virtual void OnDelete(StateMachineAction*, InteractionEvent*);
/*
* @brief Clear all seed points and call UpdatePreview to reset the segmentation Preview
*/
void ClearSeeds();
/**
* @brief Overriden method from the tool manager to execute the segmentation
* Implementation:
* 1. Creates Hash for input image from current plane geometry.
* 2. Transfers image pointer to python service along with the hash code.
* 3. Creates seed points as CSV string & transfers to python service
* 3. Retrieves resulting segmentation Image pointer from python service 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;
/**
* @brief Get the Points from positive and negative pointsets as std::vector.
*
* @return std::vector<std::pair<mitk::Point2D, std::string>>
*/
- std::vector<std::pair<mitk::Point2D, std::string>> GetPointsAsVector(const mitk::BaseGeometry*);
+ std::vector<std::pair<mitk::Point2D, std::string>> GetPointsAsVector(const mitk::BaseGeometry *baseGeometry) const;
/**
* @brief Get the Points from positive and negative pointsets as csv string.
*
* @param baseGeometry
- * @return std::stringstream
+ * @return std::string
*/
- virtual std::stringstream GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry);
+ virtual std::string GetPointsAsCSVString(const mitk::BaseGeometry *baseGeometry) const;
/**
* @brief Get the Hash For Current Plane from current working plane geometry.
*
* @return std::string
*/
- std::string GetHashForCurrentPlane(const mitk::LevelWindow&);
+ std::string GetHashForCurrentPlane(const mitk::LevelWindow &levelWindow) const;
/**
* @brief Emits message to connected Listnerers.
*
*/
- void EmitSAMStatusMessageEvent(const std::string&);
+ void EmitSAMStatusMessageEvent(const std::string &status);
/**
* @brief Cleans up segmentation preview and clears all seeds.
*
*/
void ConfirmCleanUp() override;
/**
* @brief Applies ITK IntensityWindowing Filter to input image;
*
*/
template <typename TPixel, unsigned int VImageDimension>
void ITKWindowing(const itk::Image<TPixel, VImageDimension>*, mitk::Image*, ScalarType, ScalarType);
/**
* @brief Convert 3D world coordinates to 2D indices.
*
- * @param baseGeometry
- * @param mitk::Point3D
+ * @param baseGeometry Base Geometry of image
+ * @param mitk::Point3D 3D world coordinates
* @return mitk::Point2D
*/
- static mitk::Point2D Get2DIndicesfrom3DWorld(const mitk::BaseGeometry*, const mitk::Point3D&);
+ static mitk::Point2D Get2DIndicesfrom3DWorld(const mitk::BaseGeometry *baseGeometry, const mitk::Point3D &point3d);
std::unique_ptr<SegmentAnythingPythonService> m_PythonService;
private:
std::string m_MitkTempDir;
std::string m_PythonPath;
std::string m_ModelType;
std::string m_CheckpointPath;
std::string m_Backend;
int m_GpuId = 0;
PointSet::Pointer m_PointSetPositive;
PointSet::Pointer m_PointSetNegative;
DataNode::Pointer m_PointSetNodePositive;
DataNode::Pointer m_PointSetNodeNegative;
bool m_IsGenerateEmbeddings = true;
bool m_IsReady = false;
int m_PointSetCount = 0;
long m_TimeOutLimit = -1;
const Label::PixelType MASK_VALUE = 1;
};
} // namespace
#endif
diff --git a/Modules/SegmentationUI/Qmitk/QmitkMedSAMToolGUI.cpp b/Modules/SegmentationUI/Qmitk/QmitkMedSAMToolGUI.cpp
index de3bf97038..ad8a1f4e9a 100644
--- a/Modules/SegmentationUI/Qmitk/QmitkMedSAMToolGUI.cpp
+++ b/Modules/SegmentationUI/Qmitk/QmitkMedSAMToolGUI.cpp
@@ -1,284 +1,284 @@
/*============================================================================
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 "QmitkMedSAMToolGUI.h"
#include <QApplication>
#include <QmitkStyleManager.h>
#include <mitkCoreServices.h>
#include <mitkIPreferencesService.h>
#include <mitkMedSAMTool.h>
MITK_TOOL_GUI_MACRO(MITKSEGMENTATIONUI_EXPORT, QmitkMedSAMToolGUI, "")
namespace
{
mitk::IPreferences *GetPreferences()
{
auto *preferencesService = mitk::CoreServices::GetPreferencesService();
return preferencesService->GetSystemPreferences()->Node("org.mitk.views.segmentation");
}
}
QmitkMedSAMToolGUI::QmitkMedSAMToolGUI() : QmitkSegWithPreviewToolGUIBase(true)
{
m_EnableConfirmSegBtnFnc = [this](bool enabled)
{
bool result = false;
auto tool = this->GetConnectedToolAs<mitk::MedSAMTool>();
if (nullptr != tool)
{
result = enabled && tool->HasPicks();
}
return result;
};
m_Preferences = GetPreferences();
m_Preferences->OnPropertyChanged +=
mitk::MessageDelegate1<QmitkMedSAMToolGUI, const mitk::IPreferences::ChangeEvent &>(
this, &QmitkMedSAMToolGUI::OnPreferenceChangedEvent);
}
QmitkMedSAMToolGUI::~QmitkMedSAMToolGUI()
{
- auto tool = this->GetConnectedToolAs<mitk::SegmentAnythingTool>(); // check -ashis
+ auto tool = this->GetConnectedToolAs<mitk::SegmentAnythingTool>();
if (nullptr != tool)
{
tool->SAMStatusMessageEvent -=
mitk::MessageDelegate1<QmitkMedSAMToolGUI, const std::string &>(this, &QmitkMedSAMToolGUI::StatusMessageListener);
}
}
void QmitkMedSAMToolGUI::EnableAll(bool isEnable)
{
m_Controls.activateButton->setEnabled(isEnable);
}
void QmitkMedSAMToolGUI::WriteStatusMessage(const QString &message)
{
m_Controls.statusLabel->setText(message);
m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: white");
qApp->processEvents();
}
void QmitkMedSAMToolGUI::WriteErrorMessage(const QString &message)
{
m_Controls.statusLabel->setText(message);
m_Controls.statusLabel->setStyleSheet("font-weight: bold; color: red");
qApp->processEvents();
}
void QmitkMedSAMToolGUI::ShowProgressBar(bool enabled)
{
m_Controls.samProgressBar->setEnabled(enabled);
m_Controls.samProgressBar->setVisible(enabled);
}
void QmitkMedSAMToolGUI::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 QmitkMedSAMToolGUI::InitializeUI(QBoxLayout *mainLayout)
{
m_Controls.setupUi(this);
m_Controls.statusLabel->setTextFormat(Qt::RichText);
QString welcomeText;
if (m_GpuLoader.GetGPUCount() != 0)
{
welcomeText = "<b>STATUS: </b><i>Welcome to MedSAM Anything tool. You're in luck: " +
QString::number(m_GpuLoader.GetGPUCount()) + " GPU(s) were detected.</i>";
}
else
{
welcomeText = "<b>STATUS: </b><i>Welcome to MedSAM Anything tool. Sorry, " +
QString::number(m_GpuLoader.GetGPUCount()) + " GPUs were detected.</i>";
}
connect(m_Controls.previewButton, SIGNAL(clicked()), this, SLOT(OnPreviewBtnClicked()));
connect(m_Controls.activateButton, SIGNAL(clicked()), this, SLOT(OnActivateBtnClicked()));
connect(m_Controls.resetButton, SIGNAL(clicked()), this, SLOT(OnResetPicksClicked()));
QIcon arrowIcon = QmitkStyleManager::ThemeIcon(
QStringLiteral(":/org_mitk_icons/icons/tango/scalable/actions/media-playback-start.svg"));
m_Controls.activateButton->setIcon(arrowIcon);
bool isInstalled = this->ValidatePrefences();
if (isInstalled)
{
QString modelType = QString::fromStdString(m_Preferences->Get("sam modeltype", ""));
welcomeText += " MedSAM is already found installed.";
}
else
{
welcomeText += " MedSAM tool is not configured correctly. Please go to Preferences (Cntl+P) > Segment Anything to "
"configure and/or install SAM & MedSAM.";
}
this->EnableAll(isInstalled);
this->WriteStatusMessage(welcomeText);
this->ShowProgressBar(false);
m_Controls.samProgressBar->setMaximum(0);
mainLayout->addLayout(m_Controls.verticalLayout);
Superclass::InitializeUI(mainLayout);
}
bool QmitkMedSAMToolGUI::ValidatePrefences()
{
const QString storageDir = QString::fromStdString(m_Preferences->Get("sam python path", ""));
bool isInstalled = QmitkSegmentAnythingToolGUI::IsSAMInstalled(storageDir);
std::string modelType = m_Preferences->Get("sam modeltype", "");
std::string path = m_Preferences->Get("sam parent path", "");
return (isInstalled && !modelType.empty() && !path.empty());
}
void QmitkMedSAMToolGUI::StatusMessageListener(const std::string &message)
{
if (message.rfind("Error", 0) == 0)
{
this->EnableAll(true);
this->WriteErrorMessage(QString::fromStdString(message));
}
else if (message == "TimeOut")
{ // trying to re init the daemon
this->WriteErrorMessage(QString("<b>STATUS: </b><i>Sorry, operation timed out. Reactivating MedSAM tool...</i>"));
if (this->ActivateSAMDaemon())
{
this->WriteStatusMessage(QString("<b>STATUS: </b><i>MedSAM tool re-initialized.</i>"));
}
else
{
this->WriteErrorMessage(QString("<b>STATUS: </b><i>Couldn't init tool backend.</i>"));
this->EnableAll(true);
}
}
else
{
this->WriteStatusMessage(QString::fromStdString(message));
}
}
bool QmitkMedSAMToolGUI::ActivateSAMDaemon()
{
auto tool = this->GetConnectedToolAs<mitk::MedSAMTool>();
if (nullptr == tool)
{
return false;
}
this->ShowProgressBar(true);
qApp->processEvents();
try
{
tool->InitSAMPythonProcess();
while (!tool->IsPythonReady())
{
qApp->processEvents();
}
tool->IsReadyOn();
}
catch (...)
{
tool->IsReadyOff();
}
this->ShowProgressBar(false);
return tool->GetIsReady();
}
void QmitkMedSAMToolGUI::OnActivateBtnClicked()
{
auto tool = this->GetConnectedToolAs<mitk::MedSAMTool>();
if (nullptr == tool)
{
return;
}
try
{
this->EnableAll(false);
qApp->processEvents();
QString pythonPath = QString::fromStdString(m_Preferences->Get("sam python path", ""));
if (!QmitkSegmentAnythingToolGUI::IsSAMInstalled(pythonPath))
{
throw std::runtime_error(WARNING_SAM_NOT_FOUND);
}
tool->SetPythonPath(pythonPath.toStdString());
tool->SetGpuId(m_Preferences->GetInt("sam gpuid", -1));
tool->SetModelType("vit_b"); // MedSAM only works with vit_b
tool->SetTimeOutLimit(m_Preferences->GetInt("sam timeout", 300));
tool->SetCheckpointPath(m_Preferences->Get("sam parent path", ""));
tool->SetBackend("MedSAM");
this->WriteStatusMessage(QString("<b>STATUS: </b><i>Initializing MedSAM...</i>"));
tool->SAMStatusMessageEvent +=
mitk::MessageDelegate1<QmitkMedSAMToolGUI, const std::string &>(this, &QmitkMedSAMToolGUI::StatusMessageListener);
if (this->ActivateSAMDaemon())
{
this->WriteStatusMessage(QString("<b>STATUS: </b><i>MedSAM tool initialized.</i>"));
}
else
{
this->WriteErrorMessage(QString("<b>STATUS: </b><i>Couldn't init tool backend.</i>"));
this->EnableAll(true);
}
}
catch (const std::exception &e)
{
std::stringstream errorMsg;
errorMsg << "<b>STATUS: </b>Error while processing parameters for MedSAM segmentation. Reason: " << e.what();
this->ShowErrorMessage(errorMsg.str());
this->WriteErrorMessage(QString::fromStdString(errorMsg.str()));
this->EnableAll(true);
return;
}
catch (...)
{
std::string errorMsg = "Unkown error occured while generation MedSAM segmentation.";
this->ShowErrorMessage(errorMsg);
this->EnableAll(true);
return;
}
}
void QmitkMedSAMToolGUI::OnPreviewBtnClicked()
{
auto tool = this->GetConnectedToolAs<mitk::MedSAMTool>();
if (nullptr != tool)
{
tool->UpdatePreview();
}
}
void QmitkMedSAMToolGUI::OnResetPicksClicked()
{
auto tool = this->GetConnectedToolAs<mitk::MedSAMTool>();
if (nullptr != tool)
{
tool->ClearPicks();
}
}
void QmitkMedSAMToolGUI::OnPreferenceChangedEvent(const mitk::IPreferences::ChangeEvent &event)
{
const std::string property = event.GetProperty();
const std::string modelType = "modeltype";
if (property.compare(property.size() - modelType.size(), modelType.size(), modelType) == 0)
return; // Model type change ignored.
this->EnableAll(true);
this->WriteStatusMessage("A Preference change was detected. Please initialize the tool again.");
auto tool = this->GetConnectedToolAs<mitk::MedSAMTool>();
if (nullptr != tool)
{
tool->IsReadyOff();
}
}
diff --git a/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.h b/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.h
index 1827bae624..faf7a3de86 100644
--- a/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.h
+++ b/Plugins/org.mitk.gui.qt.segmentation/src/QmitkSegmentAnythingPreferencePage.h
@@ -1,108 +1,108 @@
/*============================================================================
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 QmitkSegmentAnythingPreferencePage_h
#define QmitkSegmentAnythingPreferencePage_h
#include <berryIQtPreferencePage.h>
#include <QmitknnUNetGPU.h>
#include <QString>
#include <QmitkSetupVirtualEnvUtil.h>
#include <QStandardPaths>
#include <ui_QmitkSegmentAnythingPreferencePage.h>
#include <QDir>
class QWidget;
namespace Ui
{
class QmitkSegmentAnythingPreferencePage;
}
class QmitkSAMInstaller : public QmitkSetupVirtualEnvUtil
{
public:
const QString VENV_NAME = ".sam";
const QString SAM_VERSION = "1.0"; // currently, unused
- const std::vector<QString> PACKAGES = {QString("git+https://github.com/ASHISRAVINDRAN/sam-mitk.git")};
+ const std::vector<QString> PACKAGES = {QString("git+https://github.com/MIC-DKFZ/agent-sam.git@v0.2")};
const QString STORAGE_DIR;
inline QmitkSAMInstaller(
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;
};
class QmitkSegmentAnythingPreferencePage : public QObject, public berry::IQtPreferencePage
{
Q_OBJECT
Q_INTERFACES(berry::IPreferencePage)
public:
QmitkSegmentAnythingPreferencePage();
~QmitkSegmentAnythingPreferencePage() override;
void Init(berry::IWorkbench::Pointer workbench) override;
void CreateQtControl(QWidget* parent) override;
QWidget* GetQtControl() const override;
bool PerformOk() override;
void PerformCancel() override;
void Update() override;
private slots:
void OnInstallBtnClicked();
void OnClearInstall();
std::pair<QString, QString> OnSystemPythonChanged(const QString &);
protected:
/**
* @brief Searches and parses paths of python virtual enviroments
* from predefined lookout locations
*/
void AutoParsePythonPaths();
/**
* @brief Get the virtual env path from UI combobox removing any
* extra special characters.
*
* @return QString
*/
QString GetPythonPathFromUI(const QString &) const;
/**
* @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 Returns GPU id of the selected GPU from the Combo box.
* @return int
*/
int FetchSelectedGPUFromUI() const;
void WriteStatusMessage(const QString &);
void WriteErrorMessage(const QString &);
private:
Ui::QmitkSegmentAnythingPreferencePage* m_Ui;
QmitkSAMInstaller m_Installer;
QWidget* m_Control;
QmitkGPULoader m_GpuLoader;
QString m_PythonPath;
const QString CPU_ID = "cpu";
const QStringList VALID_MODELS = {"vit_b", "vit_l", "vit_h"};
};
#endif