Working on DeepSSM command

Author: akenmorris

Applications/shapeworks/CMakeLists.txt

@@ -23,7 +23,7 @@ target_include_directories(shapeworks_exe PUBLIC
 
 target_link_libraries(shapeworks_exe
   Mesh ${VTK_LIBRARIES} Optimize Utils trimesh2 Particles
-  pybind11::embed Project Image Groom Analyze
+  pybind11::embed Project Image Groom Analyze Application
   )
 
 message(STATUS "opt libs ${OPTIMIZE_LIBRARIES}")

Applications/shapeworks/Command.h

@@ -103,6 +103,15 @@ class ParticleSystemCommand : public Command
 private:
 };
 
+class DeepSSMCommandGroup : public Command
+{
+public:
+  const std::string type() override { return "DeepSSM"; }
+
+private:
+};
+
+
 class ShapeworksCommand : public Command
 {
 public:

Applications/shapeworks/Commands.cpp

@@ -1,18 +1,18 @@
 #include "Commands.h"
 
 #include <Analyze/Analyze.h>
+#include <Application/DeepSSM/DeepSSMJob.h>
 #include <Groom/Groom.h>
 #include <Logging.h>
 #include <Optimize/Optimize.h>
 #include <Optimize/OptimizeParameterFile.h>
 #include <Optimize/OptimizeParameters.h>
+#include <Profiling.h>
 #include <ShapeworksUtils.h>
 #include <Utils/StringUtils.h>
 
 #include <boost/filesystem.hpp>
 
-#include <Profiling.h>
-
 namespace shapeworks {
 
 // boilerplate for a command. Copy this to start a new command
@@ -43,8 +43,6 @@ bool Example::execute(const optparse::Values &options, SharedCommandData &shared
 }
 #endif
 
-
-
 ///////////////////////////////////////////////////////////////////////////////
 // Seed
 ///////////////////////////////////////////////////////////////////////////////
@@ -331,4 +329,110 @@ bool ConvertProjectCommand::execute(const optparse::Values& options, SharedComma
     return false;
   }
 }
+
+///////////////////////////////////////////////////////////////////////////////
+// DeepSSM
+///////////////////////////////////////////////////////////////////////////////
+void DeepSSMCommand::buildParser() {
+  const std::string prog = "deepssm";
+  const std::string desc = "run deepssm steps";
+  parser.prog(prog).description(desc);
+
+  parser.add_option("--name").action("store").type("string").set_default("").help(
+      "Path to input project file (xlsx or swproj).");
+
+  // Create a vector of choices first
+  std::vector<std::string> prep_choices = {"all", "groom_training", "optimize_training", "optimize_validation",
+                                           "groom_images"};
+
+  // --prep option with choices
+  parser.add_option("--prep")
+      .action("store")
+      .type("choice")
+      .choices(prep_choices.begin(), prep_choices.end())
+      .set_default("all")
+      .help("Preparation step to run");
+
+  // Boolean flag options
+  parser.add_option("--augment").action("store_true").help("Run data augmentation");
+
+  parser.add_option("--train").action("store_true").help("Run training");
+
+  parser.add_option("--test").action("store_true").help("Run testing");
+
+  parser.add_option("--all").action("store_true").help("Run all steps");
+
+  Command::buildParser();
+}
+
+bool DeepSSMCommand::execute(const optparse::Values& options, SharedCommandData& sharedData) {
+  // Handle project file: either from --name or first positional argument
+  std::string project_file;
+  if (options.is_set_by_user("name")) {
+    // User explicitly provided --name
+    project_file = options["name"];
+  } else if (!parser.args().empty()) {
+    // Use first positional argument
+    project_file = parser.args()[0];
+  } else {
+    // No project file provided at all
+    parser.error("Project file must be provided either as --name or as a positional argument");
+  }
+
+  std::cout << "DeepSSM: Using project file: " << project_file << std::endl;
+
+  bool do_prep = options.is_set("prep") || options.is_set("all");
+  std::string prep_step = options["prep"];
+  bool do_augment = options.is_set("augment") || options.is_set("all");
+  bool do_train = options.is_set("train") || options.is_set("all");
+  bool do_test = options.is_set("test") || options.is_set("all");
+  if (!do_prep && !do_augment && !do_train && !do_test) {
+    do_prep = true;
+    do_augment = true;
+    do_train = true;
+    do_test = true;
+  }
+
+  ProjectHandle project = std::make_shared<Project>();
+  project->load(project_file);
+
+  DeepSSMJob job(project, DeepSSMJob::JobType::DeepSSM_PrepType);
+
+  if (do_prep) {
+    if (prep_step == "all") {
+      job.set_prep_step(DeepSSMJob::PrepStep::NOT_STARTED);
+    } else if (prep_step == "groom_training") {
+      job.set_prep_step(DeepSSMJob::PrepStep::GROOM_TRAINING);
+    } else if (prep_step == "optimize_training") {
+      job.set_prep_step(DeepSSMJob::PrepStep::OPTIMIZE_TRAINING);
+    } else if (prep_step == "optimize_validation") {
+      job.set_prep_step(DeepSSMJob::PrepStep::OPTIMIZE_VALIDATION);
+    } else if (prep_step == "groom_images") {
+      job.set_prep_step(DeepSSMJob::PrepStep::GROOM_IMAGES);
+    } else {
+      SW_ERROR("Unknown prep step: {}", prep_step);
+      return false;
+    }
+    std::cout << "Running DeepSSM preparation step...\n";
+    job.run_prep();
+  }
+  if (do_augment) {
+    std::cout << "Running DeepSSM data augmentation...\n";
+    job.run_augmentation();
+  }
+  if (do_train) {
+    std::cout << "Running DeepSSM training...\n";
+    job.run_training();
+  }
+  if (do_test) {
+    std::cout << "Running DeepSSM testing...\n";
+    job.run_testing();
+  }
+
+  project->save();
+
+  SW_ERROR("DeepSSM command is not implemented yet.");
+  return false;
+}
+
 }  // namespace shapeworks

