Segmentation Tool (#2332)

* Add skeleton segmentation tool

* Display segmentation

* Restyle

* Working on segmentation painting tool

* Working on segmentation tool

* 2D circle cursor

* Reformat

* Segmentation working.

* Segmentation working.

* Segmentation working.

* More functionality needed for segmentation support.

* Cleanup

* Fix up paint tool for constraints

* Prioritize segmentation over contour tool

* Improve segmentation cursor

* Updates for projects with only images.

* Fix crash with segmentation tools when no images are loaded
Show message and disable tools when images are not present in project.

* Fix default segmentation mode, fix layout.

Author: akenmorris

Libs/Analyze/MeshGenerator.cpp

@@ -98,6 +98,11 @@ MeshHandle MeshGenerator::build_mesh_from_image(ImageType::Pointer image, float
   MeshHandle mesh(new StudioMesh);
 
   try {
+    // only interested in 1's and 0's
+    Image itk_image = Image(image);
+    itk_image.binarize(0, 1);
+    image = itk_image.getITKImage();
+
     // connect to VTK
     vtkSmartPointer<vtkImageImport> vtk_image = vtkSmartPointer<vtkImageImport>::New();
     itk::VTKImageExport<ImageType>::Pointer itk_exporter = itk::VTKImageExport<ImageType>::New();
@@ -177,7 +182,7 @@ MeshHandle MeshGenerator::build_mesh_from_file(std::string filename, float iso_v
       mesh = this->build_mesh_from_image(image, iso_value);
 
       if (mesh->get_poly_data()->GetNumberOfPoints() == 0) {
-        SW_ERROR("Mesh generated for {} with iso-level {} is empty",filename, iso_value);
+        SW_ERROR("Mesh generated for {} with iso-level {} is empty", filename, iso_value);
       }
 
     } catch (itk::ExceptionObject& excep) {

Libs/Analyze/Shape.cpp

@@ -107,6 +107,55 @@ MeshGroup Shape::get_original_meshes(bool wait) {
   return original_meshes_;
 }
 
+//---------------------------------------------------------------------------
+void Shape::recompute_original_surface() {
+  auto seg = get_segmentation();
+  if (!seg) {
+    return;
+  }
+  if (original_meshes_.meshes().empty()) {
+    return;
+  }
+  auto meshes = original_meshes_.meshes();
+  Image copy = *seg;
+  copy.binarize(0, 1);
+  Mesh mesh = copy.toMesh(0.001);
+  MeshHandle mesh_handle = std::make_shared<StudioMesh>();
+  mesh_handle->set_poly_data(mesh.getVTKMesh());
+  original_meshes_.set_mesh(0, mesh_handle);
+}
+
+
+//---------------------------------------------------------------------------
+void Shape::ensure_segmentation()
+{
+  if (get_segmentation()) {
+    return;
+  }
+
+  if (!subject_) {
+    return;
+  }
+
+  if (subject_->get_feature_filenames().empty()) {
+    return;
+  }
+
+  // get image volume
+  auto image_name = subject_->get_feature_filenames().begin()->first;
+  auto base_image = get_image_volume(image_name);
+  if (!base_image) {
+    return;
+  }
+
+  // deep copy
+  Image blank = *base_image;
+  blank.fill(0);
+  segmentation_ = std::make_shared<Image>(blank);
+  // remove extension and add "_seg.nrrd"
+  segmentation_filename_ = StringUtils::removeExtension(image_volume_filename_) + "_seg.nrrd";
+}
+
 //---------------------------------------------------------------------------
 MeshGroup Shape::get_groomed_meshes(bool wait) {
   if (!subject_) {
@@ -593,6 +642,33 @@ std::shared_ptr<Image> Shape::get_image_volume(std::string image_volume_name) {
   return nullptr;
 }
 
+//---------------------------------------------------------------------------
+std::shared_ptr<Image> Shape::get_segmentation() {
+  if (segmentation_) {
+    return segmentation_;
+  }
+  if (!subject_) {
+    return nullptr;
+  }
+  auto filenames = subject_->get_original_filenames();
+  if (filenames.empty()) {
+    return nullptr;
+  }
+  auto filename = filenames[0];
+
+  if (Image::isSupportedType(filename)) {
+    try {
+      std::shared_ptr<Image> image = std::make_shared<Image>(filename);
+      segmentation_ = image;
+      segmentation_filename_ = filename;
+    } catch (std::exception& ex) {
+      SW_ERROR("Unable to open file \"{}\": {}", filename, ex.what());
+    }
+  }
+
+  return segmentation_;
+}
+
 //---------------------------------------------------------------------------
 void Shape::apply_feature_to_points(std::string feature, ImageType::Pointer image) {
   using LinearInterpolatorType = itk::LinearInterpolateImageFunction<ImageType, double>;

Libs/Analyze/Shape.h

@@ -167,6 +167,9 @@ class Shape {
 
   std::shared_ptr<Image> get_image_volume(std::string image_volume_name);
 
+  std::shared_ptr<Image> get_segmentation();
+  std::string get_segmentation_filename() { return segmentation_filename_; }
+
   Eigen::VectorXd get_point_features(std::string feature);
 
   void set_point_features(std::string feature, Eigen::VectorXd values);
@@ -186,6 +189,11 @@ class Shape {
 
   std::vector<std::shared_ptr<MeshWrapper>> get_groomed_mesh_wrappers();
 
+  void recompute_original_surface();
+
+  //! If a segmentation doesn't exist, create a blank canvas
+  void ensure_segmentation();
+
  private:
   void generate_meshes(std::vector<std::string> filenames, MeshGroup& mesh_list, bool save_transform,
                        bool wait = false);
@@ -225,6 +233,9 @@ class Shape {
   std::shared_ptr<Image> image_volume_;
   std::string image_volume_filename_;
 
+  std::shared_ptr<Image> segmentation_;
+  std::string segmentation_filename_;
+
   std::vector<Constraints> constraints_;  // one set for each domain
   int alignment_type_;
 };

Libs/Image/Image.cpp

@@ -1,5 +1,6 @@
 #include "Image.h"
 
+#include <Logging.h>
 #include <itkAntiAliasBinaryImageFilter.h>
 #include <itkBinaryFillholeImageFilter.h>
 #include <itkBinaryThresholdImageFilter.h>
@@ -905,6 +906,11 @@ Image& Image::isolate() {
   return *this;
 }
 
+double Image::get_largest_dimension_size() const {
+  auto our_size = size();
+  return std::max(our_size[0], std::max(our_size[1], our_size[2]));
+}
+
 double Image::get_minimum_spacing() const {
   auto spacing = this->spacing();
   return std::min(spacing[0], std::min(spacing[1], spacing[2]));
@@ -1038,12 +1044,129 @@ Image::PixelType Image::evaluate(Point p) {
   return interpolator_->Evaluate(p);
 }
 
+//-------------------------------------------------------------------------
+void Image::paintSphere(Point p, double radius, PixelType value) {
+  painted_ = true;
+
+  // Convert the center and radius to the index space
+  ImageType::IndexType centerIndex;
+  itk_image_->TransformPhysicalPointToIndex(p, centerIndex);
+
+  // Calculate the bounding box of the sphere in the index space
+  const double radiusSquared = radius * radius;
+  ImageType::RegionType region = itk_image_->GetLargestPossibleRegion();
+
+  ImageType::SizeType size = region.GetSize();
+  ImageType::IndexType start = region.GetIndex();
+
+  for (unsigned int i = 0; i < 3; ++i) {
+    start[i] = std::max<long>(start[i], centerIndex[i] - static_cast<long>(std::ceil(radius)));
+    size[i] = std::min<long>(size[i], centerIndex[i] + static_cast<long>(std::ceil(radius)) + 1);
+  }
+
+  // Iterate only within the bounding box of the sphere
+  ImageType::IndexType index;
+  for (index[2] = start[2]; index[2] < start[2] + size[2]; ++index[2]) {
+    for (index[1] = start[1]; index[1] < start[1] + size[1]; ++index[1]) {
+      for (index[0] = start[0]; index[0] < start[0] + size[0]; ++index[0]) {
+        // Check if the current index is within the radius
+        Point3 q = logicalToPhysical(index);
+        if ((p - q).GetSquaredNorm() <= radiusSquared) {
+          itk_image_->SetPixel(index, value);
+        }
+      }
+    }
+  }
+}
+
+//-------------------------------------------------------------------------
+void Image::paintCircle(Point p, double radius, unsigned int axis, PixelType value) {
+  painted_ = true;
+  // Convert the center to index space
+  ImageType::IndexType centerIndex;
+  itk_image_->TransformPhysicalPointToIndex(p, centerIndex);
+
+  ImageType::SpacingType spacing = itk_image_->GetSpacing();
+
+  // Calculate the radius squared
+  const double radiusSquared = radius * radius;
+
+  // Define the plane's coordinate pairs based on the given VTK axis (dim1, dim2 are the in-plane dimensions)
+  unsigned int dim1, dim2;
+  switch (axis) {
+    case 2:  // XY plane
+      dim1 = 0;
+      dim2 = 1;
+      break;
+    case 1:  // XZ plane
+      dim1 = 0;
+      dim2 = 2;
+      break;
+    case 0:  // YZ plane
+      dim1 = 1;
+      dim2 = 2;
+      break;
+    default:
+      throw std::invalid_argument("Invalid axis, must be 0 (YZ), 1 (XZ), or 2 (XY) according to VTK conventions");
+  }
+
+  // Get the image region
+  ImageType::RegionType region = itk_image_->GetLargestPossibleRegion();
+  ImageType::IndexType startIndex = region.GetIndex();
+  ImageType::SizeType size = region.GetSize();
+  ImageType::IndexType endIndex;
+  endIndex[0] = startIndex[0] + size[0] - 1;
+  endIndex[1] = startIndex[1] + size[1] - 1;
+  endIndex[2] = startIndex[2] + size[2] - 1;
+
+  // Calculate bounding box of the circle in the specified plane, considering spacing
+  startIndex[dim1] =
+      std::max<long>(startIndex[dim1], centerIndex[dim1] - static_cast<long>(std::ceil(radius / spacing[dim1])));
+  startIndex[dim2] =
+      std::max<long>(startIndex[dim2], centerIndex[dim2] - static_cast<long>(std::ceil(radius / spacing[dim2])));
+  endIndex[dim1] =
+      std::min<long>(endIndex[dim1], centerIndex[dim1] + static_cast<long>(std::ceil(radius / spacing[dim1])));
+  endIndex[dim2] =
+      std::min<long>(endIndex[dim2], centerIndex[dim2] + static_cast<long>(std::ceil(radius / spacing[dim2])));
+
+  // Set the orthogonal dimension index to the center slice
+  ImageType::IndexType index;
+  index[(axis == 2) ? 2 : (axis == 1) ? 1 : 0] = centerIndex[(axis == 2) ? 2 : (axis == 1) ? 1 : 0];
+
+  // Iterate only within the bounding box of the circle
+  for (index[dim2] = startIndex[dim2]; index[dim2] <= endIndex[dim2]; ++index[dim2]) {
+    for (index[dim1] = startIndex[dim1]; index[dim1] <= endIndex[dim1]; ++index[dim1]) {
+      // Calculate the 2D distance on the specified plane
+      Point3 q = logicalToPhysical(index);
+      double distanceSquared = std::pow((p[dim1] - q[dim1]), 2) + std::pow((p[dim2] - q[dim2]), 2);
+
+      // If it's within the radius, set the pixel value
+      if (distanceSquared <= radiusSquared) {
+        itk_image_->SetPixel(index, value);
+      }
+    }
+  }
+}
+
+//-------------------------------------------------------------------------
+Image& Image::fill(PixelType value) {
+  itk::ImageRegionIterator<ImageType> imageIterator(itk_image_, itk_image_->GetLargestPossibleRegion());
+  while (!imageIterator.IsAtEnd()) {
+    imageIterator.Set(value);
+    ++imageIterator;
+  }
+  painted_ = true;
+  return *this;
+}
+
+//-------------------------------------------------------------------------
 TransformPtr Image::createCenterOfMassTransform() {
   AffineTransformPtr xform(AffineTransform::New());
   xform->Translate(-(center() - centerOfMass()));  // ITK translations go in a counterintuitive direction
   return xform;
 }
 
+//-------------------------------------------------------------------------
 TransformPtr Image::createRigidRegistrationTransform(const Image& target_dt, float isoValue, unsigned iterations) {
   if (!target_dt.itk_image_) {
     throw std::invalid_argument("Invalid target. Expected distance transform image");
@@ -1067,6 +1190,7 @@ TransformPtr Image::createRigidRegistrationTransform(const Image& target_dt, flo
   return AffineTransform::New();
 }
 
+//-------------------------------------------------------------------------
 std::ostream& operator<<(std::ostream& os, const Image& img) {
   return os << "{\n\tdims: " << img.dims() << ",\n\torigin: " << img.origin() << ",\n\tsize: " << img.size()
             << ",\n\tspacing: " << img.spacing() << "\n}";

Libs/Image/Image.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <StringUtils.h>
 #include <itkImage.h>
 #include <itkImageRegionIterator.h>
 #include <itkLinearInterpolateImageFunction.h>
@@ -145,7 +146,7 @@ class Image {
                         const ImageType::DirectionType direction, InterpolationType interp = NearestNeighbor);
 
   /// applies the given transformation to the image by using resampling filter with reference image
-  Image &applyTransform(const TransformPtr transform, const Image& referenceImage, InterpolationType interp = Linear);
+  Image& applyTransform(const TransformPtr transform, const Image& referenceImage, InterpolationType interp = Linear);
 
   /// extracts/isolates a specific voxel label from a given multi-label volume and outputs the corresponding binary
   /// image
@@ -216,6 +217,9 @@ class Image {
   /// physical dimensions of the image (dims * spacing)
   Point3 size() const { return toPoint(spacing()) * toPoint(dims()); }
 
+  /// largest dimension size
+  double get_largest_dimension_size() const;
+
   /// physical spacing of the image
   Vector spacing() const { return itk_image_->GetSpacing(); }
 
@@ -290,11 +294,33 @@ class Image {
   //! Evaluates the image at a given position
   Image::PixelType evaluate(Point p);
 
+  //! Paints a sphere in the image
+  void paintSphere(Point p, double radius, PixelType value);
+
+  //! Paints a circle in the image
+  void paintCircle(Point p, double radius, unsigned int axis, PixelType value);
+
+  //! Returns if the image has been painted
+  bool isPainted() const { return painted_; }
+
+  //! fill with value
+  Image& fill(PixelType value);
+
   //! Return supported file types
   static std::vector<std::string> getSupportedTypes() {
     return {"nrrd", "nii", "nii.gz", "mhd", "tiff", "jpeg", "jpg", "png", "dcm", "ima"};
   }
 
+  //! Return if the file type is supported
+  static bool isSupportedType(const std::string& filename) {
+    for (const auto& type : Image::getSupportedTypes()) {
+      if (StringUtils::hasSuffix(filename, type)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
  private:
   friend struct SharedCommandData;
   Image()
@@ -318,6 +344,8 @@ class Image {
 
   ImageType::Pointer itk_image_;
 
+  bool painted_ = false;
+
   InterpolatorType::Pointer interpolator_;
 };