AlignTwoPolyDatasUsingTPS.tar

AlignTwoPolyDatasUsingTPS/0000755000076500000240000000000013424703035016771 5ustar  lorensenstaff00000000000000AlignTwoPolyDatasUsingTPS/CMakeLists.txt0000644000076500000240000000225713424703035021537 0ustar  lorensenstaff00000000000000
cmake_minimum_required(VERSION 3.3 FATAL_ERROR)

project(AlignTwoPolyDatasUsingTPS)

find_package(VTK COMPONENTS 
  vtkCommonColor
  vtkCommonCore
  vtkCommonDataModel
  vtkCommonTransforms
  vtkFiltersGeneral
  vtkFiltersModeling
  vtkFiltersSources
  vtkIOGeometry
  vtkIOLegacy
  vtkIOPLY
  vtkIOXML
  vtkInteractionStyle
  vtkRenderingContextOpenGL2
  vtkRenderingCore
  vtkRenderingFreeType
  vtkRenderingGL2PSOpenGL2
  vtkRenderingOpenGL2 QUIET)
if (NOT VTK_FOUND)
  message("Skipping AlignTwoPolyDatasUsingTPS: ${VTK_NOT_FOUND_MESSAGE}")
  return ()
endif()
message (STATUS "VTK_VERSION: ${VTK_VERSION}")
if (VTK_VERSION VERSION_LESS "8.90.0")
  # old system
  include(${VTK_USE_FILE})
  add_executable(AlignTwoPolyDatasUsingTPS MACOSX_BUNDLE AlignTwoPolyDatasUsingTPS.cxx )
  target_link_libraries(AlignTwoPolyDatasUsingTPS PRIVATE ${VTK_LIBRARIES})
else ()
  # include all components
  add_executable(AlignTwoPolyDatasUsingTPS MACOSX_BUNDLE AlignTwoPolyDatasUsingTPS.cxx )
  target_link_libraries(AlignTwoPolyDatasUsingTPS PRIVATE ${VTK_LIBRARIES})
  # vtk_module_autoinit is needed
  vtk_module_autoinit(
    TARGETS AlignTwoPolyDatasUsingTPS
    MODULES ${VTK_LIBRARIES}
    )
endif () 
AlignTwoPolyDatasUsingTPS/build/0000755000076500000240000000000013424703035020070 5ustar  lorensenstaff00000000000000AlignTwoPolyDatasUsingTPS/AlignTwoPolyDatasUsingTPS.cxx0000644000076500000240000002663113424703035024467 0ustar  lorensenstaff00000000000000#include <vtkSmartPointer.h>
#include <vtkThinPlateSplineTransform.h>
#include <vtkIterativeClosestPointTransform.h>

#include <vtkLandmarkTransform.h>
#include <vtkPoints.h>
#include <vtkTransformPolyDataFilter.h>
#include <vtkTransform.h>
#include <vtkHausdorffDistancePointSetFilter.h>
#include <vtkFieldData.h>

#include <vtkOBBTree.h>
#include <vtkBYUReader.h>
#include <vtkOBJReader.h>
#include <vtkPLYReader.h>
#include <vtkPolyDataReader.h>
#include <vtkSTLReader.h>
#include <vtkXMLPolyDataReader.h>
#include <vtkSphereSource.h>

#include <vtkActor.h>
#include <vtkDataSetMapper.h>
#include <vtkProperty.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkRenderer.h>
#include <vtkCamera.h>

#include <vtkNamedColors.h>
#include <vtksys/SystemTools.hxx>

#include <string>
#include <algorithm>
#include <random>
#include <array>

namespace
{
vtkSmartPointer<vtkPolyData> ReadPolyData(const char *fileName);
void AlignBoundingBoxes (vtkPolyData *, vtkPolyData *);
void BestBoundingBox (
  std::string axis,
  vtkPolyData *target,
  vtkPolyData *source,
  vtkPolyData *targetLandmarks,
  vtkPolyData *sourceLandmarks,
  double      &distance,
  vtkPoints   *bestPoints);
}