Applications/shapeworks/Commands.h

@@ -101,5 +101,6 @@ COMMAND_DECLARE(OptimizeCommand, OptimizeCommandGroup);
 COMMAND_DECLARE(GroomCommand, GroomCommandGroup);
 COMMAND_DECLARE(AnalyzeCommand, AnalyzeCommandGroup);
 COMMAND_DECLARE(ConvertProjectCommand, ProjectCommandGroup);
+COMMAND_DECLARE(DeepSSMCommand, DeepSSMCommandGroup);
 
 } // shapeworks

Applications/shapeworks/shapeworks.cpp

@@ -110,6 +110,7 @@ int main(int argc, char *argv[])
   shapeworks.addCommand(GroomCommand::getCommand());
   shapeworks.addCommand(AnalyzeCommand::getCommand());
   shapeworks.addCommand(ConvertProjectCommand::getCommand());
+  shapeworks.addCommand(DeepSSMCommand::getCommand());
 
   try {
     TIME_START("shapeworks");

Libs/Application/DeepSSM/DeepSSMJob.cpp

@@ -22,26 +22,26 @@ using namespace pybind11::literals;  // to bring in the `_a` literal
 namespace shapeworks {
 
 //---------------------------------------------------------------------------
-DeepSSMJob::DeepSSMJob(std::shared_ptr<Project> project, DeepSSMJob::ToolMode tool_mode, DeepSSMJob::PrepStep prep_step)
-    : project_(project), tool_mode_(tool_mode), prep_step_(prep_step) {}
+DeepSSMJob::DeepSSMJob(std::shared_ptr<Project> project, DeepSSMJob::JobType tool_mode, DeepSSMJob::PrepStep prep_step)
+    : project_(project), job_type_(tool_mode), prep_step_(prep_step) {}
 
 //---------------------------------------------------------------------------
 DeepSSMJob::~DeepSSMJob() {}
 
 //---------------------------------------------------------------------------
 void DeepSSMJob::run() {
   try {
-    switch (tool_mode_) {
-      case DeepSSMJob::ToolMode::DeepSSM_PrepType:
+    switch (job_type_) {
+      case DeepSSMJob::JobType::DeepSSM_PrepType:
         run_prep();
         break;
-      case DeepSSMJob::ToolMode::DeepSSM_AugmentationType:
+      case DeepSSMJob::JobType::DeepSSM_AugmentationType:
         run_augmentation();
         break;
-      case DeepSSMJob::ToolMode::DeepSSM_TrainingType:
+      case DeepSSMJob::JobType::DeepSSM_TrainingType:
         run_training();
         break;
-      case DeepSSMJob::ToolMode::DeepSSM_TestingType:
+      case DeepSSMJob::JobType::DeepSSM_TestingType:
         run_testing();
         break;
     }
@@ -52,17 +52,17 @@ void DeepSSMJob::run() {
 
 //---------------------------------------------------------------------------
 QString DeepSSMJob::name() {
-  switch (tool_mode_) {
-    case DeepSSMJob::ToolMode::DeepSSM_PrepType:
+  switch (job_type_) {
+    case DeepSSMJob::JobType::DeepSSM_PrepType:
       return "DeepSSM: Prep";
       break;
-    case DeepSSMJob::ToolMode::DeepSSM_AugmentationType:
+    case DeepSSMJob::JobType::DeepSSM_AugmentationType:
       return "DeepSSM: Augmentation";
       break;
-    case DeepSSMJob::ToolMode::DeepSSM_TrainingType:
+    case DeepSSMJob::JobType::DeepSSM_TrainingType:
       return "DeepSSM: Training";
       break;
-    case DeepSSMJob::ToolMode::DeepSSM_TestingType:
+    case DeepSSMJob::JobType::DeepSSM_TestingType:
       return "DeepSSM: Testing";
       break;
   }

Libs/Application/DeepSSM/DeepSSMJob.h

@@ -16,7 +16,7 @@ class DeepSSMJob : public Job {
   Q_OBJECT;
 
  public:
-  enum class ToolMode {
+  enum class JobType {
     DeepSSM_PrepType = 0,
     DeepSSM_AugmentationType = 1,
     DeepSSM_TrainingType = 2,
@@ -34,7 +34,7 @@ class DeepSSMJob : public Job {
 
   enum class SplitType { TRAIN, VAL, TEST };
 
-  DeepSSMJob(std::shared_ptr<Project> project, DeepSSMJob::ToolMode tool_mode,
+  DeepSSMJob(std::shared_ptr<Project> project, DeepSSMJob::JobType tool_mode,
              DeepSSMJob::PrepStep prep_step = DeepSSMJob::NOT_STARTED);
   ~DeepSSMJob();
 
@@ -51,13 +51,18 @@ class DeepSSMJob : public Job {
 
   static std::vector<int> get_split(ProjectHandle project, DeepSSMJob::SplitType split_type);
 
+  void set_prep_step(DeepSSMJob::PrepStep step) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    prep_step_ = step;
+  }
+
  private:
   void update_prep_stage(DeepSSMJob::PrepStep step);
   void process_test_results();
 
   std::shared_ptr<Project> project_;
 
-  DeepSSMJob::ToolMode tool_mode_;
+  DeepSSMJob::JobType job_type_;
 
   QString prep_message_;
   DeepSSMJob::PrepStep prep_step_{DeepSSMJob::NOT_STARTED};