From 00a69e11e800995e10daeeb7641574471b623177 Mon Sep 17 00:00:00 2001 From: Ma Nan Date: Tue, 9 Apr 2024 22:19:38 +0800 Subject: [PATCH] init --- .gitattributes | 1 + .gitignore | 132 ++++++++ mednodes/__init__.py | 33 ++ mednodes/assets/Nodes.blend | 3 + mednodes/auto_load.py | 155 +++++++++ mednodes/custom_nodes/__init__.py | 1 + mednodes/custom_nodes/node_menu.py | 131 ++++++++ mednodes/custom_nodes/nodes.py | 164 ++++++++++ mednodes/dicom.py | 200 ++++++++++++ mednodes/external_package/__init__.py | 1 + mednodes/external_package/package.py | 390 +++++++++++++++++++++++ mednodes/external_package/preferences.py | 72 +++++ mednodes/nodes.py | 75 +++++ mednodes/preferences.py | 23 ++ mednodes/requirements.txt | 1 + 15 files changed, 1382 insertions(+) create mode 100644 .gitattributes create mode 100644 .gitignore create mode 100644 mednodes/__init__.py create mode 100644 mednodes/assets/Nodes.blend create mode 100644 mednodes/auto_load.py create mode 100644 mednodes/custom_nodes/__init__.py create mode 100644 mednodes/custom_nodes/node_menu.py create mode 100644 mednodes/custom_nodes/nodes.py create mode 100644 mednodes/dicom.py create mode 100644 mednodes/external_package/__init__.py create mode 100644 mednodes/external_package/package.py create mode 100644 mednodes/external_package/preferences.py create mode 100644 mednodes/nodes.py create mode 100644 mednodes/preferences.py create mode 100644 mednodes/requirements.txt diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..80a1944 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.blend filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c41ab26 --- /dev/null +++ b/.gitignore @@ -0,0 +1,132 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +pythonlib* + +*.blend1 + +blendcache_* + +*.ipynb \ No newline at end of file diff --git a/mednodes/__init__.py b/mednodes/__init__.py new file mode 100644 index 0000000..a27affd --- /dev/null +++ b/mednodes/__init__.py @@ -0,0 +1,33 @@ +from .nodes import custom_nodes +from . import auto_load +from .dicom import MEDNODES_add_topbar_menu +import bpy + + +bl_info = { + "name": "MedNodes", + "author": "Ma Nan", + "description": "", + "blender": (4, 0, 0), + "version": (0, 1, 0), + "location": "File -> Import", + "warning": "", + "category": "Import-Export" +} + +auto_load.init() + + +def register(): + auto_load.register() + custom_nodes.register() + bpy.types.TOPBAR_MT_file_import.prepend(MEDNODES_add_topbar_menu) + + +def unregister(): + try: + bpy.types.TOPBAR_MT_file_import.remove(MEDNODES_add_topbar_menu) + custom_nodes.unregister() + auto_load.unregister() + except RuntimeError: + pass diff --git a/mednodes/assets/Nodes.blend b/mednodes/assets/Nodes.blend new file mode 100644 index 0000000..1933a41 --- /dev/null +++ b/mednodes/assets/Nodes.blend @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fad3bd63d4c86693814b67c9efe09efc445415fa9e67ff7ab9eff20e51c70a53 +size 1769197 diff --git a/mednodes/auto_load.py b/mednodes/auto_load.py new file mode 100644 index 0000000..86fc69c --- /dev/null +++ b/mednodes/auto_load.py @@ -0,0 +1,155 @@ +import bpy +import typing +import inspect +import pkgutil +import importlib +from pathlib import Path + +__all__ = ( + "init", + "register", + "unregister", +) + +blender_version = bpy.app.version + +modules = None +ordered_classes = None + +def init(): + global modules + global ordered_classes + + modules = get_all_submodules(Path(__file__).parent) + ordered_classes = get_ordered_classes_to_register(modules) + +def register(): + for cls in ordered_classes: + bpy.utils.register_class(cls) + + for module in modules: + if module.__name__ == __name__: + continue + if hasattr(module, "register"): + module.register() + +def unregister(): + for cls in reversed(ordered_classes): + bpy.utils.unregister_class(cls) + + for module in modules: + if module.__name__ == __name__: + continue + if hasattr(module, "unregister"): + module.unregister() + + +# Import modules +################################################# + +def get_all_submodules(directory): + return list(iter_submodules(directory, directory.name)) + +def iter_submodules(path, package_name): + for name in sorted(iter_submodule_names(path)): + yield importlib.import_module("." + name, package_name) + +def iter_submodule_names(path, root=""): + for _, module_name, is_package in pkgutil.iter_modules([str(path)]): + if is_package: + sub_path = path / module_name + sub_root = root + module_name + "." + yield from iter_submodule_names(sub_path, sub_root) + else: + yield root + module_name + + +# Find classes to register +################################################# + +def get_ordered_classes_to_register(modules): + return toposort(get_register_deps_dict(modules)) + +def get_register_deps_dict(modules): + my_classes = set(iter_my_classes(modules)) + my_classes_by_idname = {cls.bl_idname : cls for cls in my_classes if hasattr(cls, "bl_idname")} + + deps_dict = {} + for cls in my_classes: + deps_dict[cls] = set(iter_my_register_deps(cls, my_classes, my_classes_by_idname)) + return deps_dict + +def iter_my_register_deps(cls, my_classes, my_classes_by_idname): + yield from iter_my_deps_from_annotations(cls, my_classes) + yield from iter_my_deps_from_parent_id(cls, my_classes_by_idname) + +def iter_my_deps_from_annotations(cls, my_classes): + for value in typing.get_type_hints(cls, {}, {}).values(): + dependency = get_dependency_from_annotation(value) + if dependency is not None: + if dependency in my_classes: + yield dependency + +def get_dependency_from_annotation(value): + if blender_version >= (2, 93): + if isinstance(value, bpy.props._PropertyDeferred): + return value.keywords.get("type") + else: + if isinstance(value, tuple) and len(value) == 2: + if value[0] in (bpy.props.PointerProperty, bpy.props.CollectionProperty): + return value[1]["type"] + return None + +def iter_my_deps_from_parent_id(cls, my_classes_by_idname): + if bpy.types.Panel in cls.__bases__: + parent_idname = getattr(cls, "bl_parent_id", None) + if parent_idname is not None: + parent_cls = my_classes_by_idname.get(parent_idname) + if parent_cls is not None: + yield parent_cls + +def iter_my_classes(modules): + base_types = get_register_base_types() + for cls in get_classes_in_modules(modules): + if any(base in base_types for base in cls.__bases__): + if not getattr(cls, "is_registered", False): + yield cls + +def get_classes_in_modules(modules): + classes = set() + for module in modules: + for cls in iter_classes_in_module(module): + classes.add(cls) + return classes + +def iter_classes_in_module(module): + for value in module.__dict__.values(): + if inspect.isclass(value): + yield value + +def get_register_base_types(): + return set(getattr(bpy.types, name) for name in [ + "Panel", "Operator", "PropertyGroup", + "AddonPreferences", "Header", "Menu", + "Node", "NodeSocket", "NodeTree", + "UIList", "RenderEngine", + "Gizmo", "GizmoGroup", + ]) + + +# Find order to register to solve dependencies +################################################# + +def toposort(deps_dict): + sorted_list = [] + sorted_values = set() + while len(deps_dict) > 0: + unsorted = [] + for value, deps in deps_dict.items(): + if len(deps) == 0: + sorted_list.append(value) + sorted_values.add(value) + else: + unsorted.append(value) + deps_dict = {value : deps_dict[value] - sorted_values for value in unsorted} + return sorted_list diff --git a/mednodes/custom_nodes/__init__.py b/mednodes/custom_nodes/__init__.py new file mode 100644 index 0000000..43d1830 --- /dev/null +++ b/mednodes/custom_nodes/__init__.py @@ -0,0 +1 @@ +from .node_menu import CustomNodes \ No newline at end of file diff --git a/mednodes/custom_nodes/node_menu.py b/mednodes/custom_nodes/node_menu.py new file mode 100644 index 0000000..9280b95 --- /dev/null +++ b/mednodes/custom_nodes/node_menu.py @@ -0,0 +1,131 @@ +import bpy +from pathlib import Path +from .nodes import AddCustomNode + + +class CustomNodes(): + def __init__( + self, + menu_items, + nodes_file, + root_label='CustomNodes', + root_icon=88 + ) -> None: + if not Path(nodes_file).is_file(): + raise FileNotFoundError(Path(nodes_file).resolve().as_posix()) + + self.menu_items = menu_items + self.nodes_file = Path(nodes_file).resolve().as_posix() + + self.root_label = root_label + self.root_icon = root_icon + + menu_classes = [] + self._create_menu_class( + items=menu_items, + label=root_label, + menu_classes=menu_classes + ) + self.menu_classes = menu_classes + + idname = f"CUSTOMNODES_MT_NODES__{root_label.replace(' ', '').upper()}" + + def add_node_menu(self, context): + if ('GeometryNodeTree' == bpy.context.area.spaces[0].tree_type): + layout = self.layout + layout.separator() + layout.menu(idname, icon_value=root_icon) + + self.add_node_menu = add_node_menu + + def _create_menu_class(self, menu_classes, items, label='CustomNodes', icon=0, idname_namespace=None): + nodes_file = self.nodes_file + idname_namespace = idname_namespace or "CUSTOMNODES_MT_NODES_" + idname = f"{idname_namespace}_{label.replace(' ', '').upper()}" + + # create submenu class if item is menu. + for item in items: + item_items = item.get('items') if item != 'separator' else None + if item_items: + menu_class = self._create_menu_class( + menu_classes=menu_classes, + items=item_items, + label=item.get('label') or 'CustomNodes', + icon=item.get('icon') or 0, + idname_namespace=idname + ) + item['menu_class'] = menu_class + + # create menu class + class Menu(bpy.types.Menu): + bl_idname = idname + bl_label = label + + def draw(self, context): + layout = self.layout + + for item in items: + # print(item) + if item == "separator": + layout.separator() + elif item.get('menu_class'): + layout.menu( + item.get('menu_class').bl_idname, + icon_value=item.get('icon') or 0 + ) + else: + op = layout.operator( + 'customnodes.add_custom_node', text=item.get('label')) + op.nodes_file = nodes_file + op.node_type = item['node_type'] + op.node_label = item.get('label') or "" + op.node_description = item.get( + 'node_description') or "Add Custom Node." + op.node_driver = item.get('node_driver') or "" + + menu_classes.append(Menu) + return Menu + + def _find_item(self, found_items, menu_items, node_type: str): + + for item in menu_items: + if item.get("node_type") == node_type: + found_items.append(item) + + item_items = item.get('items') if item != 'separator' else None + if item_items: + self._find_item(found_items, item_items, node_type) + + def find_item(self, node_type: str): + found_items = [] + self._find_item(found_items, self.menu_items, node_type) + return found_items[0] if len(found_items) > 0 else None + + def add_node(self, node_tree, node_type: str): + item = self.find_item(node_type) + if item: + op = AddCustomNode() + op.nodes_file = self.nodes_file + op.node_type = item['node_type'] + op.node_label = item.get('label') or "" + op.node_description = item.get( + 'node_description') or "Add Custom Node." + op.node_driver = item.get('node_driver') or "" + return op.add_node(node_tree) + else: + raise RuntimeError("No custom node type found.") + + def register(self): + for cls in self.menu_classes: + bpy.utils.register_class(cls) + + bpy.types.NODE_MT_add.append(self.add_node_menu) + + def unregister(self): + try: + for cls in reversed(self.menu_classes): + bpy.utils.unregister_class(cls) + + bpy.types.NODE_MT_add.remove(self.add_node_menu) + except RuntimeError: + pass diff --git a/mednodes/custom_nodes/nodes.py b/mednodes/custom_nodes/nodes.py new file mode 100644 index 0000000..a15cf4c --- /dev/null +++ b/mednodes/custom_nodes/nodes.py @@ -0,0 +1,164 @@ +import bpy + + +def add_driver(target_prop, var_sources, expression): + + driver = target_prop.driver_add("default_value") + is_vector = isinstance(driver, list) + drivers = driver if is_vector else [driver] + + for i, driver in enumerate(drivers): + for j, var_source in enumerate(var_sources): + + source = var_source['source'] + prop = var_source['prop'] + + var = driver.driver.variables.new() + var.name = f"var{j}" + + var.targets[0].id_type = source.id_type + var.targets[0].id = source + var.targets[0].data_path = f'["{prop}"][{i}]'\ + if is_vector else f'["{prop}"]' + + driver.driver.expression = expression + + +def add_direct_driver(target, target_prop, source, source_prop): + target_prop = target.inputs.get(target_prop) + drivers = [ + { + "source": source, + "prop": source_prop + } + ] + expression = "var0" + add_driver(target_prop, drivers, expression) + + +class AddCustomNode(): + + nodes_file: bpy.props.StringProperty( + name="nodes_file", + subtype='FILE_PATH', + default="" + ) # type: ignore + + node_type: bpy.props.StringProperty( + name='node_type', + description='', + default='', + subtype='NONE', + maxlen=0 + ) # type: ignore + + node_label: bpy.props.StringProperty( + name='node_label', + default='' + ) # type: ignore + + node_driver: bpy.props.StringProperty( + name='node_driver', + default='' + ) # type: ignore + + node_description: bpy.props.StringProperty( + name="node_description", + description="", + default="Add custom node group.", + subtype="NONE" + ) # type: ignore + + node_link: bpy.props.BoolProperty( + name='node_link', + default=True + ) # type: ignore + + @classmethod + def description(cls, context, properties): + return properties.node_description + + def assign_node_tree(self, node): + node.node_tree = bpy.data.node_groups[self.node_type] + + node.width = 200.0 + node.label = self.node_label or self.node_type + node.name = self.node_type + + if self.node_driver: + source_prop = self.node_driver.split(":")[0] + target_prop = self.node_driver.split(":")[1] + add_direct_driver( + target=node, + target_prop=target_prop, + source=bpy.context.active_object, + source_prop=source_prop + ) + + return node + + def append_node_tree(self, node_type): + # try to get node from current file if exists + node_tree = bpy.data.node_groups.get(node_type) + # if not exists, get it from asset file. + if not node_tree: + bpy.ops.wm.append( + 'EXEC_DEFAULT', + directory=f"{self.nodes_file}/NodeTree", + filename=node_type, + link=False, + use_recursive=True + ) + + node_tree = bpy.data.node_groups.get(node_type) + if node_tree: + # self.recursive_append_material(node_tree) + return node_tree + else: + raise RuntimeError('No custom node found') + + def recursive_append_material(self, node_tree): + for child in node_tree.nodes: + material_socket = child.inputs.get('Material') + if material_socket: + print(material_socket.default_value) + try: + self.recursive_append_material(child.node_tree) + except: + ... + + def add_node(self, node_tree): + self.append_node_tree(self.node_type) + + node = node_tree.new("GeometryNodeGroup") + + self.assign_node_tree(node) + + return node + + +class CUSTOMNODES_OT_Add_Custom_Node(bpy.types.Operator, AddCustomNode): + bl_idname = "customnodes.add_custom_node" + bl_label = "Add Custom Node" + bl_options = {"REGISTER", "UNDO"} + + def execute(self, context): + self.append_node_tree(self.node_type) + + # intended to be called upon button press in the node tree + prev_context = bpy.context.area.type + bpy.context.area.type = 'NODE_EDITOR' + # actually invoke the operator to add a node to the current node tree + # use_transform=True ensures it appears where the user's mouse is and is currently + # being moved so the user can place it where they wish + bpy.ops.node.add_node( + 'INVOKE_DEFAULT', + type='GeometryNodeGroup', + use_transform=True + ) + bpy.context.area.type = prev_context + node = bpy.context.active_node + + self.assign_node_tree(node) + + return {"FINISHED"} diff --git a/mednodes/dicom.py b/mednodes/dicom.py new file mode 100644 index 0000000..0871763 --- /dev/null +++ b/mednodes/dicom.py @@ -0,0 +1,200 @@ +import bpy +from bpy_extras.io_utils import axis_conversion, ImportHelper +import pyopenvdb as vdb +import numpy as np +from pathlib import Path +from uuid import uuid4 +import mathutils +from .nodes import custom_nodes + + +class ImportDICOM(bpy.types.Operator, ImportHelper): + bl_idname = "mednodes.import_dicom" + bl_label = "DICOM as VDB (.dcm)" + bl_description = "Load DICOM as VDB" + bl_options = {'UNDO'} + + filepath: bpy.props.StringProperty( + subtype="FILE_PATH" + ) # type: ignore + + global_scale: bpy.props.FloatProperty( + name="Scale", + soft_min=0.001, soft_max=100.0, + min=1e-6, max=1e6, + default=0.01, + ) # type: ignore + + do_resample: bpy.props.BoolProperty( + name="Uniform Spacing (Resample)", + default=True, + ) # type: ignore + + do_change_render_setting: bpy.props.BoolProperty( + name="Change Render Setting", + default=True, + ) # type: ignore + + def execute(self, context): + import SimpleITK as sitk + + preferences = context.preferences.addons[__package__].preferences + + file_path: Path = Path(self.filepath).resolve() + name = file_path.parent.stem + + reader = sitk.ImageSeriesReader() + dicom_names = reader.GetGDCMSeriesFileNames( + file_path.parent.as_posix()) + reader.SetFileNames(dicom_names) + image = reader.Execute() + + print("Spacing:", image.GetSpacing()) + print("Origin:", image.GetOrigin()) + print("Direction:", image.GetDirection()) + + image = sitk.DICOMOrient(image, 'LPS') + + size = image.GetSize() + spacing = image.GetSpacing() + origin = image.GetOrigin() + direction = image.GetDirection() + + print("Size:", size) + print("Spacing:", spacing) + print("Origin:", origin) + print("Direction:", direction) + + if self.do_resample: + resampled_spacing = (1, 1, 1) + resampled_size = ( + int(size[0]*spacing[0]), + int(size[1]*spacing[1]), + int(size[2]*spacing[2]), + ) + + print("Resampling...") + image = sitk.Resample( + image1=image, + size=resampled_size, + transform=sitk.Transform(), + interpolator=sitk.sitkLinear, + outputOrigin=origin, + outputSpacing=resampled_spacing, + outputDirection=direction, + defaultPixelValue=0, + outputPixelType=image.GetPixelID(), + ) + print("Resampled Size:", resampled_size) + + spacing = resampled_spacing + size = resampled_size + + array = sitk.GetArrayFromImage(image) + + # ITK indices, by convention, are [i,j,k] while NumPy indices are [k,j,i] + # https://www.slicer.org/wiki/Coordinate_systems + + # ITK Numpy 3D + # L (eft) i -> k -> x + # P (osterior) j -> j -> y + # S (uperior) k -> i -> z + + array = np.transpose(array) + value_max = float(np.max(array)) + value_min = float(np.min(array)) + + value_offset = 0.0 + if value_min < 0: + value_offset = -value_min + array = array + np.full_like(array, value_offset) + value_max = float(np.max(array)) + value_min = float(np.min(array)) + print("offset max value:", value_max) + print("offset min value:", value_min) + + # Build VDB + grid = vdb.FloatGrid() + grid.copyFromArray(array.copy()) + grid.name = "density" + + vdb_dirpath = Path(preferences.cache_dir, 'VDBs') + vdb_dirpath.mkdir(parents=True, exist_ok=True) + vdb_path = Path(vdb_dirpath, f"{uuid4()}.vdb") + vdb.write(vdb_path.as_posix(), grids=[grid]) + + # Read VDB + bpy.ops.object.volume_import( + filepath=vdb_path.as_posix(), align='WORLD', location=(0, 0, 0), scale=(1, 1, 1)) + + vdb_obj = bpy.context.active_object + + # Set props to VDB object + vdb_obj.name = name + vdb_obj.data.name = name + + vdb_obj['value_max'] = value_max + vdb_obj['value_min'] = value_min + vdb_obj['value_offset'] = value_offset + + vdb_obj['size'] = size + vdb_obj['spacing'] = spacing + vdb_obj['origin'] = origin + + # Make transformation + + # (S)uperior -Z -> Y + # (P)osterior -Y -> Z + axis_rot = axis_conversion( + from_forward='-Z', + from_up='-Y', + to_forward='Y', + to_up='Z' + ).to_4x4() + + mat_loc = mathutils.Matrix.Translation( + mathutils.Vector(origin) * self.global_scale) + mat_sca = mathutils.Matrix.Scale(self.global_scale, 4) + mat_rot = mathutils.Matrix( + np.array(direction).reshape((3, 3)) + ).to_4x4() + + vdb_obj.matrix_world = axis_rot @ mat_loc @ mat_rot @ mat_sca + + # Wrapper a Locator + bpy.ops.object.empty_add( + type='PLAIN_AXES', + align='WORLD', + location=(0, 0, 0), + scale=(1, 1, 1) + ) + loc_obj = bpy.context.active_object + loc_obj.name = f"LOC_{name}" + vdb_obj.parent = loc_obj + + # Create MedNodes to VDB object + bpy.context.view_layer.objects.active = vdb_obj + bpy.ops.node.new_geometry_nodes_modifier() + modifier = vdb_obj.modifiers[0] + nodes = modifier.node_group.nodes + links = modifier.node_group.links + + input_node = nodes.get("Group Input") + output_node = nodes.get("Group Output") + segment_node = custom_nodes.add_node(nodes, 'MedNodes_Segment') + + links.new(input_node.outputs[0], segment_node.inputs[0]) + links.new(segment_node.outputs[0], output_node.inputs[0]) + + # Change render setting for better result + if self.do_change_render_setting: + bpy.context.scene.render.engine = 'CYCLES' + bpy.context.scene.cycles.volume_bounces = 12 + bpy.context.scene.cycles.transparent_max_bounces = 64 + + return {'FINISHED'} + + +def MEDNODES_add_topbar_menu(self, context): + layout = self.layout + layout.operator(ImportDICOM.bl_idname) diff --git a/mednodes/external_package/__init__.py b/mednodes/external_package/__init__.py new file mode 100644 index 0000000..c00e7fb --- /dev/null +++ b/mednodes/external_package/__init__.py @@ -0,0 +1 @@ +from .preferences import ExternalPackagePreferences \ No newline at end of file diff --git a/mednodes/external_package/package.py b/mednodes/external_package/package.py new file mode 100644 index 0000000..36bb68a --- /dev/null +++ b/mednodes/external_package/package.py @@ -0,0 +1,390 @@ +""" +Handling installation of external python packages inside of Blender's bundled python. +""" +import subprocess +import sys +import logging +from importlib.metadata import version as get_version, PackageNotFoundError +import bpy +from pathlib import Path + +ADDON_NAME = __package__.split(".")[0] + +PYPI_MIRROR = { + # the original. + 'Default': '', + # two mirrors in China Mainland to help those poor victims under GFW. + 'BFSU (Beijing)': 'https://mirrors.bfsu.edu.cn/pypi/web/simple', + 'TUNA (Beijing)': 'https://pypi.tuna.tsinghua.edu.cn/simple', + # append more if necessary. +} + + +class InstallationError(Exception): + """ + Exception raised when there is an error installing a package. + + Attributes + ---------- + package_name : str + The name of the package that failed to install. + error_message : str + The error message returned by pip. + + """ + + def __init__(self, package_name, error_message): + self.package_name = package_name + self.error_message = error_message + super().__init__(f"Failed to install {package_name}: {error_message}") + + +class PackageInstaller(): + def __init__( + self, + pypi_mirror_provider='Default', + log_dir: str = None, + requirements_dir: str = None + ) -> None: + self.pypi_mirror_provider = pypi_mirror_provider + self.log_path: Path = Path(log_dir) if log_dir else \ + Path(Path.home(), '.externalpackage', 'logs') + self.requirements_path: Path = Path(requirements_dir) if requirements_dir else \ + Path(__file__).parent + + def start_logging(self, logfile_name: str = 'side-packages-install') -> logging.Logger: + """ + Configure and start logging to a file. + + Parameters + ---------- + logfile_name : str, optional + The name of the log file. Defaults to 'side-packages-install'. + + Returns + ------- + logging.Logger + A Logger object that can be used to write log messages. + + This function sets up a logging configuration with a specified log file name and logging level. + The log file will be created in the `ADDON_DIR/logs` directory. If the directory + does not exist, it will be created. The function returns a Logger object that can be used to + write log messages. + + """ + # Create the logs directory if it doesn't exist + self.log_path.mkdir(parents=True, exist_ok=True) + + # Set up logging configuration + logfile_path = Path(self.log_path, f"{logfile_name}.log") + logging.basicConfig(filename=logfile_path, level=logging.INFO) + + # Return logger object + return logging.getLogger() + + @property + def pypi_mirror_url(self) -> str: + """ + Process a PyPI mirror provider and return the corresponding URL. + + Parameters + ---------- + pypi_mirror_provider : str + The PyPI mirror provider to process. + + Returns + ------- + str + The URL of the PyPI mirror. + + Raises + ------ + ValueError + If the provided PyPI mirror provider is invalid. + + """ + if self.pypi_mirror_provider.startswith('https:'): + return self.pypi_mirror_provider + elif self.pypi_mirror_provider in PYPI_MIRROR.keys(): + return PYPI_MIRROR[self.pypi_mirror_provider] + else: + raise ValueError( + f"Invalid PyPI mirror provider: {self.pypi_mirror_provider}") + + @property + def packages(self) -> dict: + """ + Read a requirements file and extract package information into a dictionary. + + Parameters + ---------- + requirements : str, optional + The path to the requirements file. If not provided, the function looks for a `requirements.txt` + file in the same directory as the script. + + Returns + ------- + dict + A dictionary containing package information. Each element of the dictionary is a dictionary containing the package name, version, and description. + + Example + ------- + Given the following requirements file: + ```python + Flask==1.1.2 # A micro web framework for Python + pandas==1.2.3 # A fast, powerful, flexible, and easy-to-use data analysis and manipulation tool + numpy==1.20.1 # Fundamental package for scientific computing + ``` + The function would return the following dictionary: + ```python + [ + { + "package": "Flask", + "version": "1.1.2", + "description": "A micro web framework for Python" + }, + { + "package": "pandas", + "version": "1.2.3", + "description": "A fast, powerful, flexible, and easy-to-use data analysis and manipulation tool" + }, + { + "package": "numpy", + "version": "1.20.1", + "description": "Fundamental package for scientific computing" + } + ] + ``` + """ + requirements_filepath = self.requirements_path / "requirements.txt" + if requirements_filepath.is_file(): + with requirements_filepath.open('r') as f: + lines = f.read().splitlines() + packages = {} + for line in lines: + try: + package, description = line.split('#') + package_meta = package.split('==') + name = package_meta[0].strip() + packages[name] = { + "name": name, + "version": package_meta[1].strip(), + "description": description.strip() + } + except ValueError: + # Skip line if it doesn't have the expected format + pass + return packages + else: + raise FileNotFoundError(requirements_filepath) + + def is_installed(self, package_name: str) -> bool: + """ + Check if the specified package is installed and the version matches that specified + in the `requirements.txt` file. + + Parameters + ---------- + package : str + The name of the package to check. + + Returns + ------- + bool + True if the package is the current version, False otherwise. + + """ + package = self.packages.get(package_name) + try: + available_version = get_version(package['name']) + return available_version == package['version'] + except PackageNotFoundError: + return False + + def run_python(self, cmd_list: list = None, timeout: int = 600): + """ + Runs pip command using the specified command list and returns the command output. + + Parameters + ---------- + cmd_list : list, optional + List of pip commands to be executed. Defaults to None. + mirror_url : str, optional + URL of a package repository mirror to be used for the command. Defaults to ''. + timeout : int, optional + Time in seconds to wait for the command to complete. Defaults to 600. + + Returns + ------- + tuple + A tuple containing the command list, command return code, command standard output, + and command standard error. + + Example + ------- + Install numpy using pip and print the command output + ```python + cmd_list = ["-m", "pip", "install", "numpy"] + mirror_url = 'https://pypi.org/simple/' + cmd_output = run_python(cmd_list, timeout=300) + print(cmd_output) + ``` + + """ + + # path to python.exe + python_exe = Path(sys.executable).resolve() + + # build the command list + cmd_list = [python_exe] + cmd_list + + # add mirror to the command list if it's valid + if self.pypi_mirror_url and self.pypi_mirror_url.startswith('https'): + cmd_list += ['-i', self.pypi_mirror_url] + + log = self.start_logging() + log.info(f"Running Pip: '{cmd_list}'") + + # run the command and capture the output + result = subprocess.run(cmd_list, timeout=timeout, + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + + if result.returncode != 0: + log.error('Command failed: %s', cmd_list) + log.error('stdout: %s', result.stdout.decode()) + log.error('stderr: %s', result.stderr.decode()) + else: + log.info('Command succeeded: %s', cmd_list) + log.info('stdout: %s', result.stdout.decode()) + # return the command list, return code, stdout, and stderr as a tuple + return result + + def install_package(self, package: str) -> list: + """ + Install a Python package and its dependencies using pip. + + Parameters + ---------- + package : str + The name of the package to install. + pypi_mirror_provider : str, optional + The name/url of the PyPI mirror provider to use. Default is 'Default'. + + Returns + ------- + list + A list of tuples containing the command list, return code, stdout, and stderr + for each pip command run. + + Raises + ------ + ValueError + If the package name is not provided. + + Example + ------- + To install the package 'requests' from the PyPI mirror 'MyMirror', use: + ``` + install_package('requests', 'MyMirror') + ``` + + """ + if not package: + raise ValueError("Package name must be provided.") + + print(f"Installing {package}...") + print(f"Using PyPI mirror:\ + {self.pypi_mirror_provider} {self.pypi_mirror_url}") + + self.run_python(["-m", "ensurepip"]), + self.run_python(["-m", "pip", "install", "--upgrade", "pip"]) + result = self.run_python(["-m", "pip", "install", package]) + + return result + + def install_all_packages(self, pypi_mirror_provider: str = 'Default') -> list: + """ + Install all packages listed in the 'requirements.txt' file. + + Parameters + ---------- + pypi_mirror_provider : str, optional + The PyPI mirror to use for package installation. Defaults to 'Default', + which uses the official PyPI repository. + + Returns + ------- + list + A list of tuples containing the installation results for each package. + + Raises + ------ + InstallationError + If there is an error during package installation. + + Example + ------- + To install all packages listed in the 'requirements.txt' file, run the following command: + ``` + install_all_packages(pypi_mirror_provider='https://pypi.org/simple/') + ``` + + """ + results = [] + for package in self.packages.items(): + + try: + result = self.install_package( + package=f"{package.get('name')}=={package.get('version')}" + ) + results.append(result) + except InstallationError as e: + raise InstallationError( + f"Error installing package {package.get('name')}: {str(e)}") + return results + + +class EXTERNALPACKAGE_OT_Install_Package(bpy.types.Operator): + bl_idname = 'externalpackage.install_package' + bl_label = 'Install Given Python Package' + bl_options = {'REGISTER', 'INTERNAL'} + + package: bpy.props.StringProperty( + name='Python Package', + description='Python Package to Install' + ) # type: ignore + version: bpy.props.StringProperty( + name='Python Package', + description='Python Package to Install' + ) # type: ignore + + description: bpy.props.StringProperty( + name='Operator description', + default='Install specified python package.' + ) # type: ignore + + @classmethod + def description(cls, context, properties): + return properties.description + + def execute(self, context): + preferences = context.preferences.addons[ADDON_NAME].preferences + installer = PackageInstaller( + pypi_mirror_provider=preferences.pypi_mirror_provider, + log_dir=preferences.log_dir, + requirements_dir=preferences.requirements_dir + ) + + result = installer.install_package(f"{self.package}=={self.version}") + if result.returncode == 0 and installer.is_installed(self.package): + self.report( + {'INFO'}, + f"Successfully installed {self.package} v{self.version}" + ) + else: + self.report( + {'ERROR'}, + f"Error installing package. Please check the log files in \ + '{preferences.log_dir}'." + ) + return {'FINISHED'} diff --git a/mednodes/external_package/preferences.py b/mednodes/external_package/preferences.py new file mode 100644 index 0000000..357e2a2 --- /dev/null +++ b/mednodes/external_package/preferences.py @@ -0,0 +1,72 @@ +import bpy +from pathlib import Path +from .package import PackageInstaller, PYPI_MIRROR + +# Defines the preferences panel for the addon, which shows the buttons for +# installing and reinstalling the required python packages defined in 'requirements.txt' + + +def get_pypi_mirrors(self, context, edit_text): + return PYPI_MIRROR.keys() + + +class ExternalPackagePreferences(): + requirements_dir: bpy.props.StringProperty( + name="requirements.txt Directory", + subtype='DIR_PATH', + default='' + ) # type: ignore + + log_dir: bpy.props.StringProperty( + name="Python Log Directory", + subtype='DIR_PATH', + default='' + ) # type: ignore + + pypi_mirror_provider: bpy.props.StringProperty( + name='pypi_mirror_provider', + description='PyPI Mirror Provider', + options={'TEXTEDIT_UPDATE', 'LIBRARY_EDITABLE'}, + default='Default', + subtype='NONE', + search=get_pypi_mirrors, + ) # type: ignore + + def draw(self, context): + layout = self.layout + layout.label(text="Install the required packages.") + + col_main = layout.column(heading='', align=False) + row_import = col_main.row() + row_import.prop(self, 'pypi_mirror_provider', text='Set PyPI Mirror') + + installer = PackageInstaller( + pypi_mirror_provider=self.pypi_mirror_provider, + log_dir=self.log_dir, + requirements_dir=self.requirements_dir + ) + + for package in installer.packages.values(): + + name = package.get('name') + version = package.get('version') + description = package.get('description') + + if installer.is_installed(name): + row = layout.row() + row.label(text=f"{name} version {version} is installed.") + op = row.operator('externalpackage.install_package', + text=f'Reinstall {name}') + op.package = name + op.version = version + op.description = f'Reinstall {name}' + else: + row = layout.row(heading=f"Package: {name}") + col = row.column() + col.label(text=str(description)) + col = row.column() + op = col.operator('externalpackage.install_package', + text=f'Install {name}') + op.package = name + op.version = version + op.description = f'Install required python package: {name}' diff --git a/mednodes/nodes.py b/mednodes/nodes.py new file mode 100644 index 0000000..46a2ca9 --- /dev/null +++ b/mednodes/nodes.py @@ -0,0 +1,75 @@ +from .custom_nodes import CustomNodes + +MENU_ITEMS = [ + { + 'label': 'Present', + 'items': [ + { + 'label': 'Segment', + 'node_type': 'MedNodes_Segment', + 'node_description': '', + 'node_driver': 'value_offset:Offset' + } + ] + }, + { + 'label': 'Segment', + 'items': [ + { + 'label': 'Segment by Level', + 'node_type': 'MedNodes_SegmentByLevel', + 'node_description': '', + 'node_driver': 'value_offset:Offset' + }, + { + 'label': 'Segment by Range', + 'node_type': 'MedNodes_SegmentByRange', + 'node_description': '', + 'node_driver': 'value_offset:Offset' + }, + { + 'label': 'Segment by Layers', + 'node_type': 'MedNodes_SegmentsByLayers', + 'node_description': '', + 'node_driver': 'value_offset:Offset' + } + ] + }, + { + 'label': 'Shader', + 'items': [ + { + 'label': 'Segment Shader', + 'node_type': 'MedNodes_SegmentShader', + 'node_description': '' + } + ] + }, + { + 'label': 'Slicer', + 'items': [ + { + 'label': 'LPS Slicer', + 'node_type': 'MedNodes_LPS-Slicer', + 'node_description': '', + 'node_driver': 'origin:Origin' + }, + { + 'label': 'Axis Slicer', + 'node_type': 'MedNodes_AxisSlicer', + 'node_description': '', + }, + { + 'label': 'Box Slicer', + 'node_type': 'MedNodes_BoxSlicer', + 'node_description': '', + } + ] + }, +] + +custom_nodes = CustomNodes( + menu_items=MENU_ITEMS, + nodes_file="assets/Nodes.blend", + root_label='MedNodes' +) diff --git a/mednodes/preferences.py b/mednodes/preferences.py new file mode 100644 index 0000000..7b1cde4 --- /dev/null +++ b/mednodes/preferences.py @@ -0,0 +1,23 @@ +import bpy +from pathlib import Path +from .external_package import ExternalPackagePreferences + + +class MedNodesPreferences(bpy.types.AddonPreferences, ExternalPackagePreferences): + bl_idname = __package__ + + cache_dir: bpy.props.StringProperty( + name="Python Log Directory", + subtype='DIR_PATH', + default=Path(Path.home(), '.mednodes').as_posix() + ) # type: ignore + + + def draw(self, context): + layout = self.layout + layout.label(text="Configuration") + layout.prop(self, 'cache_dir', text='Set Cache Directory') + + # ExternalPackagePreferences Config + self.requirements_dir = Path(__file__).parent.as_posix() + super().draw(context) diff --git a/mednodes/requirements.txt b/mednodes/requirements.txt new file mode 100644 index 0000000..2a2ba13 --- /dev/null +++ b/mednodes/requirements.txt @@ -0,0 +1 @@ +SimpleITK==2.3.1 # Insight Segmentation and Registration Toolkit. \ No newline at end of file