Added command line to generate cranioplasty implant

app.py

@@ -345,6 +345,8 @@ def parse_command_line():
         help="Debug navigated TMS E-field computation",
     )
 
+    parser.add_argument("--cranioplasty", help="Creates an AI-based cranioplasty implant.")
+
     args = parser.parse_args()
     return args
 
@@ -359,6 +361,7 @@ def use_cmd_optargs(args):
             Publisher.sendMessage("Save project", filepath=os.path.abspath(args.save))
             exit(0)
 
+        check_for_segmentation(args)
         check_for_export(args)
 
         return True
@@ -408,6 +411,11 @@ def use_cmd_optargs(args):
     return False
 
 
+def check_for_segmentation(args):
+    if args.cranioplasty:
+        Publisher.sendMessage("Create implant for cranioplasty")
+
+
 def sanitize(text):
     text = str(text).strip().replace(" ", "_")
     return re.sub(r"(?u)[^-\w.]", "", text)

invesalius/control.py

@@ -45,6 +45,7 @@
 from invesalius import inv_paths, plugins
 from invesalius.i18n import tr as _
 from invesalius.pubsub import pub as Publisher
+from invesalius.segmentation.deep_learning import segment
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -124,6 +125,9 @@ def __bind_events(self) -> None:
 
         Publisher.subscribe(self.LoadProject, "Load project data")
 
+        # for call cranioplasty implant by command line
+        Publisher.subscribe(segment.run_cranioplasty_implant, "Create implant for cranioplasty")
+
     def SetBitmapSpacing(self, spacing: Tuple[float, float, float]) -> None:
         proj = prj.Project()
         proj.spacing = spacing

invesalius/segmentation/deep_learning/segment.py

@@ -16,12 +16,55 @@
 from invesalius.data import imagedata_utils
 from invesalius.data.converters import to_vtk
 from invesalius.net.utils import download_url_to_file
+from invesalius.pubsub import pub as Publisher
 from invesalius.utils import new_name_by_pattern
 
 from . import utils
 
 SIZE = 48
 
+
+def run_cranioplasty_implant():
+    """
+    This function was created to allow the creation of implants for
+    cranioplasty to be called by command line.
+    """
+    image = slc.Slice().matrix
+    backend = "pytorch"
+    device_id = list(utils.get_torch_devices().values())[0]
+    apply_wwwl = False
+    create_new_mask = True
+    use_gpu = True
+    resize_by_spacing = True
+    window_width = slc.Slice().window_width
+    window_level = slc.Slice().window_level
+    overlap = 50
+    patch_size = 480
+    method = 0  # binary
+
+    seg = ImplantCTSegmentProcess
+
+    ps = seg(
+        image,
+        create_new_mask,
+        backend,
+        device_id,
+        use_gpu,
+        overlap,
+        apply_wwwl,
+        window_width,
+        window_level,
+        method=method,
+        patch_size=patch_size,
+        resize_by_spacing=True,
+        image_spacing=slc.Slice().spacing,
+    )
+    ps._run_segmentation()
+    ps.apply_segment_threshold(0.75)
+    slc.Slice().discard_all_buffers()
+    Publisher.sendMessage("Reload actual slice")
+
+
 patch_type = Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]
 
 

invesalius/segmentation/deep_learning/utils.py

@@ -3,6 +3,28 @@
 import sys
 
 
+def get_torch_devices():
+    TORCH_DEVICES = {}
+
+    try:
+        import torch
+
+        HAS_TORCH = True
+    except ImportError:
+        HAS_TORCH = False
+
+    if HAS_TORCH:
+        TORCH_DEVICES = {}
+        if torch.cuda.is_available():
+            for i in range(torch.cuda.device_count()):
+                name = torch.cuda.get_device_name()
+                device_id = f"cuda:{i}"
+                TORCH_DEVICES[name] = device_id
+        TORCH_DEVICES["CPU"] = "cpu"
+
+    return TORCH_DEVICES
+
+
 def prepare_plaidml():
     # Linux if installed plaidml with pip3 install --user
     if sys.platform.startswith("linux"):