├── .gitignore ├── CMakeLists.txt ├── LICENSE ├── PyTorch.png ├── PyTorchUtils ├── CMakeLists.txt ├── PyTorchUtils.py ├── Resources │ ├── Icons │ │ └── PyTorchUtils.png │ └── UI │ │ └── PyTorchUtils.ui └── Testing │ ├── CMakeLists.txt │ └── Python │ └── CMakeLists.txt ├── PyTorch_128.png ├── README.md ├── logo_gimp.xcf └── project_week_diagram.png /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | .vscode 132 | .idea 133 | 134 | .DS_Store 135 | 136 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.13.4) 2 | 3 | project(PyTorch) 4 | 5 | #----------------------------------------------------------------------------- 6 | # Extension meta-information 7 | set(EXTENSION_HOMEPAGE "https://github.com/fepegar/SlicerPyTorch") 8 | set(EXTENSION_CATEGORY "Utilities") 9 | set(EXTENSION_CONTRIBUTORS "Fernando Pérez-García (University College London and King's College London), Andras Lasso (PerkLab Queen's University)") 10 | set(EXTENSION_DESCRIPTION "Utilities to install and use PyTorch within 3D Slicer.") 11 | set(EXTENSION_ICONURL "https://raw.githubusercontent.com/fepegar/SlicerPyTorch/master/PyTorch.png") 12 | set(EXTENSION_SCREENSHOTURLS "https://raw.githubusercontent.com/fepegar/SlicerPyTorch/master/project_week_diagram.png") 13 | set(EXTENSION_DEPENDS "NA") # Specified as a list or "NA" if no dependencies 14 | 15 | #----------------------------------------------------------------------------- 16 | # Extension dependencies 17 | find_package(Slicer REQUIRED) 18 | include(${Slicer_USE_FILE}) 19 | 20 | #----------------------------------------------------------------------------- 21 | # Extension modules 22 | add_subdirectory(PyTorchUtils) 23 | ## NEXT_MODULE 24 | 25 | #----------------------------------------------------------------------------- 26 | include(${Slicer_EXTENSION_GENERATE_CONFIG}) 27 | include(${Slicer_EXTENSION_CPACK}) 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | For more information, please see: 3 | 4 | http://www.slicer.org 5 | 6 | The 3D Slicer license below is a BSD style license, with extensions 7 | to cover contributions and other issues specific to 3D Slicer. 8 | 9 | 10 | 3D Slicer Contribution and Software License Agreement ("Agreement") 11 | Version 1.0 (December 20, 2005) 12 | 13 | This Agreement covers contributions to and downloads from the 3D 14 | Slicer project ("Slicer") maintained by The Brigham and Women's 15 | Hospital, Inc. ("Brigham"). Part A of this Agreement applies to 16 | contributions of software and/or data to Slicer (including making 17 | revisions of or additions to code and/or data already in Slicer). Part 18 | B of this Agreement applies to downloads of software and/or data from 19 | Slicer. Part C of this Agreement applies to all transactions with 20 | Slicer. If you distribute Software (as defined below) downloaded from 21 | Slicer, all of the paragraphs of Part B of this Agreement must be 22 | included with and apply to such Software. 23 | 24 | Your contribution of software and/or data to Slicer (including prior 25 | to the date of the first publication of this Agreement, each a 26 | "Contribution") and/or downloading, copying, modifying, displaying, 27 | distributing or use of any software and/or data from Slicer 28 | (collectively, the "Software") constitutes acceptance of all of the 29 | terms and conditions of this Agreement. If you do not agree to such 30 | terms and conditions, you have no right to contribute your 31 | Contribution, or to download, copy, modify, display, distribute or use 32 | the Software. 33 | 34 | PART A. CONTRIBUTION AGREEMENT - License to Brigham with Right to 35 | Sublicense ("Contribution Agreement"). 36 | 37 | 1. As used in this Contribution Agreement, "you" means the individual 38 | contributing the Contribution to Slicer and the institution or 39 | entity which employs or is otherwise affiliated with such 40 | individual in connection with such Contribution. 41 | 42 | 2. This Contribution Agreement applies to all Contributions made to 43 | Slicer, including without limitation Contributions made prior to 44 | the date of first publication of this Agreement. If at any time you 45 | make a Contribution to Slicer, you represent that (i) you are 46 | legally authorized and entitled to make such Contribution and to 47 | grant all licenses granted in this Contribution Agreement with 48 | respect to such Contribution; (ii) if your Contribution includes 49 | any patient data, all such data is de-identified in accordance with 50 | U.S. confidentiality and security laws and requirements, including 51 | but not limited to the Health Insurance Portability and 52 | Accountability Act (HIPAA) and its regulations, and your disclosure 53 | of such data for the purposes contemplated by this Agreement is 54 | properly authorized and in compliance with all applicable laws and 55 | regulations; and (iii) you have preserved in the Contribution all 56 | applicable attributions, copyright notices and licenses for any 57 | third party software or data included in the Contribution. 58 | 59 | 3. Except for the licenses granted in this Agreement, you reserve all 60 | right, title and interest in your Contribution. 61 | 62 | 4. You hereby grant to Brigham, with the right to sublicense, a 63 | perpetual, worldwide, non-exclusive, no charge, royalty-free, 64 | irrevocable license to use, reproduce, make derivative works of, 65 | display and distribute the Contribution. If your Contribution is 66 | protected by patent, you hereby grant to Brigham, with the right to 67 | sublicense, a perpetual, worldwide, non-exclusive, no-charge, 68 | royalty-free, irrevocable license under your interest in patent 69 | rights covering the Contribution, to make, have made, use, sell and 70 | otherwise transfer your Contribution, alone or in combination with 71 | any other code. 72 | 73 | 5. You acknowledge and agree that Brigham may incorporate your 74 | Contribution into Slicer and may make Slicer available to members 75 | of the public on an open source basis under terms substantially in 76 | accordance with the Software License set forth in Part B of this 77 | Agreement. You further acknowledge and agree that Brigham shall 78 | have no liability arising in connection with claims resulting from 79 | your breach of any of the terms of this Agreement. 80 | 81 | 6. YOU WARRANT THAT TO THE BEST OF YOUR KNOWLEDGE YOUR CONTRIBUTION 82 | DOES NOT CONTAIN ANY CODE THAT REQURES OR PRESCRIBES AN "OPEN 83 | SOURCE LICENSE" FOR DERIVATIVE WORKS (by way of non-limiting 84 | example, the GNU General Public License or other so-called 85 | "reciprocal" license that requires any derived work to be licensed 86 | under the GNU General Public License or other "open source 87 | license"). 88 | 89 | PART B. DOWNLOADING AGREEMENT - License from Brigham with Right to 90 | Sublicense ("Software License"). 91 | 92 | 1. As used in this Software License, "you" means the individual 93 | downloading and/or using, reproducing, modifying, displaying and/or 94 | distributing the Software and the institution or entity which 95 | employs or is otherwise affiliated with such individual in 96 | connection therewith. The Brigham and Women's Hospital, 97 | Inc. ("Brigham") hereby grants you, with right to sublicense, with 98 | respect to Brigham's rights in the software, and data, if any, 99 | which is the subject of this Software License (collectively, the 100 | "Software"), a royalty-free, non-exclusive license to use, 101 | reproduce, make derivative works of, display and distribute the 102 | Software, provided that: 103 | 104 | (a) you accept and adhere to all of the terms and conditions of this 105 | Software License; 106 | 107 | (b) in connection with any copy of or sublicense of all or any portion 108 | of the Software, all of the terms and conditions in this Software 109 | License shall appear in and shall apply to such copy and such 110 | sublicense, including without limitation all source and executable 111 | forms and on any user documentation, prefaced with the following 112 | words: "All or portions of this licensed product (such portions are 113 | the "Software") have been obtained under license from The Brigham and 114 | Women's Hospital, Inc. and are subject to the following terms and 115 | conditions:" 116 | 117 | (c) you preserve and maintain all applicable attributions, copyright 118 | notices and licenses included in or applicable to the Software; 119 | 120 | (d) modified versions of the Software must be clearly identified and 121 | marked as such, and must not be misrepresented as being the original 122 | Software; and 123 | 124 | (e) you consider making, but are under no obligation to make, the 125 | source code of any of your modifications to the Software freely 126 | available to others on an open source basis. 127 | 128 | 2. The license granted in this Software License includes without 129 | limitation the right to (i) incorporate the Software into 130 | proprietary programs (subject to any restrictions applicable to 131 | such programs), (ii) add your own copyright statement to your 132 | modifications of the Software, and (iii) provide additional or 133 | different license terms and conditions in your sublicenses of 134 | modifications of the Software; provided that in each case your use, 135 | reproduction or distribution of such modifications otherwise 136 | complies with the conditions stated in this Software License. 137 | 138 | 3. This Software License does not grant any rights with respect to 139 | third party software, except those rights that Brigham has been 140 | authorized by a third party to grant to you, and accordingly you 141 | are solely responsible for (i) obtaining any permissions from third 142 | parties that you need to use, reproduce, make derivative works of, 143 | display and distribute the Software, and (ii) informing your 144 | sublicensees, including without limitation your end-users, of their 145 | obligations to secure any such required permissions. 146 | 147 | 4. The Software has been designed for research purposes only and has 148 | not been reviewed or approved by the Food and Drug Administration 149 | or by any other agency. YOU ACKNOWLEDGE AND AGREE THAT CLINICAL 150 | APPLICATIONS ARE NEITHER RECOMMENDED NOR ADVISED. Any 151 | commercialization of the Software is at the sole risk of the party 152 | or parties engaged in such commercialization. You further agree to 153 | use, reproduce, make derivative works of, display and distribute 154 | the Software in compliance with all applicable governmental laws, 155 | regulations and orders, including without limitation those relating 156 | to export and import control. 157 | 158 | 5. The Software is provided "AS IS" and neither Brigham nor any 159 | contributor to the software (each a "Contributor") shall have any 160 | obligation to provide maintenance, support, updates, enhancements 161 | or modifications thereto. BRIGHAM AND ALL CONTRIBUTORS SPECIFICALLY 162 | DISCLAIM ALL EXPRESS AND IMPLIED WARRANTIES OF ANY KIND INCLUDING, 163 | BUT NOT LIMITED TO, ANY WARRANTIES OF MERCHANTABILITY, FITNESS FOR 164 | A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL 165 | BRIGHAM OR ANY CONTRIBUTOR BE LIABLE TO ANY PARTY FOR DIRECT, 166 | INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES 167 | HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY ARISING IN ANY WAY 168 | RELATED TO THE SOFTWARE, EVEN IF BRIGHAM OR ANY CONTRIBUTOR HAS 169 | BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. TO THE MAXIMUM 170 | EXTENT NOT PROHIBITED BY LAW OR REGULATION, YOU FURTHER ASSUME ALL 171 | LIABILITY FOR YOUR USE, REPRODUCTION, MAKING OF DERIVATIVE WORKS, 172 | DISPLAY, LICENSE OR DISTRIBUTION OF THE SOFTWARE AND AGREE TO 173 | INDEMNIFY AND HOLD HARMLESS BRIGHAM AND ALL CONTRIBUTORS FROM AND 174 | AGAINST ANY AND ALL CLAIMS, SUITS, ACTIONS, DEMANDS AND JUDGMENTS 175 | ARISING THEREFROM. 176 | 177 | 6. None of the names, logos or trademarks of Brigham or any of 178 | Brigham's affiliates or any of the Contributors, or any funding 179 | agency, may be used to endorse or promote products produced in 180 | whole or in part by operation of the Software or derived from or 181 | based on the Software without specific prior written permission 182 | from the applicable party. 183 | 184 | 7. Any use, reproduction or distribution of the Software which is not 185 | in accordance with this Software License shall automatically revoke 186 | all rights granted to you under this Software License and render 187 | Paragraphs 1 and 2 of this Software License null and void. 188 | 189 | 8. This Software License does not grant any rights in or to any 190 | intellectual property owned by Brigham or any Contributor except 191 | those rights expressly granted hereunder. 192 | 193 | PART C. MISCELLANEOUS 194 | 195 | This Agreement shall be governed by and construed in accordance with 196 | the laws of The Commonwealth of Massachusetts without regard to 197 | principles of conflicts of law. This Agreement shall supercede and 198 | replace any license terms that you may have agreed to previously with 199 | respect to Slicer. 200 | -------------------------------------------------------------------------------- /PyTorch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fepegar/SlicerPyTorch/e1c85cc0a62e92fa6c38223cd88c611906adb60e/PyTorch.png -------------------------------------------------------------------------------- /PyTorchUtils/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #----------------------------------------------------------------------------- 2 | set(MODULE_NAME PyTorchUtils) 3 | 4 | #----------------------------------------------------------------------------- 5 | set(MODULE_PYTHON_SCRIPTS 6 | ${MODULE_NAME}.py 7 | ) 8 | 9 | set(MODULE_PYTHON_RESOURCES 10 | Resources/Icons/${MODULE_NAME}.png 11 | Resources/UI/${MODULE_NAME}.ui 12 | ) 13 | 14 | #----------------------------------------------------------------------------- 15 | slicerMacroBuildScriptedModule( 16 | NAME ${MODULE_NAME} 17 | SCRIPTS ${MODULE_PYTHON_SCRIPTS} 18 | RESOURCES ${MODULE_PYTHON_RESOURCES} 19 | WITH_GENERIC_TESTS 20 | ) 21 | 22 | #----------------------------------------------------------------------------- 23 | if(BUILD_TESTING) 24 | 25 | # Register the unittest subclass in the main script as a ctest. 26 | # Note that the test will also be available at runtime. 27 | slicer_add_python_unittest(SCRIPT ${MODULE_NAME}.py) 28 | 29 | # Additional build-time testing 30 | add_subdirectory(Testing) 31 | endif() 32 | -------------------------------------------------------------------------------- /PyTorchUtils/PyTorchUtils.py: -------------------------------------------------------------------------------- 1 | import qt 2 | import logging 3 | 4 | import slicer 5 | from slicer.ScriptedLoadableModule import ( 6 | ScriptedLoadableModule, 7 | ScriptedLoadableModuleWidget, 8 | ScriptedLoadableModuleLogic, 9 | ScriptedLoadableModuleTest, 10 | ) 11 | 12 | 13 | class PyTorchUtils(ScriptedLoadableModule): 14 | def __init__(self, parent): 15 | ScriptedLoadableModule.__init__(self, parent) 16 | self.parent.title = "PyTorch Utils" 17 | self.parent.categories = ['Utilities'] 18 | self.parent.dependencies = [] 19 | self.parent.contributors = [ 20 | "Fernando Perez-Garcia (University College London and King's College London)", 21 | "Andras Lasso (PerkLab Queen's University)", 22 | ] 23 | self.parent.helpText = 'This hidden module containing some tools to work with PyTorch inside Slicer.' 24 | self.parent.acknowledgementText = ( 25 | 'This work was funded by the Engineering and Physical Sciences' 26 | ' Research Council (EPSRC) and supported by the UCL Centre for Doctoral' 27 | ' Training in Intelligent, Integrated Imaging in Healthcare, the UCL' 28 | ' Wellcome / EPSRC Centre for Interventional and Surgical Sciences (WEISS),' 29 | ' and the School of Biomedical Engineering & Imaging Sciences (BMEIS)' 30 | " of King's College London." 31 | ) 32 | 33 | 34 | class PyTorchUtilsWidget(ScriptedLoadableModuleWidget): 35 | 36 | def setup(self): 37 | super().setup() 38 | 39 | self.logic = PyTorchUtilsLogic() 40 | 41 | # Load widget from .ui file (created by Qt Designer). 42 | # Additional widgets can be instantiated manually and added to self.layout. 43 | uiWidget = slicer.util.loadUI(self.resourcePath('UI/PyTorchUtils.ui')) 44 | self.layout.addWidget(uiWidget) 45 | self.ui = slicer.util.childWidgetVariables(uiWidget) 46 | 47 | self.ui.detectPushButton.clicked.connect(self.onDetect) 48 | self.ui.installPushButton.clicked.connect(self.onInstallTorch) 49 | self.ui.uninstallPushButton.clicked.connect(self.onUninstallTorch) 50 | self.ui.restartPushButton.clicked.connect(self.onApplicationRestart) 51 | 52 | self.updateVersionInformation() 53 | 54 | def onDetect(self): 55 | with slicer.util.tryWithErrorDisplay("Failed to detect compatible computation backends.", waitCursor=True): 56 | torchVersionRequirement = self.ui.torchVersionLineEdit.text 57 | backends = PyTorchUtilsLogic.getCompatibleComputationBackends(torchVersionRequirement=torchVersionRequirement) 58 | currentBackend = self.ui.backendComboBox.currentText 59 | self.ui.backendComboBox.clear() 60 | self.ui.backendComboBox.addItem("automatic") 61 | for backend in backends: 62 | self.ui.backendComboBox.addItem(backend) 63 | self.ui.backendComboBox.currentText = currentBackend 64 | self.ui.backendComboBox.showPopup() 65 | self.updateVersionInformation() 66 | 67 | def onInstallTorch(self): 68 | with slicer.util.tryWithErrorDisplay("Failed to install PyTorch. Some PyTorch files may be in use or corrupted. Please restart the application, uninstall PyTorch, and try installing again.", waitCursor=True): 69 | if PyTorchUtilsLogic.torchInstalled(): 70 | torch = self.logic.torch 71 | slicer.util.delayDisplay(f'PyTorch {torch.__version__} is already installed, using {self.logic.getDevice()}.', autoCloseMsec=2000) 72 | else: 73 | backend = self.ui.backendComboBox.currentText 74 | automaticBackend = (backend == "automatic") 75 | askConfirmation = automaticBackend 76 | torchVersionRequirement = self.ui.torchVersionLineEdit.text 77 | torchvisionVersionRequirement = self.ui.torchvisionVersionLineEdit.text 78 | torch = self.logic.installTorch(askConfirmation, None if automaticBackend else backend, torchVersionRequirement, torchvisionVersionRequirement) 79 | if torch is not None: 80 | slicer.util.delayDisplay(f'PyTorch {torch.__version__} installed successfully using {self.logic.getDevice()}.', autoCloseMsec=2000) 81 | self.updateVersionInformation() 82 | 83 | def onUninstallTorch(self): 84 | with slicer.util.tryWithErrorDisplay("Failed to uninstall PyTorch. Probably PyTorch is already in use. Please restart the application and try again.", waitCursor=True): 85 | self.logic.uninstallTorch() 86 | slicer.util.delayDisplay(f'PyTorch uninstalled successfully.', autoCloseMsec=2000) 87 | self.updateVersionInformation() 88 | 89 | def updateVersionInformation(self): 90 | try: 91 | self.ui.torchVersionInformation.text = self.logic.torchVersionInformation 92 | except Exception as e: 93 | logging.error(str(e)) 94 | self.ui.torchVersionInformation.text = "unknown (corrupted installation?)" 95 | try: 96 | self.ui.torchvisionVersionInformation.text = self.logic.torchvisionVersionInformation 97 | except Exception as e: 98 | logging.error(str(e)) 99 | self.ui.torchvisionVersionInformation.text = "unknown (corrupted installation?)" 100 | try: 101 | info = self.logic.nvidiaDriverVersionInformation 102 | self.ui.nvidiaVersionInformation.text = info if info else "not found" 103 | except Exception as e: 104 | logging.error(str(e)) 105 | self.ui.nvidiaVersionInformation.text = "not found" 106 | 107 | def onApplicationRestart(self): 108 | slicer.util.restart() 109 | 110 | class PyTorchUtilsLogic(ScriptedLoadableModuleLogic): 111 | def __init__(self): 112 | self._torch = None 113 | 114 | @property 115 | def nvidiaDriverVersionInformation(self): 116 | """Get NVIDIA driver version information as a string that can be displayed to the user. 117 | If light-the-torch is not installed yet then empty string is returned. 118 | """ 119 | 120 | try: 121 | import light_the_torch._cb as computationBackend 122 | version = computationBackend._detect_nvidia_driver_version() 123 | if version is None: 124 | return "" 125 | else: 126 | return f"installed version {str(version)}" 127 | except Exception as e: 128 | # Don't install light-the-torch just for getting the NVIDIA driver version 129 | return "" 130 | 131 | @property 132 | def torchVersionInformation(self): 133 | """Get PyTorch version information as a string that can be displayed to the user. 134 | """ 135 | if not PyTorchUtilsLogic.torchInstalled(): 136 | return "not installed" 137 | import torch 138 | return f"installed version {torch.__version__}" 139 | 140 | @property 141 | def torchvisionVersionInformation(self): 142 | """Get TorchVision version information as a string that can be displayed to the user. 143 | """ 144 | if not PyTorchUtilsLogic.torchvisionInstalled(): 145 | return "not installed" 146 | import torchvision 147 | return f"installed version {torchvision.__version__}" 148 | 149 | @property 150 | def torch(self): 151 | """``torch`` Python module. it will be installed if necessary.""" 152 | if self._torch is None: 153 | logging.info('Importing torch...') 154 | self._torch = self.importTorch() 155 | return self._torch 156 | 157 | @staticmethod 158 | def torchInstalled(): 159 | # Attempt to import torch could load some files, which could prevent uninstalling a corrupted pytorch install 160 | import importlib.metadata 161 | try: 162 | metadataPath = [p for p in importlib.metadata.files('torch') if 'METADATA' in str(p)][0] 163 | except importlib.metadata.PackageNotFoundError as e: 164 | return False 165 | try: 166 | import torch 167 | installed = True 168 | except ModuleNotFoundError: 169 | installed = False 170 | return installed 171 | 172 | @staticmethod 173 | def torchvisionInstalled(): 174 | try: 175 | import torchvision 176 | installed = True 177 | except ModuleNotFoundError: 178 | installed = False 179 | return installed 180 | 181 | def importTorch(self): 182 | """Import the ``torch`` Python module, installing it if necessary.""" 183 | if PyTorchUtilsLogic.torchInstalled(): 184 | import torch 185 | else: 186 | torch = self.installTorch() 187 | if torch is None: 188 | logging.warning('PyTorch was not installed') 189 | else: 190 | logging.info(f'PyTorch {torch.__version__} imported successfully') 191 | logging.info(f'CUDA available: {torch.cuda.is_available()}') 192 | return torch 193 | 194 | def installTorch(self, askConfirmation=False, forceComputationBackend=None, torchVersionRequirement=None, torchvisionVersionRequirement=None): 195 | """Install PyTorch and return the ``torch`` Python module. 196 | 197 | :param forceComputationBackend: optional parameter to set computation backend (cpu, cu116, cu117, ...) 198 | :param torchVersionRequirement: optional version requirement for torch (e.g., ">=1.12") 199 | :param torchvisionVersionRequirement: optional version requirement for torchvision (e.g., ">=0.8") 200 | 201 | If computation backend is not specified then the ``light-the-torch`` Python package is used to get the most recent version of 202 | PyTorch compatible with the installed NVIDIA drivers. If CUDA-compatible device is not found, a version compiled for CPU will be installed. 203 | """ 204 | 205 | args = PyTorchUtilsLogic._getPipInstallArguments(forceComputationBackend, torchVersionRequirement, torchvisionVersionRequirement) 206 | 207 | if askConfirmation and not slicer.app.commandOptions().testingEnabled: 208 | install = slicer.util.confirmOkCancelDisplay( 209 | f'PyTorch will be downloaded and installed using light-the-torch (ltt {" ".join(args)}).' 210 | ' The process might take some minutes.' 211 | ) 212 | if not install: 213 | logging.info('Installation of PyTorch aborted by user') 214 | return None 215 | 216 | logging.info(f"Install PyTorch using light-the-torch with arguments: {args}") 217 | 218 | try: 219 | import light_the_torch._patch 220 | except: 221 | PyTorchUtilsLogic._installLightTheTorch() 222 | import light_the_torch._patch 223 | 224 | slicer.util._executePythonModule('light_the_torch', args) 225 | import torch 226 | logging.info(f'PyTorch {torch.__version__} installed successfully.') 227 | return torch 228 | 229 | def uninstallTorch(self, askConfirmation=False, forceComputationBackend=None): 230 | """Uninstall PyTorch""" 231 | slicer.util.pip_uninstall('torch torchvision') 232 | logging.info(f'PyTorch uninstalled successfully.') 233 | 234 | @staticmethod 235 | def _getPipInstallArguments(forceComputationBackend=None, torchVersionRequirement=None, torchvisionVersionRequirement=None): 236 | if torchVersionRequirement is None: 237 | torchVersionRequirement = "" 238 | if torchvisionVersionRequirement is None: 239 | torchvisionVersionRequirement = "" 240 | args = ["install", "torch"+torchVersionRequirement, "torchvision"+torchvisionVersionRequirement] 241 | if forceComputationBackend is not None: 242 | args.append(f"--pytorch-computation-backend={forceComputationBackend}") 243 | return args 244 | 245 | @staticmethod 246 | def _installLightTheTorch(): 247 | slicer.util.pip_install('light-the-torch>=0.5') 248 | 249 | @staticmethod 250 | def getCompatibleComputationBackends(forceComputationBackend=None, torchVersionRequirement=None): 251 | """Get the list of computation backends compatible with the available hardware. 252 | 253 | :param forceComputationBackend: optional parameter to set computation backend (cpu, cu116, cu117, ...) 254 | :param torchVersionRequirement: optional version requirement for torch (e.g., ">=1.12") 255 | 256 | If computation backend is not specified then the ``light-the-torch`` is used to get the most recent version of 257 | PyTorch compatible with the installed NVIDIA drivers. 258 | """ 259 | try: 260 | import light_the_torch._patch 261 | except: 262 | PyTorchUtilsLogic._installLightTheTorch() 263 | import light_the_torch._patch 264 | 265 | args = PyTorchUtilsLogic._getPipInstallArguments(forceComputationBackend, torchVersionRequirement) 266 | try: 267 | backends = sorted(light_the_torch._patch.LttOptions.from_pip_argv(args).computation_backends) 268 | except Exception as e: 269 | logging.warning(str(e)) 270 | raise ValueError(f"Failed to get computation backend. Requested computation backend: `{forceComputationBackend}`.") 271 | 272 | return backends 273 | 274 | def getPyTorchHubModel(self, repoOwner, repoName, modelName, addPretrainedKwarg=True, *args, **kwargs): 275 | """Use PyTorch Hub to download a PyTorch model, typically pre-trained. 276 | 277 | More information can be found at https://pytorch.org/hub/. 278 | """ 279 | repo = f'{repoOwner}/{repoName}' 280 | if addPretrainedKwarg: 281 | kwargs['pretrained'] = True 282 | model = self.torch.hub.load(repo, modelName, *args, **kwargs) 283 | return model 284 | 285 | def getDevice(self): 286 | """Get CUDA device if available and CPU otherwise.""" 287 | return self.torch.device('cuda') if self.torch.cuda.is_available() else 'cpu' 288 | 289 | @property 290 | def cuda(self): 291 | """Return True if a CUDA-compatible device is available.""" 292 | return self.getDevice() != 'cpu' 293 | 294 | 295 | class PyTorchUtilsTest(ScriptedLoadableModuleTest): 296 | 297 | def runTest(self): 298 | self.test_PyTorchUtils() 299 | 300 | def _delayDisplay(self, message): 301 | if not slicer.app.testingEnabled(): 302 | self.delayDisplay(message) 303 | 304 | def test_PyTorchUtils(self): 305 | self._delayDisplay('Starting the test') 306 | logic = PyTorchUtilsLogic() 307 | self._delayDisplay(f'CUDA available: {logic.torch.cuda.is_available()}') 308 | self._delayDisplay('Test passed!') 309 | -------------------------------------------------------------------------------- /PyTorchUtils/Resources/Icons/PyTorchUtils.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fepegar/SlicerPyTorch/e1c85cc0a62e92fa6c38223cd88c611906adb60e/PyTorchUtils/Resources/Icons/PyTorchUtils.png -------------------------------------------------------------------------------- /PyTorchUtils/Resources/UI/PyTorchUtils.ui: -------------------------------------------------------------------------------- 1 | 2 | 3 | PyTorchUtils 4 | 5 | 6 | 7 | 0 8 | 0 9 | 394 10 | 378 11 | 12 | 13 | 14 | 15 | 16 | 17 | QFrame::StyledPanel 18 | 19 | 20 | QFrame::Raised 21 | 22 | 23 | 24 | 25 | 26 | Torch: 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | TorchVision: 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | NVIDIA driver: 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | Install/uninstall 72 | 73 | 74 | 75 | 76 | 77 | Torch version requirement: 78 | 79 | 80 | 81 | 82 | 83 | 84 | Optional version requirement for installing Torch. For example, to install TorchVision version to earlier than 2.0, set this to "<2.0". The field has no effect if PyTorch is already installed. 85 | 86 | 87 | default 88 | 89 | 90 | 91 | 92 | 93 | 94 | TorchVision version requirement: 95 | 96 | 97 | 98 | 99 | 100 | 101 | Optional version requirement for installing TorchVision. For example, to install TorchVision version 1.12 or later, set this to ">=1.12". The field has no effect if PyTorch is already installed. 102 | 103 | 104 | default 105 | 106 | 107 | 108 | 109 | 110 | 111 | Computation backend: 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 0 122 | 0 123 | 124 | 125 | 126 | Computation backend (cpu, cu116, cu117...). If set to "automatic" then the backend most suitable to the available hardware will be used. 127 | 128 | 129 | true 130 | 131 | 132 | 133 | automatic 134 | 135 | 136 | 137 | 138 | cpu 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | Detect all compatible computation backends and add them to the drop-down list. 147 | 148 | 149 | show all compatible 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | Install Torch and TorchVision 159 | 160 | 161 | Install PyTorch 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | Uninstall Torch and TorchVision. Application restart may be required. 171 | 172 | 173 | Uninstall PyTorch 174 | 175 | 176 | 177 | 178 | 179 | 180 | Restart Slicer. If Torch is in use, the application must be restarted before it can be uninstalled. 181 | 182 | 183 | Restart the application 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | Qt::Vertical 196 | 197 | 198 | 199 | 20 200 | 40 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | ctkCollapsibleButton 210 | QWidget 211 |
ctkCollapsibleButton.h
212 | 1 213 | 214 | 215 | qMRMLCollapsibleButton 216 | ctkCollapsibleButton 217 |
qMRMLCollapsibleButton.h
218 | 1 219 | 220 | 221 | qMRMLWidget 222 | QWidget 223 |
qMRMLWidget.h
224 | 1 225 | 226 | 227 | 228 | 229 | 230 | -------------------------------------------------------------------------------- /PyTorchUtils/Testing/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(Python) 2 | -------------------------------------------------------------------------------- /PyTorchUtils/Testing/Python/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | #slicer_add_python_unittest(SCRIPT ${MODULE_NAME}ModuleTest.py) 3 | -------------------------------------------------------------------------------- /PyTorch_128.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fepegar/SlicerPyTorch/e1c85cc0a62e92fa6c38223cd88c611906adb60e/PyTorch_128.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SlicerPyTorch 2 | 3 |