int main (int argc, char *argv[])
{
  // Vis Pipeline
  auto colors = vtkSmartPointer<vtkNamedColors>::New();

  auto renderer = vtkSmartPointer<vtkRenderer>::New();

  auto renderWindow = vtkSmartPointer<vtkRenderWindow>::New();
  renderWindow->SetSize(640, 480);
  renderWindow->AddRenderer(renderer);

  auto interactor = vtkSmartPointer<vtkRenderWindowInteractor>::New();
  interactor->SetRenderWindow(renderWindow);

  renderer->SetBackground(colors->GetColor3d("sea_green_light").GetData());
  renderer->UseHiddenLineRemovalOn();

  std::cout << "Loading source: " << argv[1] << std::endl;
  auto sourcePolyData = ReadPolyData(argv[1]);

  // Save the source polydata in case the align does not improve
  // segmentation
  auto originalSourcePolyData = vtkSmartPointer<vtkPolyData>::New();
  originalSourcePolyData->DeepCopy(sourcePolyData);

  std::cout << "Loading target: " << argv[2] << std::endl;
  auto targetPolyData = ReadPolyData(argv[2]);

  auto distance =
    vtkSmartPointer<vtkHausdorffDistancePointSetFilter>::New();
  distance->SetInputData(0, targetPolyData);
  distance->SetInputData(1, sourcePolyData);
  distance->Update();

  double distanceBeforeAlign =
    static_cast<vtkPointSet*>(distance->GetOutput(0))->GetFieldData()->GetArray("HausdorffDistance")->GetComponent(0,0);

  // Get initial alignment using oriented bounding boxes
  AlignBoundingBoxes(sourcePolyData, targetPolyData);

  distance->SetInputData(0, targetPolyData);
  distance->SetInputData(1, sourcePolyData);
  distance->Modified();
  distance->Update();
  double distanceAfterAlign =
    static_cast<vtkPointSet*>(distance->GetOutput(0))->GetFieldData()->GetArray("HausdorffDistance")->GetComponent(0,0);

  double bestDistance = std::min(distanceBeforeAlign, distanceAfterAlign);

  if (distanceAfterAlign > distanceBeforeAlign)
  {
    sourcePolyData->DeepCopy(originalSourcePolyData);
  }

  auto icp = vtkSmartPointer<vtkIterativeClosestPointTransform>::New();
  icp->SetSource(sourcePolyData);
  icp->SetTarget(targetPolyData);
  icp->GetLandmarkTransform()->SetModeToRigidBody();
  icp->SetMaximumNumberOfLandmarks(500);
  icp->SetMaximumMeanDistance(.00001);
  icp->SetMaximumNumberOfIterations(500);
  icp->CheckMeanDistanceOn();
  icp->StartByMatchingCentroidsOn();
  icp->Update();

  auto transformICP = vtkSmartPointer<vtkTransformPolyDataFilter>::New();
  transformICP->SetInputData(sourcePolyData);
  transformICP->SetTransform(icp);
  transformICP->Update();

  // Refine the alignment using ThinPLateSplineTransform
  auto tps = vtkSmartPointer<vtkThinPlateSplineTransform>::New();
  tps->SetSourceLandmarks(icp->GetLandmarkTransform()->GetSourceLandmarks());
  tps->SetTargetLandmarks(icp->GetLandmarkTransform()->GetTargetLandmarks());
  tps->SetBasisToR();
  tps->SetSigma(20);
  tps->Update();

//  tps->Print(std::cout);

  auto transformTPS = vtkSmartPointer<vtkTransformPolyDataFilter>::New();
  transformTPS->SetInputData(transformICP->GetOutput());
  transformTPS->SetTransform(tps);
  transformTPS->Update();

  distance->SetInputData(0, targetPolyData);
  distance->SetInputData(1, transformTPS->GetOutput());
  distance->Update();

  double distanceAfterTPS =
    static_cast<vtkPointSet*>(distance->GetOutput(0))->GetFieldData()->GetArray("HausdorffDistance")->GetComponent(0,0);

  if (distanceAfterTPS < bestDistance)
  {
    bestDistance = distanceAfterTPS;
  }

  std::cout << "Distance before, after align, after TPS, min: "
            << distanceBeforeAlign << ", "
            << distanceAfterAlign << ", "
            << distanceAfterTPS << ", "
            << bestDistance << std::endl;

  // Select
  auto sourceMapper = vtkSmartPointer<vtkDataSetMapper>::New();
  if (bestDistance == distanceBeforeAlign)
  {
    sourceMapper->SetInputData(originalSourcePolyData);
    std::cout << "Using original alignment" << std::endl;
  }
  else if (bestDistance == distanceAfterAlign)
  {
    sourceMapper->SetInputData(sourcePolyData);
    std::cout << "Using alignment by OBB" << std::endl;
  }
  else
  {
    sourceMapper->SetInputConnection(transformTPS->GetOutputPort());
    std::cout << "Using alignment by TPS" << std::endl;
  }
  sourceMapper->ScalarVisibilityOff();

  auto sourceActor = vtkSmartPointer<vtkActor>::New();
  sourceActor->SetMapper(sourceMapper);
  sourceActor->GetProperty()->SetOpacity(.6);
  sourceActor->GetProperty()->SetDiffuseColor(
    colors->GetColor3d("White").GetData());
  renderer->AddActor(sourceActor);

  auto targetMapper = vtkSmartPointer<vtkDataSetMapper>::New();
  targetMapper->SetInputData(targetPolyData);
  targetMapper->ScalarVisibilityOff();

  auto targetActor = vtkSmartPointer<vtkActor>::New();
  targetActor->SetMapper(targetMapper);
  targetActor->GetProperty()->SetDiffuseColor(
    colors->GetColor3d("Tomato").GetData());
  renderer->AddActor(targetActor);

  renderWindow->AddRenderer(renderer);

  renderWindow->Render();
  interactor->Start();

  return EXIT_SUCCESS;
}
namespace
{
vtkSmartPointer<vtkPolyData> ReadPolyData(const char *fileName)
{
  vtkSmartPointer<vtkPolyData> polyData;
  std::string extension =
    vtksys::SystemTools::GetFilenameLastExtension(std::string(fileName));

  // Drop the case of the extension
  std::transform(extension.begin(), extension.end(),
                 extension.begin(), ::tolower);

  if (extension == ".ply")
  {
    auto reader = vtkSmartPointer<vtkPLYReader>::New();
    reader->SetFileName (fileName);
    reader->Update();
    polyData = reader->GetOutput();
  }
  else if (extension == ".vtp")
  {
    auto reader = vtkSmartPointer<vtkXMLPolyDataReader>::New();
    reader->SetFileName (fileName);
    reader->Update();
    polyData = reader->GetOutput();
  }
  else if (extension == ".obj")
  {
    auto reader = vtkSmartPointer<vtkOBJReader>::New();
    reader->SetFileName (fileName);
    reader->Update();
    polyData = reader->GetOutput();
  }
  else if (extension == ".stl")
  {
    auto reader = vtkSmartPointer<vtkSTLReader>::New();
    reader->SetFileName (fileName);
    reader->Update();
    polyData = reader->GetOutput();
  }
  else if (extension == ".vtk")
  {
    auto reader = vtkSmartPointer<vtkPolyDataReader>::New();
    reader->SetFileName (fileName);
    reader->Update();
    polyData = reader->GetOutput();
  }
  else if (extension == ".g")
  {
    auto reader = vtkSmartPointer<vtkBYUReader>::New();
    reader->SetGeometryFileName (fileName);
    reader->Update();
    polyData = reader->GetOutput();
  }
  else
  {
    auto source = vtkSmartPointer<vtkSphereSource>::New();
    source->Update();
    polyData = source->GetOutput();
  }
  return polyData;
}

void AlignBoundingBoxes (vtkPolyData *source, vtkPolyData *target)
{
  // Use OBBTree to create an oriented bounding box for target and source
  auto sourceOBBTree = vtkSmartPointer<vtkOBBTree>::New();
  sourceOBBTree->SetDataSet(source);
  sourceOBBTree->SetMaxLevel(1);
  sourceOBBTree->BuildLocator();

  auto targetOBBTree = vtkSmartPointer<vtkOBBTree>::New();
  targetOBBTree->SetDataSet(target);
  targetOBBTree->SetMaxLevel(1);
  targetOBBTree->BuildLocator();

  auto sourceLandmarks = vtkSmartPointer<vtkPolyData>::New();
  sourceOBBTree->GenerateRepresentation(0, sourceLandmarks);

  auto targetLandmarks = vtkSmartPointer<vtkPolyData>::New();
  targetOBBTree->GenerateRepresentation(0, targetLandmarks);

  auto lmTransform = vtkSmartPointer<vtkLandmarkTransform>::New();
  lmTransform->SetModeToSimilarity();
  lmTransform->SetTargetLandmarks(targetLandmarks->GetPoints());
  auto lmTransformPD = vtkSmartPointer<vtkTransformPolyDataFilter>::New();
  double bestDistance = VTK_DOUBLE_MAX;
  auto bestPoints = vtkSmartPointer<vtkPoints>::New();
  BestBoundingBox(
    "X",
    target,
    source,
    targetLandmarks,
    sourceLandmarks,
    bestDistance,
    bestPoints);
  BestBoundingBox(
    "Y",
    target,
    source,
    targetLandmarks,
    sourceLandmarks,
    bestDistance,
    bestPoints);
  BestBoundingBox(
    "Z",
    target,
    source,
    targetLandmarks,
    sourceLandmarks,
    bestDistance,
    bestPoints);

  lmTransform->SetSourceLandmarks(bestPoints);
  lmTransform->Modified();
  
  auto transformPD = vtkSmartPointer<vtkTransformPolyDataFilter>::New();
  transformPD->SetInputData(source);
  transformPD->SetTransform(lmTransform);
  transformPD->Update();

  source->DeepCopy(transformPD->GetOutput());

}
void BestBoundingBox (
  std::string axis,
  vtkPolyData *target,
  vtkPolyData *source,
  vtkPolyData *targetLandmarks,
  vtkPolyData *sourceLandmarks,
  double      &bestDistance,
  vtkPoints   *bestPoints)
{
  auto distance =
    vtkSmartPointer<vtkHausdorffDistancePointSetFilter>::New();
  auto testTransform = vtkSmartPointer<vtkTransform>::New();
  auto testTransformPD = vtkSmartPointer<vtkTransformPolyDataFilter>::New();
  auto lmTransform = vtkSmartPointer<vtkLandmarkTransform>::New();
  auto lmTransformPD = vtkSmartPointer<vtkTransformPolyDataFilter>::New();

  lmTransform->SetModeToSimilarity();
  lmTransform->SetTargetLandmarks(targetLandmarks->GetPoints());

  double sourceCenter[3];
  sourceLandmarks->GetCenter(sourceCenter);

  auto delta = 90.0;
  for (auto i = 0; i < 4; ++i)
  {
    auto angle = delta * i;
    // Rotate about center
    testTransform->Identity();
    testTransform->Translate(
      sourceCenter[0], sourceCenter[1], sourceCenter[2]);
    if (axis == "X")
    {
      testTransform->RotateX(angle);
    }
    else if (axis == "Y")
    {
      testTransform->RotateY(angle);
    }
    else
    {
      testTransform->RotateZ(angle);
    }
    testTransform->Translate(
      -sourceCenter[0], -sourceCenter[1], -sourceCenter[2]);

    testTransformPD->SetTransform(testTransform);
    testTransformPD->SetInputData(sourceLandmarks);
    testTransformPD->Update();

    lmTransform->SetSourceLandmarks(testTransformPD->GetOutput()->GetPoints());
    lmTransform->Modified();

    lmTransformPD->SetInputData(source);
    lmTransformPD->SetTransform(lmTransform);
    lmTransformPD->Update();

    distance->SetInputData(0, target);
    distance->SetInputData(1, lmTransformPD->GetOutput());
    distance->Update();
  
    double testDistance =
      static_cast<vtkPointSet*>(distance->GetOutput(0))->GetFieldData()->GetArray("HausdorffDistance")->GetComponent(0,0);
    if (testDistance < bestDistance)
    {
      bestDistance = testDistance;
      bestPoints->DeepCopy(testTransformPD->GetOutput()->GetPoints());
    }
  }
  return;
}

}