Curvatures

#include <vtkActor.h>
#include <vtkColorSeries.h>
#include <vtkColorTransferFunction.h>
#include <vtkCurvatures.h>
#include <vtkDiscretizableColorTransferFunction.h>
#include <vtkFeatureEdges.h>
#include <vtkLookupTable.h>
#include <vtkNamedColors.h>
#include <vtkNew.h>
#include <vtkPointData.h>
#include <vtkPolyDataMapper.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkRenderer.h>
#include <vtkScalarBarActor.h>
#include <vtkScalarBarRepresentation.h>
#include <vtkScalarBarWidget.h>
#include <vtkSmartPointer.h>
#include <vtkTextProperty.h>
#include <vtkVersion.h>
#include <vtkXMLPolyDataReader.h>

#include <vtk_cli11.h>
#include <vtk_fmt.h>
// clang-format off
#include VTK_FMT(fmt/format.h)
// clang-format on

#if VTK_VERSION_NUMBER >= 90020210809ULL
#define HAS_COW
#include <vtkCameraOrientationWidget.h>
#endif

// vtkGenerateIds was introduced in VTK build version 20240504
#if VTK_BUILD_VERSION >= 20240504
#define USE_USE_GENERATE_IDS
#include <vtkGenerateIds.h>
#else
#include <vtkIdFilter.h>
#endif

#include <array>
#include <filesystem>
#include <numeric>
#include <set>

namespace fs = std::filesystem;

namespace {

//! Adjust curvatures along the edges of the surface.
/*!
 * This function adjusts curvatures along the edges of the surface by replacing
 *  the value with the average value of the curvatures of points in the
 *  neighborhood.
 *
 * Remember to update the vtkCurvatures object before calling this.
 *
 * @param source - A vtkPolyData object corresponding to the vtkCurvatures
 * object.
 * @param curvatureName: The name of the curvature, "Gauss_Curvature" or
 * "Mean_Curvature".
 * @param epsilon: Curvature values less than this will be set to zero.
 * @return
 */
void AdjustEdgeCurvatures(vtkPolyData* source, std::string const& curvatureName,
                          double const& epsilon = 1.0e-08);
/**
 * Get the color transfer function to use as a lookup table.
 *
 * @param scalarRange - the scalar range.
 * @param schemeNumber - The Brewer Color Scheme number.
 * @return The color transfer function.
 */
vtkNew<vtkColorTransferFunction> GetCtfLut(double* scalarRange,
                                           const int& schemeNumber);

typedef std::map<std::string, std::array<double, 2>> TTextPosition;
typedef std::map<std::string, TTextPosition> TTextPositions;

struct ScalarBarProperties
{
  vtkSmartPointer<vtkNamedColors> colors;

  // The properties needed for scalar bars.
  vtkSmartPointer<vtkColorTransferFunction> lut;
  // These are in pixels.
  std::map<std::string, int> maximumDimensions{{"width", 100}, {"height", 260}};
  std::string titleText{""};
  int number_of_labels{5};
  std::string labelFormat{"{:0.2f}"};
  // Orientation vertical=true, horizontal=false.
  bool orientation{true};
  // Horizontal and vertical positioning.
  // These are the default positions, don't change these.
  TTextPosition defaultV = {{"p", {0.85, 0.1}}, {"p2", {0.1, 0.7}}};
  TTextPosition defaultH = {{"p", {0.10, 0.1}}, {"p2", {0.7, 0.1}}};
  // Modify these as needed.
  TTextPosition positionV = {{"p", {0.85, 0.1}}, {"p2", {0.1, 0.7}}};
  TTextPosition positionH = {{"p", {0.10, 0.1}}, {"p2", {0.7, 0.1}}};
};

/** Make a scalar bar widget.
 *
 *  @param scalar_bar_properties - The lookup table, title name, maximum
 * dimensions in pixels and position.
 *  @param textProperty - The properties for the title.
 *  @param labelTextProperty - The properties for the labels.
 *  @param ren - The vtkRenderer.
 *  @param iren - The vtkInteractor.
 *  @return The scalar bar widget.
 */
vtkNew<vtkScalarBarWidget>
MakeScalarBarWidget(ScalarBarProperties& scalarBarProperties,
                    vtkTextProperty* textProperty,
                    vtkTextProperty* labelTextProperty, vtkRenderer* ren,
                    vtkRenderWindowInteractor* iren);

} // namespace

