Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/free-drive-option' into free-dri…
Browse files Browse the repository at this point in the history
…ve-option
  • Loading branch information
v-malheiro committed Sep 24, 2024
2 parents ecf9bbe + 2c88617 commit cf01dc9
Showing 7 changed files with 193 additions and 17 deletions.
39 changes: 39 additions & 0 deletions app.py
Original file line number Diff line number Diff line change
@@ -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_cranioplasty(args)
check_for_export(args)

return True
@@ -408,6 +411,42 @@ def use_cmd_optargs(args):
return False


def check_for_cranioplasty(args):
import invesalius.constants as const
from invesalius.i18n import tr as _

surface_options = {
"method": {
"algorithm": "Default",
"options": {},
},
"options": {
"index": 0,
"name": "",
"quality": _("Optimal *"),
"fill": False,
"keep_largest": False,
"overwrite": False,
},
}

if args.cranioplasty:
from invesalius.data import slice_
from invesalius.project import Project

# create cranium mask
Publisher.sendMessage("Update threshold limits", threshold_range=(226, 3071))
Publisher.sendMessage("Appy threshold all slices")

# create implant mask
Publisher.sendMessage("Create implant for cranioplasty")

# convert masks to surfaces and exports them.
Publisher.sendMessage(
"Export all surfaces separately", folder="./", filetype=const.FILETYPE_STL
)


def sanitize(text):
text = str(text).strip().replace(" ", "_")
return re.sub(r"(?u)[^-\w.]", "", text)
4 changes: 4 additions & 0 deletions invesalius/control.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions invesalius/data/slice_.py
Original file line number Diff line number Diff line change
@@ -254,6 +254,7 @@ def __bind_events(self) -> None:
Publisher.subscribe(self._fill_holes_auto, "Fill holes automatically")

Publisher.subscribe(self._set_interpolation_method, "Set interpolation method")
Publisher.subscribe(self.do_threshold_to_all_slices, "Appy threshold all slices")

def GetMaxSliceNumber(self, orientation: str) -> int:
shape: Tuple[int, int, int] = self.matrix.shape
56 changes: 56 additions & 0 deletions invesalius/data/surface.py
Original file line number Diff line number Diff line change
@@ -228,6 +228,7 @@ def __bind_events(self):
Publisher.subscribe(self.UpdateConvertToInvFlag, "Update convert_to_inv flag")

Publisher.subscribe(self.CreateSurfaceFromPolydata, "Create surface from polydata")
Publisher.subscribe(self.export_all_surfaces_separately, "Export all surfaces separately")

def OnDuplicate(self, surface_indexes):
proj = prj.Project()
@@ -1206,6 +1207,61 @@ def OnExportSurface(self, filename, filetype, convert_to_world=False):
dlg.Destroy()
os.remove(temp_file)

def export_all_surfaces_separately(self, folder, filetype):
import invesalius.data.slice_ as slc

if filetype in (
const.FILETYPE_STL,
const.FILETYPE_VTP,
const.FILETYPE_PLY,
const.FILETYPE_STL_ASCII,
):
proj = prj.Project()

for index in list(proj.mask_dict.keys()):
if index == 1:
algorithm = "Context aware smoothing"
else:
algorithm = "Default"
surface_parameters = {
"method": {
"algorithm": algorithm,
"options": {},
},
"options": {
"index": index,
"name": "",
"quality": _("Optimal *"),
"fill": False,
"keep_largest": True,
"overwrite": False,
},
}

mask = proj.mask_dict[index]
print(mask.matrix.min(), mask.matrix.max(), mask.name)
slice_ = slc.Slice()

Publisher.sendMessage(
"Create surface",
slice_=slice_,
mask=mask,
surface_parameters=surface_parameters,
)

# hide all surfaces
for index in proj.surface_dict:
proj.surface_dict[index].is_shown = False

# Displays one surface at a time and export
for index in proj.surface_dict.keys():
print(proj.surface_dict[index].name)
proj.surface_dict[index].is_shown = True
self._export_surface(
os.path.join(folder, str(index) + ".stl"), const.FILETYPE_STL, False
)
proj.surface_dict[index].is_shown = False

def _export_surface(self, filename, filetype, convert_to_world):
if filetype in (
const.FILETYPE_STL,
10 changes: 1 addition & 9 deletions invesalius/project.py
Original file line number Diff line number Diff line change
@@ -518,17 +518,9 @@ def Extract(filename: Union[str, bytes, os.PathLike], folder: Union[str, bytes,
os.mkdir(os.path.join(folder, idir))
filelist = []
for t in tar.getmembers():
fsrc = tar.extractfile(t)
if fsrc is None:
raise Exception("Error extracting file")
tar.extract(t, path=folder, filter=tarfile.tar_filter)
fname = os.path.join(folder, decode(t.name, "utf-8"))
fdst = open(fname, "wb")
shutil.copyfileobj(fsrc, fdst)
filelist.append(fname)
fsrc.close()
fdst.close()
del fsrc
del fdst
tar.close()
return filelist

78 changes: 70 additions & 8 deletions invesalius/segmentation/deep_learning/segment.py
Original file line number Diff line number Diff line change
@@ -1,33 +1,95 @@
import itertools
import multiprocessing
import os
import pathlib
import sys
import tempfile
import traceback
from typing import Generator, Tuple

import numpy as np
from skimage.transform import resize
from vtkmodules.vtkIOXML import vtkXMLImageDataWriter

import invesalius.data.slice_ as slc
from invesalius import inv_paths
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 gen_patches(image, patch_size, overlap):
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]]


def gen_patches(
image: np.ndarray, patch_size: int, overlap: int
) -> Generator[Tuple[float, np.ndarray, patch_type], None, None]:
overlap = int(patch_size * overlap / 100)
sz, sy, sx = image.shape
i_cuts = list(
itertools.product(
range(0, sz, patch_size - overlap),
range(0, sy, patch_size - overlap),
range(0, sx, patch_size - overlap),
)
)
slices_x = [i for i in range(0, sx, patch_size - overlap) if i + patch_size <= sx]
if not slices_x:
slices_x.append(0)
elif slices_x[-1] + patch_size < sx:
slices_x.append(sx - patch_size)
slices_y = [i for i in range(0, sy, patch_size - overlap) if i + patch_size <= sy]
if not slices_y:
slices_y.append(0)
elif slices_y[-1] + patch_size < sy:
slices_y.append(sy - patch_size)
slices_z = [i for i in range(0, sz, patch_size - overlap) if i + patch_size <= sz]
if not slices_z:
slices_z.append(0)
elif slices_z[-1] + patch_size < sz:
slices_z.append(sz - patch_size)
i_cuts = list(itertools.product(slices_z, slices_y, slices_x))

sub_image = np.empty(shape=(patch_size, patch_size, patch_size), dtype="float32")
for idx, (iz, iy, ix) in enumerate(i_cuts):
sub_image[:] = 0
22 changes: 22 additions & 0 deletions invesalius/segmentation/deep_learning/utils.py
Original file line number Diff line number Diff line change
@@ -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"):

0 comments on commit cf01dc9

Please sign in to comment.