4 | SlicerPyTorch logo 5 |

6 | 7 | This is the code for the `PyTorch` extension for [3D Slicer](https://www.slicer.org/). 8 | 9 | Its main function is to install [PyTorch](https://pytorch.org/) inside Slicer. 10 | The latest version compatible with the installed drivers will be selected automatically. 11 | [CI tests](https://slicer.cdash.org/index.php?project=SlicerPreview&filtercount=1&showfilters=1&field1=buildname&compare1=63&value1=PyTorch) are run nightly. 12 | 13 | PyTorch can be installed opening the `PyTorch Utils` module and clicking on the button, or programmatically: 14 | 15 | ```python 16 | import PyTorchUtils 17 | torch = PyTorchUtils.PyTorchUtilsLogic().torch # will be installed if necessary 18 | tensor = torch.rand(50, 60) 19 | ``` 20 | 21 | 22 | Unit tests are run nightly and the results can be checked on [CDash](https://slicer.cdash.org/index.php?project=SlicerPreview&filtercount=1&showfilters=1&field1=buildname&compare1=63&value1=PyTorch). 23 | 24 | Here's a diagram of the integration between PyTorch and Slicer: 25 | 26 | ![Diagram of PyTorch and Slicer](https://raw.githubusercontent.com/NA-MIC/ProjectWeek/master/PW35_2021_Virtual/Projects/PyTorchIntegration/diagram.svg) 27 | -------------------------------------------------------------------------------- /logo_gimp.xcf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fepegar/SlicerPyTorch/e1c85cc0a62e92fa6c38223cd88c611906adb60e/logo_gimp.xcf -------------------------------------------------------------------------------- /project_week_diagram.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fepegar/SlicerPyTorch/e1c85cc0a62e92fa6c38223cd88c611906adb60e/project_week_diagram.png --------------------------------------------------------------------------------