int main(int argc, char* argv[])
{
  vtkNew<vtkNamedColors> colors;
  colors->SetColor("ParaViewBlueGrayBkg",
                   std::array<unsigned char, 4>{84, 89, 109, 255}.data());
  colors->SetColor("ParaViewWarmGrayBkg",
                   std::array<unsigned char, 4>{98, 93, 90, 255}.data());

  CLI::App app{"Compute Gaussian or Mean Curvatures."};

  // Define options.
  std::string fileName{""};
  app.add_option("fileName", fileName,
                 "The path to a file that contains vtkPolyData e.g. "
                 "cowHead.vtp");
  auto colorScheme = 16;
  app.add_option("colorScheme", colorScheme,
                 "The Brewer Color Scheme number, default is 16.");

  auto gaussianCurvature{false};
  app.add_flag("-g", gaussianCurvature,
               "Use gaussian curvature instead of mean curvature.");

  CLI11_PARSE(app, argc, argv);

  auto path = fs::path(fileName);
  if (!path.empty())
  {
    if (!fs::is_regular_file(path))
    {
      std::cerr << fmt::format("Unable to find: {:s}", path.string())
                << std::endl;
      return EXIT_FAILURE;
    }
  }

  std::string curvature = "Mean_Curvature";
  if (gaussianCurvature)
  {
    curvature = "Gauss_Curvature";
  }
  // Create a polydata.
  vtkNew<vtkXMLPolyDataReader> reader;
  reader->SetFileName(argv[1]);
  reader->Update();

  auto source = reader->GetOutput();

  vtkNew<vtkCurvatures> cc;
  cc->SetInputData(source);
  if (curvature == "Gauss_Curvature")
  {
    cc->SetCurvatureTypeToGaussian();
    cc->Update();
  }
  else
  {
    if (curvature == "Mean_Curvature")
    {
      cc->SetCurvatureTypeToMean();
      cc->Update();
    }
    else
    {
      std::cerr << "Unknown curvature" << std::endl;
      return EXIT_FAILURE;
    }
  }
  AdjustEdgeCurvatures(cc->GetOutput(), curvature);
  source->GetPointData()->AddArray(
      cc->GetOutput()->GetPointData()->GetAbstractArray(curvature.c_str()));

  auto scalarRange =
      source->GetPointData()->GetScalars(curvature.c_str())->GetRange();

  std::string curvatureTitle = curvature;
  std::replace(curvatureTitle.begin(), curvatureTitle.end(), '_', '\n');

  auto ctfLut = GetCtfLut(scalarRange, colorScheme);

  // Create a mapper and actor.
  vtkNew<vtkPolyDataMapper> mapper;
  mapper->SetInputData(source);
  mapper->SetScalarModeToUsePointFieldData();
  mapper->SelectColorArray(curvature.c_str());
  mapper->SetScalarRange(scalarRange);
  mapper->SetLookupTable(ctfLut);

  vtkNew<vtkActor> actor;
  actor->SetMapper(mapper);

  auto windowWidth = 800;
  auto windowHeight = 800;

  // Create a renderer, render window, and interactor.
  auto appFn = fs::path((app.get_name())).stem().string();
  vtkNew<vtkRenderer> renderer;
  vtkNew<vtkRenderWindow> renWin;
  renWin->AddRenderer(renderer);
  renWin->SetSize(windowWidth, windowHeight);
  if (path.empty())
  {
    renWin->SetWindowName(appFn.c_str());
  }
  else
  {
    auto winName = fmt::format("{:s} {:s}", appFn, path.filename().string());
    renWin->SetWindowName(winName.c_str());
  }

  vtkNew<vtkRenderWindowInteractor> iRen;
  iRen->SetRenderWindow(renWin);
  // Important: The interactor must be set prior to enabling the widget.
  iRen->SetRenderWindow(renWin);

  vtkNew<vtkTextProperty> textProperty;
  textProperty->SetColor(colors->GetColor3d("AliceBlue").GetData());
  textProperty->SetJustificationToLeft();
  textProperty->SetFontSize(16);
  textProperty->BoldOn();
  textProperty->ItalicOn();
  textProperty->ShadowOn();

  // Add the scalar bar.
  auto sbProperties = ScalarBarProperties();

  if (gaussianCurvature)
  {
    sbProperties.titleText = "Gaussian\nCurvature\n";
  }
  else
  {
    sbProperties.titleText = "Mean\nCurvature\n";
  }
  sbProperties.positionV = {{"p", {0.85, 0.1}}, {"p2", {0.125, 0.375}}};
  sbProperties.lut = ctfLut;
  auto sbWidgetCurv = MakeScalarBarWidget(sbProperties, textProperty,
                                          textProperty, renderer, iRen);

#ifdef HAS_COW
  vtkNew<vtkCameraOrientationWidget> camOrientManipulator;
  camOrientManipulator->SetParentRenderer(renderer);
  // Enable the widget.
  camOrientManipulator->On();
#endif

  // Add the actors to the scene.
  renderer->AddActor(actor);
  // renderer->AddViewProp(sbWidgetCurv);
  renderer->SetBackground(colors->GetColor3d("DarkSlateGray").GetData());

  // Render and interact.
  renWin->Render();
  iRen->Start();

  return EXIT_SUCCESS;
}

