Make extraction of the largest connected component optional (#305)

Author: eigenvivek

diffdrr/visualization.py

@@ -181,6 +181,7 @@ def drr_to_mesh(
     subject: Subject,  # torchio.Subject with a `volume` attribute
     method: str,  # Either `surface_nets` or `marching_cubes`
     threshold: float = 150,  # Min value for marching cubes (Hounsfield units)
+    extract_largest: bool = True,  # Extract the largest connected component from the mesh
     verbose: bool = True,  # Display progress bars for mesh processing steps
 ):
     """
@@ -190,7 +191,7 @@ def drr_to_mesh(
 
     The mesh processing steps are:
 
-    1. Keep only largest connected components
+    1. Keep only largest connected components (optional)
     2. Smooth
     3. Decimate (if `method=="marching_cubes"`)
     4. Fill any holes
@@ -229,7 +230,8 @@ def drr_to_mesh(
     mesh = mesh.transform(subject.volume.affine.squeeze())
 
     # Preprocess the mesh
-    mesh.extract_largest(inplace=True, progress_bar=verbose)
+    if extract_largest:
+        mesh.extract_largest(inplace=True, progress_bar=verbose)
     mesh.point_data.clear()
     mesh.cell_data.clear()
     mesh.smooth_taubin(

notebooks/api/04_visualization.ipynb

@@ -286,6 +286,7 @@
     "    subject: Subject,  # torchio.Subject with a `volume` attribute\n",
     "    method: str,  # Either `surface_nets` or `marching_cubes`\n",
     "    threshold: float = 150,  # Min value for marching cubes (Hounsfield units)\n",
+    "    extract_largest: bool = True,  # Extract the largest connected component from the mesh\n",
     "    verbose: bool = True,  # Display progress bars for mesh processing steps\n",
     "):\n",
     "    \"\"\"\n",
@@ -295,7 +296,7 @@
     "\n",
     "    The mesh processing steps are:\n",
     "\n",
-    "    1. Keep only largest connected components\n",
+    "    1. Keep only largest connected components (optional)\n",
     "    2. Smooth\n",
     "    3. Decimate (if `method==\"marching_cubes\"`)\n",
     "    4. Fill any holes\n",
@@ -334,7 +335,8 @@
     "    mesh = mesh.transform(subject.volume.affine.squeeze())\n",
     "\n",
     "    # Preprocess the mesh\n",
-    "    mesh.extract_largest(inplace=True, progress_bar=verbose)\n",
+    "    if extract_largest:\n",
+    "        mesh.extract_largest(inplace=True, progress_bar=verbose)\n",
     "    mesh.point_data.clear()\n",
     "    mesh.cell_data.clear()\n",
     "    mesh.smooth_taubin(\n",