Merge pull request #2486 from SCIInstitute/deepssm_refactor2

Refactors DeepSSM for reliability, memory efficiency, and testability.

Author: akenmorris

Applications/shapeworks/Command.h

@@ -33,7 +33,7 @@ class Command {
   const std::string desc() const { return parser.description(); }
 
   /// parses the arguments for this command, saving them in the parser and returning the leftovers
-  std::vector<std::string> parse_args(const std::vector<std::string> &arguments);
+  virtual std::vector<std::string> parse_args(const std::vector<std::string> &arguments);
 
   /// calls execute for this command using the parsed args, returning system exit value
   int run(SharedCommandData &sharedData);
@@ -108,6 +108,12 @@ class DeepSSMCommandGroup : public Command
 public:
   const std::string type() override { return "DeepSSM"; }
 
+  // DeepSSM is a terminal command - don't pass remaining args to other commands
+  std::vector<std::string> parse_args(const std::vector<std::string> &arguments) override {
+    Command::parse_args(arguments);
+    return {};  // return empty - DeepSSM consumes all args
+  }
+
 private:
 };
 

Applications/shapeworks/Commands.cpp

@@ -12,7 +12,7 @@
 #include <ShapeworksUtils.h>
 #include <Utils/StringUtils.h>
 
-#include <QApplication>
+#include <QCoreApplication>
 #include <boost/filesystem.hpp>
 
 namespace shapeworks {
@@ -371,18 +371,23 @@ void DeepSSMCommand::buildParser() {
       .set_default(0)
       .help("Number of data loader workers (default: 0)");
 
+  parser.add_option("--aug_processes")
+      .action("store")
+      .type("int")
+      .set_default(0)
+      .help("Number of augmentation processes (default: 0 = use all cores)");
+
   Command::buildParser();
 }
 
 bool DeepSSMCommand::execute(const optparse::Values& options, SharedCommandData& sharedData) {
-  // Create a non-gui QApplication instance
-  int argc = 3;
-  char* argv[3];
+  // QCoreApplication provides the event loop needed for PythonWorker's QThread,
+  // without requiring Qt platform plugins (which may not be available on headless CI).
+  int argc = 1;
+  char* argv[1];
   argv[0] = const_cast<char*>("shapeworks");
-  argv[1] = const_cast<char*>("-platform");
-  argv[2] = const_cast<char*>("offscreen");
 
-  QApplication app(argc, argv);
+  QCoreApplication app(argc, argv);
 
   // Handle project file: either from --name or first positional argument
   std::string project_file;
@@ -413,12 +418,14 @@ bool DeepSSMCommand::execute(const optparse::Values& options, SharedCommandData&
   bool do_test = options.is_set("test") || options.is_set("all");
 
   int num_workers = static_cast<int>(options.get("num_workers"));
+  int aug_processes = static_cast<int>(options.get("aug_processes"));
 
   std::cout << "Prep step:    " << (do_prep ? "on" : "off") << "\n";
   std::cout << "Augment step: " << (do_augment ? "on" : "off") << "\n";
   std::cout << "Train step:   " << (do_train ? "on" : "off") << "\n";
   std::cout << "Test step:    " << (do_test ? "on" : "off") << "\n";
   std::cout << "Num dataloader workers:  " << num_workers << "\n";
+  std::cout << "Augmentation processes:  " << (aug_processes == 0 ? QThread::idealThreadCount() : aug_processes) << "\n";
 
   if (!do_prep && !do_augment && !do_train && !do_test) {
     do_prep = true;
@@ -472,6 +479,7 @@ bool DeepSSMCommand::execute(const optparse::Values& options, SharedCommandData&
   if (do_augment) {
     std::cout << "Running DeepSSM data augmentation...\n";
     auto job = QSharedPointer<DeepSSMJob>::create(project, DeepSSMJob::JobType::DeepSSM_AugmentationType);
+    job->set_aug_processes(aug_processes);
     python_worker.run_job(job);
     if (!wait_for_job(job)) {
       return false;
@@ -499,7 +507,7 @@ bool DeepSSMCommand::execute(const optparse::Values& options, SharedCommandData&
 
   project->save();
 
-  return false;
+  return true;
 }
 
 }  // namespace shapeworks

Examples/Python/RunUseCase.py

@@ -69,8 +69,14 @@
     parser.add_argument("--tiny_test", help="Run as a short test", action="store_true")
     parser.add_argument("--verify", help="Run as a full test", action="store_true")
     parser.add_argument("--clean", help="Run from scratch, ignoring intermediate stages", action="store_true")
+    parser.add_argument("--exact_check", help="Save or verify exact values for refactoring verification (platform-specific)",
+                        choices=["save", "verify"])
     args = parser.parse_args()
 
+    # Validate deep_ssm-specific arguments
+    if args.exact_check and args.use_case != "deep_ssm":
+        parser.error("--exact_check is only supported for the deep_ssm use case")
+
     type = ""
     if args.tiny_test:
         type = "tiny_test_"

Examples/Python/deep_ssm.py

@@ -64,8 +64,8 @@ def Run_Pipeline(args):
         This data is comprised of femur meshes and corresponding hip CT scans.
         """
 
-        if platform.system() == "Darwin":
-            # On MacOS, CPU PyTorch is hanging with parallel
+        if platform.system() != "Linux":
+            # CPU PyTorch hangs with OpenMP parallelism on macOS and Windows
             os.environ['OMP_NUM_THREADS'] = "1"
         # If running a tiny_test, then download subset of the data
         if args.tiny_test:
@@ -396,6 +396,7 @@ def Run_Pipeline(args):
                 "c_lat": 6.3
             }
         }
+
         if args.tiny_test:
             model_parameters["trainer"]["epochs"] = 1
         # Save config file
@@ -436,17 +437,17 @@ def Run_Pipeline(args):
             val_world_particles.append(project_path + subjects[index].get_world_particle_filenames()[0])
             val_mesh_files.append(project_path + subjects[index].get_groomed_filenames()[0])
 
-        val_out_dir = output_directory + model_name + '/validation_predictions/'
         predicted_val_world_particles = DeepSSMUtils.testDeepSSM(config_file, loader='validation')
         print("Validation world predictions saved.")
-        # Generate local predictions
-        local_val_prediction_dir = val_out_dir + 'local_predictions/'
+        # Generate local predictions - create directory next to world_predictions
+        world_pred_dir = os.path.dirname(predicted_val_world_particles[0])
+        local_val_prediction_dir = world_pred_dir.replace("world_predictions", "local_predictions")
         if not os.path.exists(local_val_prediction_dir):
             os.makedirs(local_val_prediction_dir)
         predicted_val_local_particles = []
         for particle_file, transform in zip(predicted_val_world_particles, val_transforms):
             particles = np.loadtxt(particle_file)
-            local_particle_file = particle_file.replace("world_predictions/", "local_predictions/")
+            local_particle_file = particle_file.replace("world_predictions", "local_predictions")
             local_particles = sw.utils.transformParticles(particles, transform, inverse=True)
             np.savetxt(local_particle_file, local_particles)
             predicted_val_local_particles.append(local_particle_file)
@@ -462,6 +463,8 @@ def Run_Pipeline(args):
         template_mesh = project_path + subjects[reference_index].get_groomed_filenames()[0]
         template_particles = project_path + subjects[reference_index].get_local_particle_filenames()[0]
         # Get distance between clipped true and predicted meshes
+        # Get the validation output directory from the predictions path
+        val_out_dir = os.path.dirname(local_val_prediction_dir.rstrip('/')) + '/'
         mean_dist = DeepSSMUtils.analyzeMeshDistance(predicted_val_local_particles, val_mesh_files,
                                                      template_particles, template_mesh, val_out_dir,
                                                      planes=val_planes)
@@ -500,17 +503,17 @@ def Run_Pipeline(args):
             with open(plane_file) as json_file:
                 test_planes.append(json.load(json_file)['planes'][0]['points'])
 
-        test_out_dir = output_directory + model_name + '/test_predictions/'
         predicted_test_world_particles = DeepSSMUtils.testDeepSSM(config_file, loader='test')
         print("Test world predictions saved.")
-        # Generate local predictions
-        local_test_prediction_dir = test_out_dir + 'local_predictions/'
+        # Generate local predictions - create directory next to world_predictions
+        world_pred_dir = os.path.dirname(predicted_test_world_particles[0])
+        local_test_prediction_dir = world_pred_dir.replace("world_predictions", "local_predictions")
         if not os.path.exists(local_test_prediction_dir):
             os.makedirs(local_test_prediction_dir)
         predicted_test_local_particles = []
         for particle_file, transform in zip(predicted_test_world_particles, test_transforms):
             particles = np.loadtxt(particle_file)
-            local_particle_file = particle_file.replace("world_predictions/", "local_predictions/")
+            local_particle_file = particle_file.replace("world_predictions", "local_predictions")
             local_particles = sw.utils.transformParticles(particles, transform, inverse=True)
             np.savetxt(local_particle_file, local_particles)
             predicted_test_local_particles.append(local_particle_file)
@@ -524,28 +527,53 @@ def Run_Pipeline(args):
         template_mesh = project_path + subjects[reference_index].get_groomed_filenames()[0]
         template_particles = project_path + subjects[reference_index].get_local_particle_filenames()[0]
 
+        # Get the test output directory from the predictions path
+        test_out_dir = os.path.dirname(local_test_prediction_dir.rstrip('/')) + '/'
         mean_dist = DeepSSMUtils.analyzeMeshDistance(predicted_test_local_particles, test_mesh_files,
                                                      template_particles, template_mesh, test_out_dir,
                                                      planes=test_planes)
         print("Test mean mesh surface-to-surface distance: " + str(mean_dist))
 
-        DeepSSMUtils.process_test_predictions(project, config_file)
-        
+        final_mean_dist = DeepSSMUtils.process_test_predictions(project, config_file)
+
         # If tiny test or verify, check results and exit
-        check_results(args, mean_dist)
+        check_results(args, final_mean_dist, output_directory)
 
         open(status_dir + "step_12.txt", 'w').close()
 
     print("All steps complete")
 
 
 # Verification
-def check_results(args, mean_dist):
+def check_results(args, mean_dist, output_directory):
     if args.tiny_test:
         print("\nVerifying use case results.")
-        if not math.isclose(mean_dist, 10, rel_tol=1):
-            print("Test failed.")
-            exit(-1)
+
+        exact_check_file = output_directory + "exact_check_value.txt"
+
+        # Exact check for refactoring verification (platform-specific)
+        if args.exact_check == "save":
+            with open(exact_check_file, "w") as f:
+                f.write(str(mean_dist))
+            print(f"Saved exact check value to: {exact_check_file}")
+            print(f"Value: {mean_dist}")
+        elif args.exact_check == "verify":
+            if not os.path.exists(exact_check_file):
+                print(f"Error: No saved value found at {exact_check_file}")
+                print("Run with --exact_check save first to create baseline.")
+                exit(-1)
+            with open(exact_check_file, "r") as f:
+                expected_mean_dist = float(f.read().strip())
+            if mean_dist != expected_mean_dist:
+                print(f"Exact check failed: expected {expected_mean_dist}, got {mean_dist}")
+                exit(-1)
+            print(f"Exact check passed: {mean_dist}")
+        else:
+            # Relaxed check for CI/cross-platform
+            if not math.isclose(mean_dist, 10, rel_tol=1):
+                print("Test failed.")
+                exit(-1)
+
         print("Done with test, verification succeeded.")
         exit(0)
     else:

Libs/Application/DeepSSM/DeepSSMJob.cpp

@@ -241,11 +241,12 @@ void DeepSSMJob::run_augmentation() {
   py::module py_deep_ssm_utils = py::module::import("DeepSSMUtils");
   py::object run_data_aug = py_deep_ssm_utils.attr("run_data_augmentation");
 
+  int processes = aug_processes_ > 0 ? aug_processes_ : QThread::idealThreadCount();
+
   int aug_dims = run_data_aug(project_, params.get_aug_num_samples(),
                               0 /* num dims, set to zero to allow percent variability to be used */,
                               params.get_aug_percent_variability(), sampler_type.toStdString(), 0 /* mixture_num */,
-                              QThread::idealThreadCount() /* processes */
-                              )
+                              processes)
                      .cast<int>();
 
   params.set_training_num_dims(aug_dims);
@@ -394,6 +395,12 @@ void DeepSSMJob::set_num_dataloader_workers(int num_workers) { num_dataloader_wo
 //---------------------------------------------------------------------------
 int DeepSSMJob::get_num_dataloader_workers() { return num_dataloader_workers_; }
 
+//---------------------------------------------------------------------------
+void DeepSSMJob::set_aug_processes(int processes) { aug_processes_ = processes; }
+
+//---------------------------------------------------------------------------
+int DeepSSMJob::get_aug_processes() { return aug_processes_; }
+
 //---------------------------------------------------------------------------
 void DeepSSMJob::update_prep_stage(PrepStep step) {
   /*

Libs/Application/DeepSSM/DeepSSMJob.h

@@ -55,6 +55,9 @@ class DeepSSMJob : public Job {
   void set_num_dataloader_workers(int num_workers);
   int get_num_dataloader_workers();
 
+  void set_aug_processes(int processes);
+  int get_aug_processes();
+
   void set_prep_step(DeepSSMJob::PrepStep step) {
     std::lock_guard<std::mutex> lock(mutex_);
     prep_step_ = step;
@@ -72,6 +75,7 @@ class DeepSSMJob : public Job {
   DeepSSMJob::PrepStep prep_step_{DeepSSMJob::NOT_STARTED};
 
   int num_dataloader_workers_{0};
+  int aug_processes_{0};
 
   // mutex
   std::mutex mutex_;

Libs/Application/Job/PythonWorker.cpp

@@ -207,11 +207,10 @@ bool PythonWorker::init() {
         path = QString::fromStdString(line);
       }
       file.close();
+      qputenv("PATH", path.toUtf8());
+      SW_LOG("Setting PATH for Python to: " + path.toStdString());
     }
 
-    qputenv("PATH", path.toUtf8());
-    SW_LOG("Setting PATH for Python to: " + path.toStdString());
-
     // Python 3.8+ requires explicit DLL directory registration
     // PATH environment variable is no longer used for DLL search
     SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_USER_DIRS);

Libs/Groom/Groom.cpp

@@ -186,7 +186,17 @@ bool Groom::image_pipeline(std::shared_ptr<Subject> subject, size_t domain) {
   std::string groomed_name = get_output_filename(original, DomainType::Image);
 
   if (params.get_convert_to_mesh()) {
+    // Use isovalue 0.0 for distance transforms (the zero level set is the surface)
     Mesh mesh = image.toMesh(0.0);
+    if (mesh.numPoints() == 0) {
+      throw std::runtime_error("Empty mesh generated from segmentation - segmentation may have no valid data");
+    }
+    // Check for valid cells
+    auto poly_data = mesh.getVTKMesh();
+    if (poly_data->GetNumberOfCells() == 0) {
+      throw std::runtime_error("Mesh has no cells - segmentation may have no valid surface");
+    }
+    SW_DEBUG("Mesh after toMesh: {} points, {} cells", poly_data->GetNumberOfPoints(), poly_data->GetNumberOfCells());
     run_mesh_pipeline(mesh, params, original);
     groomed_name = get_output_filename(original, DomainType::Mesh);
     // save the groomed mesh
@@ -239,6 +249,9 @@ bool Groom::run_image_pipeline(Image& image, GroomParameters params) {
   // crop
   if (params.get_crop()) {
     PhysicalRegion region = image.physicalBoundingBox(0.5);
+    if (!region.valid()) {
+      throw std::runtime_error("Empty segmentation - no voxels found above threshold for cropping");
+    }
     image.crop(region);
     increment_progress();
   }
@@ -1336,7 +1349,7 @@ std::vector<std::vector<double>> Groom::get_icp_transforms(const std::vector<Mes
       matrix->Identity();
 
       Mesh source = meshes[i];
-      if (source.getVTKMesh()->GetNumberOfPoints() != 0) {
+      if (source.getVTKMesh()->GetNumberOfPoints() != 0 && reference.getVTKMesh()->GetNumberOfPoints() != 0) {
         // create copies for thread safety
         auto poly_data1 = vtkSmartPointer<vtkPolyData>::New();
         poly_data1->DeepCopy(source.getVTKMesh());