namespace {
void AdjustEdgeCurvatures(vtkPolyData* source, std::string const& curvatureName,
                          double const& epsilon)
{
  auto PointNeighbourhood =
      [&source](vtkIdType const& pId) -> std::set<vtkIdType> {
    // Extract the topological neighbors for point pId. In two steps:
    //  1) source->GetPointCells(pId, cellIds)
    //  2) source->GetCellPoints(cellId, cellPointIds) for all cellId in cellIds
    vtkNew<vtkIdList> cellIds;
    source->GetPointCells(pId, cellIds);
    std::set<vtkIdType> neighbours;
    for (vtkIdType i = 0; i < cellIds->GetNumberOfIds(); ++i)
    {
      auto cellId = cellIds->GetId(i);
      vtkNew<vtkIdList> cellPointIds;
      source->GetCellPoints(cellId, cellPointIds);
      for (vtkIdType j = 0; j < cellPointIds->GetNumberOfIds(); ++j)
      {
        neighbours.insert(cellPointIds->GetId(j));
      }
    }
    return neighbours;
  };

  auto ComputeDistance = [&source](vtkIdType const& ptIdA,
                                   vtkIdType const& ptIdB) {
    std::array<double, 3> ptA{0.0, 0.0, 0.0};
    std::array<double, 3> ptB{0.0, 0.0, 0.0};
    std::array<double, 3> ptC{0.0, 0.0, 0.0};
    source->GetPoint(ptIdA, ptA.data());
    source->GetPoint(ptIdB, ptB.data());
    std::transform(std::begin(ptA), std::end(ptA), std::begin(ptB),
                   std::begin(ptC), std::minus<double>());
    // Calculate the norm.
    auto result = std::sqrt(std::inner_product(std::begin(ptC), std::end(ptC),
                                               std::begin(ptC), 0.0));
    return result;
  };

  source->GetPointData()->SetActiveScalars(curvatureName.c_str());
  // Curvature as a vector.
  auto array = source->GetPointData()->GetAbstractArray(curvatureName.c_str());
  std::vector<double> curvatures;
  for (vtkIdType i = 0; i < source->GetNumberOfPoints(); ++i)
  {
    curvatures.push_back(array->GetVariantValue(i).ToDouble());
  }

  // Get the boundary point IDs.
  std::string name = "Ids";
#ifdef USE_USE_GENERATE_IDS
  vtkNew<vtkGenerateIds> idFilter;
#else
  vtkNew<vtkIdFilter> idFilter;
#endif
  idFilter->SetInputData(source);
  idFilter->SetPointIds(true);
  idFilter->SetCellIds(false);
  idFilter->SetPointIdsArrayName(name.c_str());
  idFilter->SetCellIdsArrayName(name.c_str());
  idFilter->Update();

  vtkNew<vtkFeatureEdges> edges;

  edges->SetInputConnection(idFilter->GetOutputPort());
  edges->BoundaryEdgesOn();
  edges->ManifoldEdgesOff();
  edges->NonManifoldEdgesOff();
  edges->FeatureEdgesOff();
  edges->Update();

  auto edgeAarray =
      edges->GetOutput()->GetPointData()->GetAbstractArray(name.c_str());
  std::vector<vtkIdType> boundaryIds;
  for (vtkIdType i = 0; i < edges->GetOutput()->GetNumberOfPoints(); ++i)
  {
    boundaryIds.push_back(edgeAarray->GetVariantValue(i).ToInt());
  }
  // Remove duplicate Ids.
  std::set<vtkIdType> pIdsSet(boundaryIds.begin(), boundaryIds.end());
  for (auto const pId : boundaryIds)
  {
    auto pIdsNeighbors = PointNeighbourhood(pId);
    std::set<vtkIdType> pIdsNeighborsInterior;
    std::set_difference(
        pIdsNeighbors.begin(), pIdsNeighbors.end(), pIdsSet.begin(),
        pIdsSet.end(),
        std::inserter(pIdsNeighborsInterior, pIdsNeighborsInterior.begin()));
    // Compute distances and extract curvature values.
    std::vector<double> curvs;
    std::vector<double> dists;
    for (auto const pIdN : pIdsNeighborsInterior)
    {
      curvs.push_back(curvatures[pIdN]);
      dists.push_back(ComputeDistance(pIdN, pId));
    }
    std::vector<vtkIdType> nonZeroDistIds;
    for (size_t i = 0; i < dists.size(); ++i)
    {
      if (dists[i] > 0)
      {
        nonZeroDistIds.push_back(i);
      }
    }
    std::vector<double> curvsNonZero;
    std::vector<double> distsNonZero;
    for (auto const i : nonZeroDistIds)
    {
      curvsNonZero.push_back(curvs[i]);
      distsNonZero.push_back(dists[i]);
    }
    // Iterate over the edge points and compute the curvature as the weighted
    // average of the neighbours.
    auto countInvalid = 0;
    auto newCurv = 0.0;
    if (curvsNonZero.size() > 0)
    {
      std::vector<double> weights;
      double sum = 0.0;
      for (auto const d : distsNonZero)
      {
        sum += 1.0 / d;
        weights.push_back(1.0 / d);
      }
      for (size_t i = 0; i < weights.size(); ++i)
      {
        weights[i] = weights[i] / sum;
      }
      newCurv = std::inner_product(curvsNonZero.begin(), curvsNonZero.end(),
                                   weights.begin(), 0.0);
    }
    else
    {
      // Corner case.
      // countInvalid += 1;
      // Assuming the curvature of the point is planar.
      newCurv = 0.0;
    }
    // Set the new curvature value.
    curvatures[pId] = newCurv;
  }

  // Set small values to zero.
  if (epsilon != 0.0)
  {
    auto eps = std::abs(epsilon);
    for (size_t i = 0; i < curvatures.size(); ++i)
    {
      if (std::abs(curvatures[i]) < eps)
      {
        curvatures[i] = 0.0;
      }
    }
  }

  if (static_cast<size_t>(source->GetNumberOfPoints()) != curvatures.size())
  {
    std::string s = curvatureName;
    s += ":\nCannot add the adjusted curvatures to the source.\n";
    s += " The number of points in source does not equal the\n";
    s += " number of point ids in the adjusted curvature array.";
    std::cerr << s << std::endl;
    return;
  }
  vtkNew<vtkDoubleArray> adjustedCurvatures;
  adjustedCurvatures->SetName(curvatureName.c_str());
  for (auto curvature : curvatures)
  {
    adjustedCurvatures->InsertNextTuple1(curvature);
  }
  source->GetPointData()->AddArray(adjustedCurvatures);
  source->GetPointData()->SetActiveScalars(curvatureName.c_str());
}

vtkNew<vtkColorTransferFunction> GetCtfLut(double* scalarRange,
                                           const int& schemeNumber)
{
  vtkNew<vtkColorSeries> colorSeries;
  colorSeries->SetColorScheme(schemeNumber);
  std::cout << fmt::format("Using color scheme: {:2d}, {:s}\n",
                           colorSeries->GetColorScheme(),
                           colorSeries->GetColorSchemeName());

  // Build a lookup table.
  vtkNew<vtkColorTransferFunction> ctfLut;
  ctfLut->SetColorSpaceToHSV();

  // Use a color series to create a transfer function.
  auto numColors = colorSeries->GetNumberOfColors();
  double dRange = scalarRange[1] - scalarRange[0];
  double nc = numColors - 1;
  for (int i = 0; i < numColors; i++)
  {
    vtkColor3ub color = colorSeries->GetColor(i);
    double dColor[3]{0.0, 0.0, 0.0};
    dColor[0] = static_cast<double>(color[0]) / 255.0;
    dColor[1] = static_cast<double>(color[1]) / 255.0;
    dColor[2] = static_cast<double>(color[2]) / 255.0;
    double t = scalarRange[0] + i * dRange / nc;
    ctfLut->AddRGBPoint(t, dColor[0], dColor[1], dColor[2]);
  }
  ctfLut->Build();

  return ctfLut;
}

vtkNew<vtkScalarBarWidget>
MakeScalarBarWidget(ScalarBarProperties& sbProperties,
                    vtkTextProperty* textProperty,
                    vtkTextProperty* labelTextProperty, vtkRenderer* ren,
                    vtkRenderWindowInteractor* iren)
{
  vtkNew<vtkScalarBarActor> sbActor;
  sbActor->SetLookupTable(sbProperties.lut);
  sbActor->SetTitle(sbProperties.titleText.c_str());
  sbActor->UnconstrainedFontSizeOn();
  sbActor->SetNumberOfLabels(sbProperties.number_of_labels);
  sbActor->SetTitleTextProperty(textProperty);
  sbActor->SetLabelTextProperty(labelTextProperty);
  sbActor->SetLabelFormat(sbProperties.labelFormat.c_str());

  vtkNew<vtkScalarBarRepresentation> sbRep;
  sbRep->EnforceNormalizedViewportBoundsOn();
  sbRep->SetOrientation(sbProperties.orientation);
  // Set the position.
  sbRep->GetPositionCoordinate()->SetCoordinateSystemToNormalizedViewport();
  sbRep->GetPosition2Coordinate()->SetCoordinateSystemToNormalizedViewport();
  if (sbProperties.orientation)
  {
    auto p1 = sbProperties.positionV["p"];
    auto p2 = sbProperties.positionV["p2"];
    sbRep->GetPositionCoordinate()->SetValue(p1.data());
    sbRep->GetPosition2Coordinate()->SetValue(p2.data());
  }
  else
  {
    auto p1 = sbProperties.positionH["p"];
    auto p2 = sbProperties.positionH["p2"];
    sbRep->GetPositionCoordinate()->SetValue(p1.data());
    sbRep->GetPosition2Coordinate()->SetValue(p2.data());
  }

  vtkNew<vtkScalarBarWidget> widget;
  widget->SetRepresentation(sbRep);
  widget->SetScalarBarActor(sbActor);
  widget->SetInteractor(iren);
  widget->SetDefaultRenderer(ren);
  widget->EnabledOn();

  return widget;
}

} // namespace

cmake_minimum_required(VERSION 3.12 FATAL_ERROR)

project(Curvatures)

find_package(VTK COMPONENTS 
  CommonColor
  CommonCore
  CommonDataModel
  FiltersCore
  FiltersGeneral
  IOXML
  InteractionStyle
  InteractionWidgets
  RenderingAnnotation
  RenderingContextOpenGL2
  RenderingCore
  RenderingFreeType
  RenderingGL2PSOpenGL2
  RenderingOpenGL2
)

if (NOT VTK_FOUND)
  message(FATAL_ERROR "Curvatures: Unable to find the VTK build folder.")
endif()

# Prevent a "command line is too long" failure in Windows.
set(CMAKE_NINJA_FORCE_RESPONSE_FILE "ON" CACHE BOOL "Force Ninja to use response files.")
add_executable(Curvatures MACOSX_BUNDLE Curvatures.cxx )
  target_link_libraries(Curvatures PRIVATE ${VTK_LIBRARIES}
)
# vtk_module_autoinit is needed
vtk_module_autoinit(
  TARGETS Curvatures
  MODULES ${VTK_LIBRARIES}
)

cd Curvatures/build

cmake ..

cmake -DVTK_DIR:PATH=/home/me/vtk_build ..

make

./Curvatures