├── .azure-pipelines ├── azure-pipelines.yml ├── ci-unix.yml ├── ci-windows.yml ├── flake8-validation.py └── syntax-validation.py ├── .bumpversion.cfg ├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── MANIFEST.in ├── README.md ├── RELEASE.md ├── pyproject.toml ├── setup.cfg ├── setup.py └── src └── icebreaker ├── KMeans_segmenter.py ├── __init__.py ├── average_mask.py ├── cli ├── __init__.py ├── ib_5fig.py ├── ib_group.py ├── ib_job.py ├── micrographs_enhance.py ├── micrographs_group.py ├── particles_group.py └── stats.py ├── correct_path.py ├── filter_designer.py ├── filter_designer_not_normalized.py ├── five_figures.py ├── ib_group_template.star ├── ib_job_template_FLATTEN_MODE.star ├── ib_job_template_GROUP_MODE.star ├── ice_groups.py ├── icebreaker_equalize_multi.py ├── icebreaker_icegroups_multi.py ├── listdir.py ├── local_mask.py ├── original_mask.py ├── original_mask_fast.py ├── plot_boxplots.py ├── show_gradient.py ├── star_appender.py └── window_mean.py /.azure-pipelines/azure-pipelines.yml: -------------------------------------------------------------------------------- 1 | trigger: 2 | branches: 3 | include: 4 | - '*' 5 | tags: 6 | include: 7 | - '*' 8 | 9 | stages: 10 | - stage: static 11 | displayName: Static Analysis 12 | jobs: 13 | - job: checks 14 | displayName: static code analysis 15 | pool: 16 | vmImage: ubuntu-20.04 17 | steps: 18 | - task: UsePythonVersion@0 19 | displayName: Set up python 20 | inputs: 21 | versionSpec: 3.9 22 | 23 | - bash: | 24 | python .azure-pipelines/syntax-validation.py 25 | displayName: Syntax validation 26 | - bash: | 27 | pip install --disable-pip-version-check flake8 28 | python .azure-pipelines/flake8-validation.py 29 | displayName: Flake8 validation 30 | - stage: build 31 | displayName: Build 32 | dependsOn: 33 | jobs: 34 | - job: build 35 | displayName: build package 36 | pool: 37 | vmImage: ubuntu-20.04 38 | steps: 39 | - task: UsePythonVersion@0 40 | displayName: Set up python 41 | inputs: 42 | versionSpec: 3.9 43 | 44 | - bash: | 45 | pip install --disable-pip-version-check collective.checkdocs wheel 46 | displayName: Install dependencies 47 | - bash: | 48 | set -ex 49 | python setup.py sdist bdist_wheel 50 | mkdir -p dist/pypi 51 | shopt -s extglob 52 | mv -v dist/!(pypi) dist/pypi 53 | git archive HEAD | gzip > dist/repo-source.tar.gz 54 | ls -laR dist 55 | displayName: Build python package 56 | - task: PublishBuildArtifacts@1 57 | displayName: Store artifact 58 | inputs: 59 | pathToPublish: dist/ 60 | artifactName: package 61 | 62 | - bash: python setup.py checkdocs 63 | displayName: Check package description 64 | 65 | #- stage: tests 66 | # displayName: Run unit tests 67 | # dependsOn: 68 | # - static 69 | # - build 70 | # jobs: 71 | # - job: linux 72 | # pool: 73 | # vmImage: ubuntu-20.04 74 | # strategy: 75 | # matrix: 76 | # python36: 77 | # PYTHON_VERSION: 3.6 78 | # python37: 79 | # PYTHON_VERSION: 3.7 80 | # python38: 81 | # PYTHON_VERSION: 3.8 82 | # python39: 83 | # PYTHON_VERSION: 3.9 84 | # steps: 85 | # - template: .azure-pipelines/ci-unix.yml 86 | 87 | # - job: macOS 88 | # pool: 89 | # vmImage: macOS-latest 90 | # strategy: 91 | # matrix: 92 | # python36: 93 | # PYTHON_VERSION: 3.6 94 | # python37: 95 | # PYTHON_VERSION: 3.7 96 | # python38: 97 | # PYTHON_VERSION: 3.8 98 | # python39: 99 | # PYTHON_VERSION: 3.9 100 | # steps: 101 | # - template: .azure-pipelines/ci-unix.yml 102 | 103 | # - job: windows 104 | # pool: 105 | # vmImage: windows-latest 106 | # strategy: 107 | # matrix: 108 | # python36: 109 | # PYTHON_VERSION: 3.6 110 | # python37: 111 | # PYTHON_VERSION: 3.7 112 | # python38: 113 | # PYTHON_VERSION: 3.8 114 | # python39: 115 | # PYTHON_VERSION: 3.9 116 | # steps: 117 | # - template: .azure-pipelines/ci-windows.yml 118 | 119 | - stage: deploy 120 | displayName: Publish release 121 | dependsOn: 122 | - static 123 | - build 124 | condition: and(succeeded(), startsWith(variables['Build.SourceBranch'], 'refs/tags/')) 125 | jobs: 126 | - job: pypi 127 | displayName: Publish pypi release 128 | pool: 129 | vmImage: ubuntu-20.04 130 | steps: 131 | - checkout: none 132 | 133 | - task: UsePythonVersion@0 134 | displayName: Set up python 135 | inputs: 136 | versionSpec: 3.8 137 | - task: DownloadBuildArtifacts@0 138 | displayName: Get pre-built package 139 | inputs: 140 | buildType: 'current' 141 | downloadType: 'single' 142 | artifactName: 'package' 143 | downloadPath: '$(System.ArtifactsDirectory)' 144 | - script: | 145 | pip install --disable-pip-version-check twine 146 | displayName: Install twine 147 | - task: TwineAuthenticate@1 148 | displayName: Set up credentials 149 | inputs: 150 | pythonUploadServiceConnection: pypi-icebreaker 151 | 152 | - bash: | 153 | python -m twine upload -r pypi-icebreaker --config-file $(PYPIRC_PATH) $(System.ArtifactsDirectory)/package/pypi/*.tar.gz $(System.ArtifactsDirectory)/package/pypi/*.whl 154 | displayName: Publish package 155 | -------------------------------------------------------------------------------- /.azure-pipelines/ci-unix.yml: -------------------------------------------------------------------------------- 1 | steps: 2 | - checkout: none 3 | 4 | - task: UsePythonVersion@0 5 | inputs: 6 | versionSpec: '$(PYTHON_VERSION)' 7 | displayName: 'Use Python $(PYTHON_VERSION)' 8 | 9 | - task: DownloadBuildArtifacts@0 10 | displayName: Get pre-built package 11 | inputs: 12 | buildType: 'current' 13 | downloadType: 'single' 14 | artifactName: 'package' 15 | downloadPath: '$(System.ArtifactsDirectory)' 16 | 17 | - task: ExtractFiles@1 18 | displayName: Checkout sources 19 | inputs: 20 | archiveFilePatterns: "$(System.ArtifactsDirectory)/package/repo-source.tar.gz" 21 | destinationFolder: "$(Pipeline.Workspace)/src" 22 | 23 | - script: | 24 | set -eux 25 | pip install --disable-pip-version-check -r "$(Pipeline.Workspace)/src/requirements_dev.txt" pytest-azurepipelines 26 | pip install --no-deps --disable-pip-version-check -e "$(Pipeline.Workspace)/src" 27 | displayName: Install package 28 | 29 | - script: | 30 | PYTHONDEVMODE=1 pytest -v -ra --regression --cov=icebreaker --cov-report=xml --cov-branch --no-coverage-upload \ 31 | || echo "##vso[task.complete result=Failed;]Some tests failed" 32 | displayName: Run tests 33 | workingDirectory: $(Pipeline.Workspace)/src 34 | 35 | - bash: bash <(curl -s https://codecov.io/bash) -t $(CODECOV_TOKEN) -n "Python $(PYTHON_VERSION) $(Agent.OS)" 36 | displayName: Publish coverage stats 37 | continueOnError: True 38 | workingDirectory: $(Pipeline.Workspace)/src 39 | timeoutInMinutes: 2 40 | -------------------------------------------------------------------------------- /.azure-pipelines/ci-windows.yml: -------------------------------------------------------------------------------- 1 | steps: 2 | - checkout: none 3 | 4 | - task: UsePythonVersion@0 5 | inputs: 6 | versionSpec: '$(PYTHON_VERSION)' 7 | displayName: 'Use Python $(PYTHON_VERSION)' 8 | 9 | - task: DownloadBuildArtifacts@0 10 | displayName: Get pre-built package 11 | inputs: 12 | buildType: 'current' 13 | downloadType: 'single' 14 | artifactName: 'package' 15 | downloadPath: '$(System.ArtifactsDirectory)' 16 | 17 | - task: ExtractFiles@1 18 | displayName: Checkout sources 19 | inputs: 20 | archiveFilePatterns: "$(System.ArtifactsDirectory)/package/repo-source.tar.gz" 21 | destinationFolder: "$(Pipeline.Workspace)/src" 22 | 23 | - script: | 24 | pip install --disable-pip-version-check -r "$(Pipeline.Workspace)/src/requirements_dev.txt" pytest-azurepipelines 25 | pip install --no-deps --disable-pip-version-check -e "$(Pipeline.Workspace)/src" 26 | displayName: Install package 27 | 28 | - script: | 29 | set PYTHONDEVMODE=1 30 | pytest -v -ra --regression ^ 31 | --cov=icebreaker --cov-report=html --cov-branch --no-coverage-upload ^ 32 | || echo "##vso[task.complete result=Failed;]Some tests failed" 33 | displayName: Run tests 34 | workingDirectory: $(Pipeline.Workspace)/src 35 | 36 | - bash: bash <(curl -s https://codecov.io/bash) -t $(CODECOV_TOKEN) -n "Python $(PYTHON_VERSION) $(Agent.OS)" 37 | displayName: Publish coverage stats 38 | continueOnError: True 39 | workingDirectory: $(Pipeline.Workspace)/src 40 | timeoutInMinutes: 2 41 | -------------------------------------------------------------------------------- /.azure-pipelines/flake8-validation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | 4 | # Flake8 validation 5 | failures = 0 6 | try: 7 | flake8 = subprocess.run( 8 | [ 9 | "flake8", 10 | "--exit-zero", 11 | ], 12 | capture_output=True, 13 | check=True, 14 | encoding="latin-1", 15 | timeout=300, 16 | ) 17 | except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e: 18 | print( 19 | "##vso[task.logissue type=error;]flake8 validation failed with", 20 | str(e.__class__.__name__), 21 | ) 22 | print(e.stdout) 23 | print(e.stderr) 24 | print("##vso[task.complete result=Failed;]flake8 validation failed") 25 | exit() 26 | for line in flake8.stdout.split("\n"): 27 | if ":" not in line: 28 | continue 29 | filename, lineno, column, error = line.split(":", maxsplit=3) 30 | errcode, error = error.strip().split(" ", maxsplit=1) 31 | filename = os.path.normpath(filename) 32 | failures += 1 33 | print( 34 | f"##vso[task.logissue type=error;sourcepath={filename};" 35 | f"linenumber={lineno};columnnumber={column};code={errcode};]" + error 36 | ) 37 | 38 | if failures: 39 | print(f"##vso[task.logissue type=warning]Found {failures} flake8 violation(s)") 40 | print(f"##vso[task.complete result=Failed;]Found {failures} flake8 violation(s)") 41 | -------------------------------------------------------------------------------- /.azure-pipelines/syntax-validation.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import os 3 | import sys 4 | 5 | print("Python", sys.version, "\n") 6 | 7 | failures = 0 8 | 9 | for base, _, files in os.walk("."): 10 | for f in files: 11 | if not f.endswith(".py"): 12 | continue 13 | filename = os.path.normpath(os.path.join(base, f)) 14 | try: 15 | with open(filename, "r") as fh: 16 | ast.parse(fh.read()) 17 | except SyntaxError as se: 18 | failures += 1 19 | print( 20 | f"##vso[task.logissue type=error;sourcepath={filename};" 21 | f"linenumber={se.lineno};columnnumber={se.offset};]" 22 | f"SyntaxError: {se.msg}" 23 | ) 24 | print(" " + se.text + " " * se.offset + "^") 25 | print(f"SyntaxError: {se.msg} in {filename} line {se.lineno}") 26 | print() 27 | 28 | if failures: 29 | print(f"##vso[task.logissue type=warning]Found {failures} syntax error(s)") 30 | print(f"##vso[task.complete result=Failed;]Found {failures} syntax error(s)") 31 | -------------------------------------------------------------------------------- /.bumpversion.cfg: -------------------------------------------------------------------------------- 1 | [bumpversion] 2 | current_version = 0.3.9 3 | commit = True 4 | tag = True 5 | 6 | [bumpversion:file:setup.cfg] 7 | search = version = {current_version} 8 | replace = version = {new_version} 9 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # editor files 7 | *.swp 8 | ~* 9 | .project 10 | .pydevproject 11 | 12 | # C extensions 13 | *.so 14 | 15 | # Distribution / packaging 16 | .Python 17 | env/ 18 | build/ 19 | develop-eggs/ 20 | dist/ 21 | downloads/ 22 | eggs/ 23 | .eggs/ 24 | lib/ 25 | lib64/ 26 | parts/ 27 | sdist/ 28 | var/ 29 | wheels/ 30 | *.egg-info/ 31 | .installed.cfg 32 | *.egg 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .coverage 48 | .coverage.* 49 | .cache 50 | nosetests.xml 51 | coverage.xml 52 | *.cover 53 | .hypothesis/ 54 | .pytest_cache/ 55 | 56 | # Translations 57 | *.mo 58 | *.pot 59 | 60 | # Django stuff: 61 | *.log 62 | local_settings.py 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 | # pyenv 81 | .python-version 82 | 83 | # celery beat schedule file 84 | celerybeat-schedule 85 | 86 | # SageMath parsed files 87 | *.sage.py 88 | 89 | # dotenv 90 | .env 91 | 92 | # virtualenv 93 | .venv 94 | venv/ 95 | ENV/ 96 | 97 | # Spyder project settings 98 | .spyderproject 99 | .spyproject 100 | 101 | # Rope project settings 102 | .ropeproject 103 | 104 | # mkdocs documentation 105 | /site 106 | 107 | # mypy 108 | .mypy_cache/ 109 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | # Syntax validation and some basic sanity checks 3 | - repo: https://github.com/pre-commit/pre-commit-hooks 4 | rev: v4.1.0 5 | hooks: 6 | - id: check-merge-conflict 7 | - id: check-ast 8 | fail_fast: True 9 | - id: check-json 10 | - id: pretty-format-json 11 | args: ['--autofix'] 12 | - id: check-added-large-files 13 | args: ['--maxkb=200'] 14 | - id: check-yaml 15 | 16 | # Automatically sort imports 17 | - repo: https://github.com/PyCQA/isort 18 | rev: 5.12.0 19 | hooks: 20 | - id: isort 21 | args: [ 22 | #'-a', 'from __future__ import annotations', # 3.7-3.11 23 | '--rm', 'from __future__ import absolute_import', # -3.0 24 | '--rm', 'from __future__ import division', # -3.0 25 | '--rm', 'from __future__ import generator_stop', # -3.7 26 | '--rm', 'from __future__ import generators', # -2.3 27 | '--rm', 'from __future__ import nested_scopes', # -2.2 28 | '--rm', 'from __future__ import print_function', # -3.0 29 | '--rm', 'from __future__ import unicode_literals', # -3.0 30 | '--rm', 'from __future__ import with_statement', # -2.6 31 | ] 32 | 33 | # Automatic source code formatting 34 | - repo: https://github.com/psf/black 35 | rev: 22.3.0 36 | hooks: 37 | - id: black 38 | args: [--safe, --quiet] 39 | 40 | # Linting 41 | - repo: https://github.com/PyCQA/flake8 42 | rev: 4.0.1 43 | hooks: 44 | - id: flake8 45 | additional_dependencies: ['flake8-comprehensions==3.8.0'] 46 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2021, Diamond Light Source, University of York 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | 3. Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE 2 | recursive-include src/icebreaker *.py *.star 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ICEBREAKER 2 | 3 | The Icebreaker software for ice gradient estimation and removal on cryoEM micrographs. 4 | 5 | External job for calling IceBreaker grouping within Relion 3.1 6 | in_parts is input star file from relion picking job 7 | in_mics is previous grouping job directory 8 | 9 | Run in main Relion project directory 10 | 11 | ## Running as a RELION External job 12 | 13 | Icebreaker provides a number of command line tools that wrap the core functionality in a way that is friendly for the RELION External job. These tools are: 14 | 15 | ``` 16 | ib_job 17 | ib_group 18 | ib_5fig 19 | ``` 20 | 21 | If Icebreaker is installed in your Python environment then these can be used directly as the executables for a RELION External job. 22 | 23 | ## Processing the micrographs 24 | 25 | To process the micrographs use the `ib_job` command line tool 26 | 27 | As an input, it takes a .star file with a list of micrographs to process, e.g. micrographs_ctf.star 28 | 29 | In the 'Params' tab, you have to add a label **mode** which can take values **group** (it will group areas with similar ice thickness) or **flatten** (it will improve contrasts locally trying to remove the ice gradient). 30 | 31 | In the Running tab use the slider to select the number of threads. 32 | 33 | The output from this job is a new set of micrographs. 34 | 35 | 36 | ## Processing the particles 37 | 38 | This is done using the `ib_group` command line tool. As 'Input micrographs' it takes micrographs 'grouped' with the previous script. As 'Input particles' any star file with particles from jobs like Extract particles, 2D classes, etc. As a result, additional parameter column (\_ibIceGroup) with estimated ice thickness value for each particle will be added. 39 | 40 | 41 | ## Using the job templates 42 | 43 | To make it easier to get started running Icebreaker from Relion, you can use one of the template star files *ib\_job\_template\_GROUP\_MODE.star*, *ib\_job\_template\_FLATTEN\_MODE.star* or *ib\_group\_template.star*. To use a template: 44 | 45 | 1. Copy one of the templates to your Relion project directory as *.gui\_externaljob.py* 46 | 2. Select the External job type 47 | 3. Click Jobs -> Load job settings 48 | 4. Change the placeholder input names to select your input files 49 | 5. Choose your running options and run the job as normal 50 | 51 | ## Runnnig outside of RELION 52 | 53 | There are a number of command line tools which can be used outside of the RELION context. These are: 54 | 55 | ``` 56 | ib.micrographs_group 57 | ib.micrographs_flatten 58 | ib.particles_group 59 | ib.five_fig 60 | ``` 61 | 62 | The functionality of these tools correspond, respectively, to `ib_job` in **group** mode, `ib_job` in **flatten** mode, `ib_group` and `ib_5fig`. 63 | The required command line options can be seen by using the `--help` flag. 64 | -------------------------------------------------------------------------------- /RELEASE.md: -------------------------------------------------------------------------------- 1 | ## How to prepare a release 2 | 3 | To prepare an Icebreaker EM release you need to install the package [bump2version](https://pypi.org/project/bump2version/): 4 | 5 | ```bash 6 | pip install bump2version 7 | ``` 8 | 9 | and then, in the repository directory, run one of the following 10 | 11 | ```bash 12 | # assuming current version is 1.2.3 13 | bumpversion patch # release version 1.2.4 14 | bumpversion minor # release version 1.3.0 15 | bumpversion major # release version 2.0.0 16 | ``` 17 | 18 | This automatically creates a release commit and a release tag. 19 | You then need to push both to the Github repository: 20 | ```bash 21 | git push # pushes the release commit 22 | git push origin v2.0.0 # pushes the release tag for version 2.0.0 23 | ``` 24 | 25 | Assuming the tests pass the release is then created by Azure and uploaded directly onto [pypi](https://pypi.org/project/icebreaker_em/). -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools >= 40.6.0", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [tool.isort] 6 | profile="black" 7 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = icebreaker_em 3 | version = 0.3.9 4 | license = BSD 5 | license_files = LICENSE 6 | classifiers = 7 | License :: OSI Approved :: BSD License 8 | Programming Language :: Python :: 3 9 | Programming Language :: Python :: 3.6 10 | Programming Language :: Python :: 3.7 11 | Programming Language :: Python :: 3.8 12 | Programming Language :: Python :: 3.9 13 | project_urls = 14 | Download = https://github.com/DiamondLightSource/python-icebreaker/releases 15 | Documentation = https://github.com/DiamondLightSource/python-icebreaker 16 | GitHub = https://github.com/DiamondLightSource/python-icebreaker 17 | Bug-Tracker = https://github.com/DiamondLightSource/python-icebreaker/issues 18 | 19 | [options] 20 | include_package_data = True 21 | install_requires = 22 | gemmi 23 | matplotlib 24 | mrcfile 25 | numpy < 2 26 | opencv-python-headless 27 | scipy 28 | packages = find: 29 | package_dir = 30 | =src 31 | python_requires = >=3.6 32 | zip_safe = False 33 | 34 | [options.entry_points] 35 | console_scripts = 36 | ib_group = icebreaker.cli.ib_group:main 37 | ib_job = icebreaker.cli.ib_job:main 38 | ib_5fig = icebreaker.cli.ib_5fig:main 39 | ib.micrographs_group = icebreaker.cli.micrographs_group:main 40 | ib.micrographs_flatten = icebreaker.cli.micrographs_enhance:main 41 | ib.particles_group = icebreaker.cli.particles_group:main 42 | ib.five_fig = icebreaker.cli.stats:main 43 | libtbx.dispatcher.script = 44 | ib_group = ib_group 45 | ib_job = ib_job 46 | ib_5fig = ib_5fig 47 | 48 | [options.packages.find] 49 | where = src 50 | 51 | [flake8] 52 | # Black disagrees with flake8 on a few points. Ignore those. 53 | ignore = E203, E266, E501, W503 54 | # E203 whitespace before ':' 55 | # E266 too many leading '#' for block comment 56 | # E501 line too long 57 | # W503 line break before binary operator 58 | 59 | max-line-length = 88 60 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from setuptools import setup 4 | 5 | setup() 6 | -------------------------------------------------------------------------------- /src/icebreaker/KMeans_segmenter.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | 5 | def segmenter(img_to_segment, num_of_segments): 6 | rolled_resizedZ = img_to_segment.copy() 7 | Z = rolled_resizedZ.reshape((-1, 1)) 8 | Z = np.float32(Z) 9 | criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 50, 20.0) 10 | K = num_of_segments 11 | ret, label, center = cv2.kmeans(Z, K, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS) 12 | 13 | # Now convert back into uint8, and make original image 14 | center = np.uint8(center) 15 | res = center[label.flatten()] 16 | res2 = res.reshape((rolled_resizedZ.shape)) 17 | ret, thresh1 = cv2.threshold(rolled_resizedZ, 90, 255, cv2.THRESH_BINARY) 18 | 19 | return res2 20 | -------------------------------------------------------------------------------- /src/icebreaker/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DiamondLightSource/python-icebreaker/46a2c925f3f4a3fadaacc964cb974018a28a511e/src/icebreaker/__init__.py -------------------------------------------------------------------------------- /src/icebreaker/average_mask.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | 5 | def average_mask(lowpass1, img2, val): 6 | reg3 = np.ones((img2.shape[0], img2.shape[1]), np.uint8) 7 | coord0 = [] 8 | coord1 = [] 9 | # k = np.zeros((img2.shape[0], img2.shape[1]), np.uint8) 10 | for u in range(img2.shape[0]): 11 | for v in range(img2.shape[1]): 12 | if img2[u, v] == val: 13 | reg3[u, v] = 1 14 | else: 15 | reg3[u, v] = 0 16 | 17 | up_reg3 = cv2.resize( 18 | reg3, (lowpass1.shape[1], lowpass1.shape[0]), interpolation=cv2.INTER_AREA 19 | ) 20 | mean_vals = cv2.mean(lowpass1, up_reg3)[0] 21 | up_reg3 = up_reg3 * lowpass1 22 | # Negative of ROI 23 | # cover3 = 1 - reg3 24 | # Upscaled negative 25 | # up_cover3 = cv2.resize(cover3, (lowpass1.shape[1], lowpass1.shape[0])) 26 | mask3 = up_reg3 27 | coord0 = np.nonzero(mask3)[0] 28 | coord1 = np.nonzero(mask3)[1] 29 | for i in range(len(coord1)): 30 | mask3[coord0[i]][coord1[i]] = int(mean_vals) 31 | # up_cover3 = cv2.resize(cover3, (lowpass1.shape[1],lowpass1.shape[0])) 32 | return mask3 # up_cover3 33 | -------------------------------------------------------------------------------- /src/icebreaker/cli/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DiamondLightSource/python-icebreaker/46a2c925f3f4a3fadaacc964cb974018a28a511e/src/icebreaker/cli/__init__.py -------------------------------------------------------------------------------- /src/icebreaker/cli/ib_5fig.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """ 3 | External job for calling IceBreaker equalize within Relion 3.1 4 | in_star is input star file from relion motion cor 5 | Run in main Relion project directory 6 | external_IBequal.py 7 | """ 8 | import argparse 9 | import json 10 | import os 11 | import os.path 12 | import pathlib 13 | 14 | import gemmi 15 | 16 | from icebreaker import five_figures 17 | 18 | 19 | def run_job(project_dir, job_dir, args_list, cpus): 20 | parser = argparse.ArgumentParser() 21 | input_group = parser.add_mutually_exclusive_group(required=True) 22 | input_group.add_argument("--in_mics", help="Input: IB_grouped star file") 23 | input_group.add_argument("--single_mic", help="A single IB_grouped micrograph") 24 | args = parser.parse_args(args_list) 25 | 26 | if args.in_mics: 27 | starfile = args.in_mics 28 | 29 | # Reading the micrographs star file from relion 30 | ctf_star = os.path.join(project_dir, starfile) 31 | in_doc = gemmi.cif.read_file(ctf_star) 32 | 33 | data_as_dict = json.loads(in_doc.as_json())["micrographs"] 34 | 35 | # Not crucial so if fails due to any reason just carry on 36 | try: 37 | with open("done_mics.txt", "r") as f: 38 | done_mics = f.read().splitlines() 39 | except Exception: 40 | done_mics = [] 41 | else: 42 | data_as_dict = {"_rlnmicrographname": [args.single_mic]} 43 | done_mics = [] 44 | 45 | pathlib.Path("IB_input").mkdir(exist_ok=True) 46 | 47 | for micrograph in data_as_dict["_rlnmicrographname"]: 48 | if os.path.split(micrograph)[-1] not in done_mics: 49 | micpath = pathlib.Path(micrograph) 50 | link_path = pathlib.Path("IB_input") / pathlib.Path( 51 | *list(micpath.parent.parts)[2:] 52 | ) 53 | if not link_path.exists(): 54 | link_path.mkdir(parents=True) 55 | try: 56 | (pathlib.Path("IB_input") / pathlib.Path(micrograph).name).symlink_to( 57 | pathlib.Path(project_dir) / micrograph 58 | ) 59 | except FileExistsError: 60 | print( 61 | f"WARNING: IB_input/{os.path.split(micrograph)[-1]} " 62 | "already exists but is not in the done micrographs" 63 | ) 64 | 65 | if args.single_mic: 66 | five_fig_csv = five_figures.single_mic_5fig( 67 | pathlib.Path(project_dir) 68 | / job_dir 69 | / "IB_input" 70 | / pathlib.Path(args.single_mic).name 71 | ) 72 | summary_results = five_fig_csv.split(",") 73 | print("Results: " + " ".join(summary_results)) 74 | else: 75 | five_fig_csv = five_figures.main( 76 | pathlib.Path(project_dir) / job_dir / "IB_input", cpus, append=True 77 | ) 78 | print("Done five figures") 79 | 80 | if args.in_mics: 81 | with open( 82 | "done_mics.txt", "a+" 83 | ) as f: # Done mics is to ensure that IB doesn't pick from already done mics 84 | for micrograph in pathlib.Path("./IB_input").glob("**/*"): 85 | if micrograph.suffix == ".mrc": 86 | f.write(micrograph.name + "\n") 87 | 88 | # Required star file 89 | out_doc = gemmi.cif.Document() 90 | output_nodes_block = out_doc.add_new_block("output_nodes") 91 | loop = output_nodes_block.init_loop( 92 | "", ["_rlnPipeLineNodeName", "_rlnPipeLineNodeType"] 93 | ) 94 | # loop.add_row([os.path.join(job_dir, 'ib_equalize.star'), '1']) 95 | loop.add_row([os.path.join(job_dir, "five_figs_test.csv"), "1"]) 96 | out_doc.write_file("RELION_OUTPUT_NODES.star") 97 | 98 | 99 | def main(): 100 | """Change to the job working directory, then call run_job()""" 101 | parser = argparse.ArgumentParser() 102 | parser.add_argument("--o", dest="out_dir", help="Output directory name") 103 | parser.add_argument("--j", help="relion stuff...", type=int) 104 | parser.add_argument("--pipeline_control", help="relion scheduler stuff...") 105 | known_args, other_args = parser.parse_known_args() 106 | 107 | project_dir = os.getcwd() 108 | job_dir = known_args.out_dir 109 | try: 110 | os.mkdir(job_dir) 111 | except FileExistsError: 112 | pass 113 | os.chdir(job_dir) 114 | try: 115 | run_job(project_dir, job_dir, other_args, known_args.j) 116 | except Exception: 117 | open("RELION_JOB_EXIT_FAILURE", "w").close() 118 | raise 119 | else: 120 | open("RELION_JOB_EXIT_SUCCESS", "w").close() 121 | 122 | 123 | if __name__ == "__main__": 124 | main() 125 | -------------------------------------------------------------------------------- /src/icebreaker/cli/ib_group.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """ 3 | External job for calling IceBreaker grouping within Relion 3.1 4 | in_parts is input star file from relion picking job 5 | in_mics is previous grouping job directory 6 | 7 | Run in main Relion project directory 8 | """ 9 | 10 | import argparse 11 | import os 12 | import os.path 13 | 14 | import gemmi 15 | 16 | from icebreaker import ice_groups as ib_igroups 17 | 18 | 19 | def run_job(project_dir, job_dir, args_list): 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--in_parts", help="Input: particle star file") 22 | parser.add_argument("--in_mics", help="Input: previous ice_group job") 23 | args = parser.parse_args(args_list) 24 | parts_star = args.in_parts 25 | group_star = os.path.join(project_dir, args.in_mics) 26 | print(group_star) 27 | # with open(group_star) as f: 28 | # lines = [line.rstrip() for line in f] 29 | # group_job = lines[1] 30 | 31 | ib_igroups.main( 32 | os.path.join(project_dir, parts_star), os.path.join(project_dir, group_star) 33 | ) 34 | 35 | # Writing a star file for Relion 36 | icegroups_doc = gemmi.cif.Document() 37 | icegroups_block = icegroups_doc.add_new_block("input_files") 38 | loop = icegroups_block.init_loop("", ["_rlnParticles", "_rlnMicrographs"]) 39 | loop.add_row( 40 | [os.path.join(project_dir, parts_star), os.path.join(project_dir, group_star)] 41 | ) 42 | icegroups_doc.write_file("ib_icegroups.star") 43 | 44 | # Required star file 45 | out_doc = gemmi.cif.Document() 46 | output_nodes_block = out_doc.add_new_block("output_nodes") 47 | loop = output_nodes_block.init_loop( 48 | "", ["_rlnPipeLineNodeName", "_rlnPipeLineNodeType"] 49 | ) 50 | loop.add_row([os.path.join(job_dir, "particles.star"), "5"]) 51 | out_doc.write_file("RELION_OUTPUT_NODES.star") 52 | 53 | 54 | def main(): 55 | """Change to the job working directory, then call run_job()""" 56 | parser = argparse.ArgumentParser() 57 | parser.add_argument("--o", dest="out_dir", help="Output directory name") 58 | parser.add_argument("--j", help="relion stuff...") 59 | parser.add_argument("--pipeline_control", help="relion scheduler stuff...") 60 | known_args, other_args = parser.parse_known_args() 61 | 62 | project_dir = os.getcwd() 63 | job_dir = known_args.out_dir 64 | try: 65 | os.mkdir(job_dir) 66 | except FileExistsError: 67 | pass 68 | os.chdir(job_dir) 69 | try: 70 | run_job(project_dir, job_dir, other_args) 71 | except Exception: 72 | open("RELION_JOB_EXIT_FAILURE", "w").close() 73 | raise 74 | else: 75 | open("RELION_JOB_EXIT_SUCCESS", "w").close() 76 | 77 | 78 | if __name__ == "__main__": 79 | main() 80 | -------------------------------------------------------------------------------- /src/icebreaker/cli/ib_job.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """ 3 | External job for calling IceBreaker equalize within Relion 3.1 4 | in_star is input star file from relion motion cor 5 | 6 | Run in main Relion project directory 7 | external_IBequal.py 8 | """ 9 | 10 | 11 | import argparse 12 | import json 13 | import os 14 | import os.path 15 | import pathlib 16 | import shutil 17 | 18 | import gemmi 19 | 20 | from icebreaker import correct_path 21 | from icebreaker import icebreaker_equalize_multi as ib_equal 22 | from icebreaker import icebreaker_icegroups_multi as ib_group 23 | from icebreaker import star_appender 24 | 25 | 26 | def run_job(project_dir, job_dir, args_list, mode, cpus): 27 | parser = argparse.ArgumentParser() 28 | input_group = parser.add_mutually_exclusive_group(required=True) 29 | input_group.add_argument("--in_mics", help="Input: Motion correction star file") 30 | input_group.add_argument("--single_mic", help="Single motion corrected micrograph") 31 | args = parser.parse_args(args_list) 32 | 33 | if args.in_mics: 34 | starfile = args.in_mics 35 | 36 | # Reading the micrographs star file from relion 37 | ctf_star = os.path.join(project_dir, starfile) 38 | in_doc = gemmi.cif.read_file(ctf_star) 39 | 40 | data_as_dict = json.loads(in_doc.as_json())["micrographs"] 41 | 42 | # Not crucial so if fails due to any reason just carry on 43 | try: 44 | with open("done_mics.txt", "r") as f: 45 | done_mics = f.read().splitlines() 46 | except Exception: 47 | done_mics = [] 48 | else: 49 | data_as_dict = {"_rlnmicrographname": [args.single_mic]} 50 | done_mics = [] 51 | 52 | pathlib.Path("IB_input").mkdir(exist_ok=True) 53 | 54 | for micrograph in data_as_dict["_rlnmicrographname"]: 55 | if os.path.split(micrograph)[-1] not in done_mics: 56 | try: 57 | (pathlib.Path("IB_input") / pathlib.Path(micrograph).name).symlink_to( 58 | pathlib.Path(project_dir) / micrograph 59 | ) 60 | except FileExistsError: 61 | print( 62 | f"WARNING: IB_input/{os.path.split(micrograph)[-1]} " 63 | "already exists but is not in the done micrographs" 64 | ) 65 | 66 | if args.in_mics: 67 | if mode == "group": 68 | ib_group.main("IB_input", cpus) 69 | print("Done equalizing") 70 | elif mode == "flatten": 71 | ib_equal.main("IB_input", cpus) 72 | else: 73 | for micrograph in data_as_dict["_rlnmicrographname"]: 74 | micrograph_name = pathlib.Path(pathlib.Path(micrograph).name).stem 75 | pathlib.Path(f"IB_input_{micrograph_name}").mkdir() 76 | ( 77 | pathlib.Path(f"IB_input_{micrograph_name}") / f"{micrograph_name}.mrc" 78 | ).symlink_to(pathlib.Path(project_dir) / micrograph) 79 | if mode == "group": 80 | ib_group.main(f"IB_input_{micrograph_name}", cpus) 81 | print("Done equalizing") 82 | elif mode == "flatten": 83 | ib_equal.main(f"IB_input_{micrograph_name}", cpus) 84 | 85 | pathlib.Path(f"IB_input/{mode}ed").mkdir(exist_ok=True) 86 | pathlib.Path( 87 | f"IB_input_{micrograph_name}/{mode}ed/{micrograph_name}_{mode}ed.mrc" 88 | ).rename(f"IB_input/{mode}ed/{micrograph_name}_{mode}ed.mrc") 89 | shutil.rmtree(f"IB_input_{micrograph_name}") 90 | 91 | if args.in_mics: 92 | with open( 93 | "done_mics.txt", "a+" 94 | ) as f: # Done mics is to ensure that IB doesn't pick from already done mics 95 | for micrograph in os.listdir("IB_input"): 96 | if micrograph.endswith("mrc"): 97 | f.write(micrograph + "\n") 98 | 99 | correct_path.correct( 100 | data_as_dict, os.path.join("IB_input", f"{mode}ed"), f"{mode}ed" 101 | ) 102 | 103 | # Writing a star file for Relion 104 | # part_doc = open('ib_equalize.star', 'w') 105 | # part_doc.write(os.path.join(project_dir, starfile)+'\n') 106 | # part_doc.write(job_dir) 107 | # part_doc.close() 108 | 109 | if args.in_mics: 110 | star_appender.mic_star(ctf_star, job_dir, mode) 111 | 112 | # Required star file 113 | out_doc = gemmi.cif.Document() 114 | output_nodes_block = out_doc.add_new_block("output_nodes") 115 | loop = output_nodes_block.init_loop( 116 | "", ["_rlnPipeLineNodeName", "_rlnPipeLineNodeType"] 117 | ) 118 | # loop.add_row([os.path.join(job_dir, 'ib_equalize.star'), '1']) 119 | loop.add_row([os.path.join(job_dir, f"{mode}ed_micrographs.star"), "1"]) 120 | out_doc.write_file("RELION_OUTPUT_NODES.star") 121 | 122 | 123 | def main(): 124 | """Change to the job working directory, then call run_job()""" 125 | parser = argparse.ArgumentParser() 126 | parser.add_argument("--o", dest="out_dir", help="Output directory name") 127 | parser.add_argument("--mode", help="ICEBREAKER mode to run i.e. flatten or group") 128 | parser.add_argument("--j", help="relion stuff...", type=int) 129 | parser.add_argument("--pipeline_control", help="relion scheduler stuff...") 130 | known_args, other_args = parser.parse_known_args() 131 | 132 | cpus = known_args.j 133 | mode = known_args.mode 134 | poss_modes = ["group", "flatten"] 135 | if mode not in poss_modes: 136 | print("IB ERROR - Mode type not found") 137 | exit() 138 | 139 | project_dir = os.getcwd() 140 | job_dir = known_args.out_dir 141 | try: 142 | os.mkdir(job_dir) 143 | except FileExistsError: 144 | pass 145 | os.chdir(job_dir) 146 | try: 147 | run_job(project_dir, job_dir, other_args, mode, cpus) 148 | except Exception: 149 | open("RELION_JOB_EXIT_FAILURE", "w").close() 150 | raise 151 | else: 152 | open("RELION_JOB_EXIT_SUCCESS", "w").close() 153 | 154 | 155 | if __name__ == "__main__": 156 | main() 157 | -------------------------------------------------------------------------------- /src/icebreaker/cli/micrographs_enhance.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from icebreaker import icebreaker_equalize_multi 4 | 5 | 6 | def main(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument( 9 | "-i", 10 | "--input", 11 | help="Directory containing micrographs (MRC files). Output will appear in a new directory called 'flattened' within this directory", 12 | dest="indir", 13 | ) 14 | parser.add_argument( 15 | "-j", help="Number of processes", dest="nproc", type=int, default=1 16 | ) 17 | args = parser.parse_args() 18 | 19 | icebreaker_equalize_multi.main(args.indir, args.nproc) 20 | -------------------------------------------------------------------------------- /src/icebreaker/cli/micrographs_group.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from icebreaker import icebreaker_icegroups_multi 4 | 5 | 6 | def main(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument( 9 | "-i", 10 | "--input", 11 | help="Directory containing micrographs (MRC files). Output will appear in a new directory called 'grouped' within this directory", 12 | dest="indir", 13 | ) 14 | parser.add_argument( 15 | "-j", help="Number of processes", dest="nproc", type=int, default=1 16 | ) 17 | args = parser.parse_args() 18 | 19 | icebreaker_icegroups_multi.main(args.indir, args.nproc) 20 | -------------------------------------------------------------------------------- /src/icebreaker/cli/particles_group.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from icebreaker import ice_groups 4 | 5 | 6 | def main(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument( 9 | "-s", 10 | "--star", 11 | help="Star file containing particles", 12 | dest="instar", 13 | ) 14 | parser.add_argument( 15 | "-i", 16 | "--input", 17 | help="Directory containing micrographs pointed to in the input star file", 18 | dest="indir", 19 | ) 20 | args = parser.parse_args() 21 | 22 | ice_groups.main(args.instar, args.indir) 23 | -------------------------------------------------------------------------------- /src/icebreaker/cli/stats.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from icebreaker import five_figures 4 | 5 | 6 | def main(): 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument( 9 | "-i", 10 | "--input", 11 | help="Directory containing grouped micrographs", 12 | dest="indir", 13 | ) 14 | parser.add_argument( 15 | "-j", help="Number of processes", dest="nproc", type=int, default=1 16 | ) 17 | args = parser.parse_args() 18 | 19 | five_figures.main(args.indir, cpus=args.nproc) 20 | -------------------------------------------------------------------------------- /src/icebreaker/correct_path.py: -------------------------------------------------------------------------------- 1 | """ 2 | Creates a folder tree that is the relion standard and then populates 3 | folder tree with the flattened micrographs produced by ICEBREAKER 4 | """ 5 | import os 6 | import pathlib 7 | import shutil 8 | 9 | 10 | def correct(data_as_dict, all_dir, ending): 11 | for i in range(len(data_as_dict["_rlnmicrographname"])): 12 | name = data_as_dict["_rlnmicrographname"][i] 13 | # New name needs to be extracted from inputs of form 14 | # JobName/JobNumber/path/to/file.suffix 15 | name_parts = list(pathlib.Path(name).parts) 16 | mic_file = name_parts[-1] 17 | full_dir = "/".join(name_parts[2:-1]) 18 | pathlib.Path(full_dir).mkdir(parents=True, exist_ok=True) 19 | picked_star = os.path.splitext(mic_file)[0] + f"_{ending}.mrc" 20 | if os.path.isfile(os.path.join(all_dir, picked_star)): 21 | try: 22 | shutil.move( 23 | os.path.join(all_dir, picked_star), 24 | os.path.join(full_dir, picked_star), 25 | ) 26 | except FileNotFoundError: 27 | print("Warning - Flattened Mic not found (Most likely already moved)") 28 | pass 29 | -------------------------------------------------------------------------------- /src/icebreaker/filter_designer.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import mrcfile 3 | import numpy as np 4 | 5 | 6 | def load_img(img_path): 7 | with mrcfile.open(img_path, "r+") as mrc: 8 | img = mrc.data 9 | # plt.imshow(img, cmap='gray') 10 | # plt.show() 11 | return img 12 | 13 | 14 | def lowpass(img, pixel_size, lowcut, edge_type, falloff, lowpass=True): 15 | mask_rad = int((pixel_size / lowcut) * (img.shape[1])) 16 | rows, cols = img.shape 17 | crow, ccol = int(rows / 2), int(cols / 2) # center 18 | 19 | filter = np.zeros((rows, cols, 2), np.float32) 20 | f_mask = np.zeros((rows, cols), np.float32) 21 | 22 | mask = cv2.circle(f_mask, (ccol, crow), mask_rad, (1), -1) 23 | 24 | if edge_type == "cos": 25 | for i in range(falloff): 26 | mask_edge = cv2.circle( 27 | mask, 28 | (ccol, crow), 29 | mask_rad + i, 30 | (np.cos(i * ((np.pi * 0.5) / falloff))), 31 | 2, 32 | ) 33 | elif edge_type == "linear": 34 | for i in range(falloff): 35 | mask_edge = cv2.circle( 36 | mask, (ccol, crow), mask_rad + i, (1 - i * (1 / falloff)), 2 37 | ) 38 | else: 39 | mask_edge = mask 40 | filter[:, :, 0] = mask_edge 41 | filter[:, :, 1] = mask_edge 42 | if lowpass is False: 43 | filter = 1 - filter 44 | # print(img.shape) 45 | # print(filter.shape) 46 | return filter 47 | 48 | 49 | def bandpass( 50 | img, 51 | pixel_size, 52 | lowcut, 53 | highcut, 54 | edge_type, 55 | falloff_in, 56 | falloff_out=1, 57 | bandpass_type=True, 58 | ): 59 | mask_rad_in = int((pixel_size / lowcut) * (img.shape[1])) - falloff_in 60 | mask_rad_out = int((pixel_size / highcut) * (img.shape[1])) 61 | rows, cols = img.shape 62 | crow, ccol = int(rows / 2), int(cols / 2) # center 63 | filter = np.zeros((rows, cols, 2), np.float32) 64 | 65 | # mask_bandpass = np.zeros((rows, cols, 2), np.float32) 66 | 67 | maskin = np.zeros((rows, cols), np.float32) 68 | maskout = np.zeros((rows, cols), np.float32) 69 | 70 | bandpass_out = cv2.circle(maskout, (ccol, crow), mask_rad_out, (1), -1) 71 | bandpass_in = cv2.circle(maskin, (ccol, crow), mask_rad_in, (1), -1) 72 | 73 | if edge_type == "cos": 74 | for i in range(falloff_out): 75 | bandpass_out = cv2.circle( 76 | bandpass_out, 77 | (ccol, crow), 78 | mask_rad_out + i, 79 | (np.cos(i * ((np.pi * 0.5) / falloff_out))), 80 | 2, 81 | ) 82 | for i in range(falloff_in): 83 | bandpass_in = cv2.circle( 84 | bandpass_in, 85 | (ccol, crow), 86 | mask_rad_in + i, 87 | (np.cos(i * ((np.pi * 0.5) / falloff_in))), 88 | 2, 89 | ) 90 | 91 | elif edge_type == "linear": 92 | for i in range(falloff_out): 93 | bandpass_out = cv2.circle( 94 | bandpass_out, 95 | (ccol, crow), 96 | mask_rad_out + i, 97 | (1 - i * (1 / falloff_out)), 98 | 2, 99 | ) 100 | for i in range(falloff_in): 101 | bandpass_in = cv2.circle( 102 | bandpass_in, 103 | (ccol, crow), 104 | mask_rad_in + i, 105 | (1 - i * (1 / falloff_in)), 106 | 2, 107 | ) 108 | 109 | else: 110 | bandpass = bandpass_out * (1 - bandpass_in) 111 | 112 | bandpass = bandpass_out * (1 - bandpass_in) 113 | 114 | filter[:, :, 0] = bandpass 115 | filter[:, :, 1] = bandpass 116 | if bandpass_type is False: 117 | filter = 1 - filter 118 | return filter 119 | 120 | 121 | def filtering(img, filter_type): 122 | img_float32 = np.float32(img) 123 | dft = cv2.dft(img_float32, flags=cv2.DFT_COMPLEX_OUTPUT) 124 | dft_shift = np.fft.fftshift(dft) 125 | fshift = dft_shift * (filter_type) 126 | f_ishift = np.fft.ifftshift(fshift) 127 | img_back = cv2.idft(f_ishift) 128 | img_back = cv2.magnitude(img_back[:, :, 0], img_back[:, :, 1]) 129 | cv2.normalize(img_back, img_back, 0.0, 255.0, cv2.NORM_MINMAX) 130 | img_back = img_back.astype(np.uint8) 131 | cv2.equalizeHist(img_back, img_back) 132 | magnitude_spectrum = 20 * np.log( 133 | cv2.magnitude(dft_shift[:, :, 0], dft_shift[:, :, 1]) + 1 134 | ) 135 | return img_back, magnitude_spectrum 136 | -------------------------------------------------------------------------------- /src/icebreaker/filter_designer_not_normalized.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import matplotlib.pyplot as plt 3 | import mrcfile 4 | import numpy as np 5 | 6 | 7 | def load_img(img_path): 8 | with mrcfile.open(img_path, "r+") as mrc: 9 | img = mrc.data 10 | plt.imshow(img, cmap="gray") 11 | plt.show() 12 | return img 13 | 14 | 15 | def lowpass(img, pixel_size, lowcut, edge_type, falloff, lowpass=True): 16 | mask_rad = int((pixel_size / lowcut) * (img.shape[1])) 17 | rows, cols = img.shape 18 | crow, ccol = int(rows / 2), int(cols / 2) # center 19 | 20 | filter = np.zeros((rows, cols, 2), np.float32) 21 | f_mask = np.zeros((rows, cols), np.float32) 22 | 23 | mask = cv2.circle(f_mask, (ccol, crow), mask_rad, (1), -1) 24 | 25 | if edge_type == "cos": 26 | for i in range(falloff): 27 | mask_edge = cv2.circle( 28 | mask, 29 | (ccol, crow), 30 | mask_rad + i, 31 | (np.cos(i * ((np.pi * 0.5) / falloff))), 32 | 2, 33 | ) 34 | elif edge_type == "linear": 35 | for i in range(falloff): 36 | mask_edge = cv2.circle( 37 | mask, (ccol, crow), mask_rad + i, (1 - i * (1 / falloff)), 2 38 | ) 39 | else: 40 | mask_edge = mask 41 | filter[:, :, 0] = mask_edge 42 | filter[:, :, 1] = mask_edge 43 | if lowpass is False: 44 | filter = 1 - filter 45 | # print(img.shape) 46 | # print(filter.shape) 47 | return filter 48 | 49 | 50 | def bandpass( 51 | img, 52 | pixel_size, 53 | lowcut, 54 | highcut, 55 | edge_type, 56 | falloff_in, 57 | falloff_out=1, 58 | bandpass_type=True, 59 | ): 60 | mask_rad_in = int((pixel_size / lowcut) * (img.shape[1])) - falloff_in 61 | mask_rad_out = int((pixel_size / highcut) * (img.shape[1])) 62 | rows, cols = img.shape 63 | crow, ccol = int(rows / 2), int(cols / 2) # center 64 | filter = np.zeros((rows, cols, 2), np.float32) 65 | 66 | # mask_bandpass = np.zeros((rows, cols, 2), np.float32) 67 | 68 | maskin = np.zeros((rows, cols), np.float32) 69 | maskout = np.zeros((rows, cols), np.float32) 70 | 71 | bandpass_out = cv2.circle(maskout, (ccol, crow), mask_rad_out, (1), -1) 72 | bandpass_in = cv2.circle(maskin, (ccol, crow), mask_rad_in, (1), -1) 73 | 74 | if edge_type == "cos": 75 | for i in range(falloff_out): 76 | bandpass_out = cv2.circle( 77 | bandpass_out, 78 | (ccol, crow), 79 | mask_rad_out + i, 80 | (np.cos(i * ((np.pi * 0.5) / falloff_out))), 81 | 2, 82 | ) 83 | for i in range(falloff_in): 84 | bandpass_in = cv2.circle( 85 | bandpass_in, 86 | (ccol, crow), 87 | mask_rad_in + i, 88 | (np.cos(i * ((np.pi * 0.5) / falloff_in))), 89 | 2, 90 | ) 91 | 92 | elif edge_type == "linear": 93 | for i in range(falloff_out): 94 | bandpass_out = cv2.circle( 95 | bandpass_out, 96 | (ccol, crow), 97 | mask_rad_out + i, 98 | (1 - i * (1 / falloff_out)), 99 | 2, 100 | ) 101 | for i in range(falloff_in): 102 | bandpass_in = cv2.circle( 103 | bandpass_in, 104 | (ccol, crow), 105 | mask_rad_in + i, 106 | (1 - i * (1 / falloff_in)), 107 | 2, 108 | ) 109 | 110 | else: 111 | bandpass = bandpass_out * (1 - bandpass_in) 112 | 113 | bandpass = bandpass_out * (1 - bandpass_in) 114 | 115 | filter[:, :, 0] = bandpass 116 | filter[:, :, 1] = bandpass 117 | if bandpass_type is False: 118 | filter = 1 - filter 119 | return filter 120 | 121 | 122 | def filtering(img, filter_type): 123 | img_float32 = np.float32(img) 124 | dft = cv2.dft(img_float32, flags=cv2.DFT_COMPLEX_OUTPUT) 125 | dft_shift = np.fft.fftshift(dft) 126 | fshift = dft_shift * (filter_type) 127 | f_ishift = np.fft.ifftshift(fshift) 128 | img_back = cv2.idft(f_ishift) 129 | img_back = cv2.magnitude(img_back[:, :, 0], img_back[:, :, 1]) 130 | # b0 = str(np.max(img)) 131 | # b01 = str(np.mean(img)) 132 | # b02 = str(np.min(img)) 133 | 134 | # b1 = str(np.max(img_back)) 135 | # b12 = str(np.mean(img_back)) 136 | # b2 = str(np.min(img_back)) 137 | # cv2.normalize(img_back, img_back, 0.0, 255.0, cv2.NORM_MINMAX) 138 | # img_back = img_back.astype(np.uint8) 139 | # cv2.equalizeHist(img_back,img_back) 140 | magnitude_spectrum = 20 * np.log( 141 | cv2.magnitude(dft_shift[:, :, 0], dft_shift[:, :, 1]) + 1 142 | ) 143 | # print(b0+'\t'+ b01 + '\t'+ b02+ '\t'+ b1+ '\t'+ b12 +'\t'+ b2 ) 144 | # cv2.normalize(img_back, img_back, 5.0, 30.0, cv2.NORM_MINMAX) 145 | return img_back, magnitude_spectrum 146 | -------------------------------------------------------------------------------- /src/icebreaker/five_figures.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from multiprocessing import Lock, Manager, Pool 3 | from pathlib import Path 4 | 5 | import mrcfile 6 | import numpy as np 7 | 8 | 9 | def single_mic_5fig(mic_name: str, r: int = 10000) -> str: 10 | with mrcfile.open(mic_name, "r", permissive=True) as mrc: 11 | img = mrc.data 12 | if not np.isnan(np.sum(img)): 13 | min = int(np.min(img) * r) 14 | q1 = int(np.quantile(img, 0.25) * r) 15 | median = int(np.median(img) * r) 16 | q3 = int(np.quantile(img, 0.75) * r) 17 | max = int(np.max(img) * r) 18 | csv_lines = f"{mic_name},{min},{q1},{median},{q3},{max}" 19 | else: 20 | csv_lines = "" 21 | return csv_lines 22 | 23 | 24 | def _process_mrc(img_path: str, csv_lines: list, lock: Lock, r: int) -> None: 25 | with mrcfile.open(img_path, "r", permissive=True) as mrc: 26 | img = mrc.data 27 | if not np.isnan(np.sum(img)): 28 | path = img_path[:-4] 29 | min = int(np.min(img) * r) 30 | q1 = int(np.quantile(img, 0.25) * r) 31 | median = int(np.median(img) * r) 32 | q3 = int(np.quantile(img, 0.75) * r) 33 | max = int(np.max(img) * r) 34 | lock.acquire() 35 | csv_lines.append(f"{path},{min},{q1},{median},{q3},{max}\n") 36 | lock.release() 37 | return None 38 | 39 | 40 | def main(grouped_mic_dir: str, cpus: int = 1, append: bool = False) -> list: 41 | manager = Manager() 42 | lock = manager.Lock() 43 | files = [str(filename) for filename in Path(grouped_mic_dir).glob("**/*.mrc")] 44 | r = 10000 45 | 46 | csv_lines = manager.list() 47 | with Pool(cpus) as p: 48 | p.starmap(_process_mrc, [(fl, csv_lines, lock, r) for fl in sorted(files)]) 49 | if append and Path("five_figs_test.csv").is_file(): 50 | with open("five_figs_test.csv", "a+") as f: 51 | for line in csv_lines: 52 | f.write(line) 53 | else: 54 | with open("five_figs_test.csv", "w") as f: 55 | f.write("path,min,q1,q2=median,q3,max\n") 56 | for line in csv_lines: 57 | f.write(line) 58 | return csv_lines 59 | 60 | 61 | if __name__ == "__main__": 62 | grouped_mic_dir = sys.argv[1] 63 | main(grouped_mic_dir) 64 | -------------------------------------------------------------------------------- /src/icebreaker/ib_group_template.star: -------------------------------------------------------------------------------- 1 | 2 | # version 30001 3 | 4 | data_job 5 | 6 | _rlnJobType 99 7 | _rlnJobIsContinue 0 8 | 9 | 10 | # version 30001 11 | 12 | data_joboptions_values 13 | 14 | loop_ 15 | _rlnJobOptionVariable #1 16 | _rlnJobOptionValue #2 17 | fn_exe ib_group.py 18 | in_mic micrographs/from/IB/JOB/GROUP/MODE.star 19 | in_part input/particles.star 20 | 21 | -------------------------------------------------------------------------------- /src/icebreaker/ib_job_template_FLATTEN_MODE.star: -------------------------------------------------------------------------------- 1 | 2 | # version 30001 3 | 4 | data_job 5 | 6 | _rlnJobType 99 7 | _rlnJobIsContinue 0 8 | 9 | 10 | # version 30001 11 | 12 | data_joboptions_values 13 | 14 | loop_ 15 | _rlnJobOptionVariable #1 16 | _rlnJobOptionValue #2 17 | fn_exe ib_job 18 | in_mic input/micrographs.star 19 | param1_label mode 20 | param1_value flatten 21 | 22 | -------------------------------------------------------------------------------- /src/icebreaker/ib_job_template_GROUP_MODE.star: -------------------------------------------------------------------------------- 1 | 2 | # version 30001 3 | 4 | data_job 5 | 6 | _rlnJobType 99 7 | _rlnJobIsContinue 0 8 | 9 | 10 | # version 30001 11 | 12 | data_joboptions_values 13 | 14 | loop_ 15 | _rlnJobOptionVariable #1 16 | _rlnJobOptionValue #2 17 | fn_exe ib_job 18 | in_mic input/micrographs.star 19 | param1_label mode 20 | param1_value group 21 | 22 | -------------------------------------------------------------------------------- /src/icebreaker/ice_groups.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import json 3 | import os 4 | import sys 5 | 6 | import gemmi 7 | import mrcfile 8 | import numpy as np 9 | 10 | from icebreaker import star_appender 11 | 12 | 13 | def splitall(path): 14 | allparts = [] 15 | while 1: 16 | parts = os.path.split(path) 17 | if parts[0] == path: # sentinel for absolute paths 18 | allparts.insert(0, parts[0]) 19 | break 20 | elif parts[1] == path: # sentinel for relative paths 21 | allparts.insert(0, parts[1]) 22 | break 23 | else: 24 | path = parts[0] 25 | allparts.insert(0, parts[1]) 26 | return allparts 27 | 28 | 29 | def main(starfile, mic_path): 30 | in_doc = gemmi.cif.read_file(starfile) 31 | data_as_dict = json.loads(in_doc.as_json())["particles"] 32 | 33 | micrographs_used = data_as_dict["_rlnmicrographname"] 34 | # print(micrographs_used) 35 | micrographs_unique = list(set(micrographs_used)) 36 | micrographs_unique.sort() 37 | num_mics = len(micrographs_unique) 38 | 39 | mic_coord = collections.OrderedDict() 40 | for mic in micrographs_unique: 41 | mic_coord[mic] = [i for i, e in enumerate(micrographs_used) if e == mic] 42 | 43 | ice_groups = [] 44 | for k in range(num_mics): 45 | print(f"{k+1} / {num_mics}") 46 | mic = micrographs_unique[k] 47 | # print(mic) 48 | # im_path = os.path.join(mic_path, os.path.split( 49 | # mic[:-4])[2:] + '_grouped.mrc') 50 | split_path = splitall(mic[:-4] + "_grouped.mrc")[2:] 51 | # print(split_path) 52 | # im_path = os.path.join(mic_path, *split_path) 53 | # print(im_path) 54 | mic_path_new = os.path.dirname(mic_path) 55 | im_path2 = os.path.join(mic_path_new, *split_path) 56 | # print(im_path2) 57 | with mrcfile.open(im_path2, "r+", permissive=True) as mrc: 58 | micro_now = mrc.data 59 | 60 | for part_ind in mic_coord[mic]: 61 | x1 = int(np.floor(data_as_dict["_rlncoordinatex"][part_ind])) 62 | y1 = int(np.floor(data_as_dict["_rlncoordinatey"][part_ind])) 63 | if micro_now is not None and np.isfinite(micro_now[y1][x1]): 64 | ice_groups.append(int(micro_now[y1][x1] * 10000)) 65 | else: 66 | ice_groups.append(-1) 67 | 68 | star_appender.update_star(starfile, ice_groups) 69 | 70 | return True 71 | 72 | 73 | if __name__ == "__main__": 74 | starfile = sys.argv[1] 75 | mic_path = sys.argv[2] 76 | main(starfile, mic_path) 77 | -------------------------------------------------------------------------------- /src/icebreaker/icebreaker_equalize_multi.py: -------------------------------------------------------------------------------- 1 | """ 2 | Input is group of motion corrected micrographs. 3 | Output is group of equalized images. 4 | """ 5 | 6 | import os 7 | import sys 8 | from multiprocessing import Pool 9 | 10 | import cv2 11 | import mrcfile 12 | import numpy as np 13 | 14 | from icebreaker import KMeans_segmenter as KMeans_seg 15 | from icebreaker import filter_designer as fd 16 | from icebreaker import local_mask as lm 17 | from icebreaker import window_mean as wm 18 | 19 | 20 | def load_img(img_path): 21 | with mrcfile.open(img_path, "r", permissive=True) as mrc: 22 | img = mrc.data 23 | return img 24 | 25 | 26 | def multigroup(filelist_full): 27 | # for filename in filelist: 28 | img = load_img(filelist_full) # (os.path.join(indir, filename)) 29 | splitpath = os.path.split(filelist_full) 30 | # Config params 31 | x_patches = 40 32 | y_patches = 40 33 | num_of_segments = 32 34 | 35 | final_image = equalize_im(img, x_patches, y_patches, num_of_segments) 36 | # final_image = img # !!!! TESTING 37 | 38 | # with mrcfile.new((path1+str(filename[:-4]) +'_' 39 | # +str(x_patches)+'x'+str(y_patches)+'x'+str(num_of_segments)+'flattened'+'.mrc'), 40 | # overwrite=True) as out_image: # Make fstring 41 | with mrcfile.new( 42 | os.path.join( 43 | splitpath[0] + "/flattened/" + splitpath[1][:-4] + "_flattened.mrc" 44 | ), 45 | overwrite=True, 46 | ) as out_image: # Make fstring 47 | out_image.set_data(final_image) 48 | 49 | 50 | def equalize_im(img, x_patches, y_patches, num_of_segments): 51 | filter_mask = fd.lowpass(img, 0.85, 20, "cos", 50) 52 | lowpass, mag = fd.filtering(img, filter_mask) 53 | lowpass = cv2.GaussianBlur(lowpass, (45, 45), 0) 54 | rolled = wm.window(lowpass, x_patches, y_patches) 55 | rolled_resized = cv2.resize(rolled, (185, 190), interpolation=cv2.INTER_NEAREST) 56 | rolled_resized = cv2.GaussianBlur(rolled_resized, (5, 5), 0) 57 | KNNsegmented = KMeans_seg.segmenter(rolled_resized, num_of_segments) 58 | 59 | # upscaled_region = cv2.resize( 60 | # KNNsegmented, (lowpass.shape[1], lowpass.shape[0]), interpolation=cv2.INTER_AREA) 61 | 62 | regions_vals = np.unique(KNNsegmented) 63 | averaged_loc = np.zeros( 64 | (lowpass.shape[0], lowpass.shape[1], num_of_segments), np.uint8 65 | ) 66 | res = np.zeros((lowpass.shape[0], lowpass.shape[1]), np.uint8) 67 | for i in range(len(regions_vals)): 68 | averaged_loc[:, :, i] = lm.local_mask(lowpass, KNNsegmented, regions_vals[i]) 69 | res[:, :] += averaged_loc[:, :, i] 70 | final_image = res # cv2.GaussianBlur(res, (45, 45), 0) 71 | 72 | return final_image 73 | 74 | 75 | def main(indir, cpus): 76 | outdir = "flattened" 77 | path1 = os.path.join(indir, outdir) 78 | 79 | try: 80 | os.mkdir(path1) 81 | except OSError: 82 | print(f"Creation of the directory {path1} failed") 83 | else: 84 | print(f"Successfully created the directory {path1}") 85 | 86 | filelist = [] 87 | for filename in os.listdir(indir): 88 | if filename.endswith(".mrc"): 89 | finalpath = os.path.join(indir, filename) 90 | filelist.append(finalpath) 91 | else: 92 | continue 93 | 94 | # cc = 0 95 | # for filename in filelist: 96 | # img = load_img(os.path.join(indir, filename)) 97 | 98 | # Config params 99 | # x_patches = 40 100 | # y_patches = 40 101 | # num_of_segments = 32 102 | 103 | # final_image = equalize_im(img, x_patches, y_patches, num_of_segments) 104 | # final_image = img # !!!! TESTING 105 | 106 | # with mrcfile.new((path1+str(filename[:-4]) +'_'+str(x_patches)+'x'+str(y_patches)+ 107 | # 'x'+str(num_of_segments)+'flattened'+'.mrc'), overwrite=True) as out_image: 108 | # Make fstring 109 | # with mrcfile.new(os.path.join(path1, filename[:-4] + f'_{outdir}.mrc'), 110 | # overwrite=True) as out_image: # Make fstring 111 | # out_image.set_data(final_image) 112 | 113 | # cc += 1 114 | # print(f'{cc}/{len(filelist)}') 115 | with Pool(cpus) as p: 116 | p.map(multigroup, filelist) 117 | 118 | return True 119 | 120 | 121 | if __name__ == "__main__": 122 | indir = sys.argv[1] 123 | batch_size = sys.argv[2] 124 | main(indir, batch_size) 125 | -------------------------------------------------------------------------------- /src/icebreaker/icebreaker_icegroups_multi.py: -------------------------------------------------------------------------------- 1 | """ 2 | Input is group of motion corrected micrographs. 3 | Output is group of icegrouped images. 4 | """ 5 | import os 6 | import sys 7 | import time 8 | from multiprocessing import Pool 9 | 10 | import cv2 11 | import mrcfile 12 | import numpy as np 13 | 14 | from icebreaker import KMeans_segmenter as KMeans_seg 15 | from icebreaker import filter_designer as fd 16 | from icebreaker import original_mask_fast as omf 17 | from icebreaker import window_mean as wm 18 | 19 | 20 | def load_img(img_path): 21 | with mrcfile.open(img_path, "r", permissive=True) as mrc: 22 | # mrc.header.map = mrcfile.constants.MAP_ID 23 | img = mrc.data 24 | return img 25 | 26 | 27 | def multigroup(filelist_full): 28 | # for filename in filelist: 29 | img = load_img(filelist_full) 30 | splitpath = os.path.split(filelist_full) 31 | # print(splitpath[0]) 32 | # Config params 33 | x_patches = 40 34 | y_patches = 40 35 | num_of_segments = 16 36 | 37 | final_image = ice_grouper(img, x_patches, y_patches, num_of_segments) 38 | # final_image = img # !!!!! FOR TESTING 39 | 40 | # with mrcfile.new((path1+str(filename[:-4]) +'_'+str(x_patches)+ 41 | # 'x'+str(y_patches)+'x'+str(num_of_segments)+'_original_mean'+'.mrc'), 42 | # overwrite=True) as out_image: 43 | with mrcfile.new( 44 | os.path.join(splitpath[0] + "/grouped/" + splitpath[1][:-4] + "_grouped.mrc"), 45 | overwrite=True, 46 | ) as out_image: # Make fstring 47 | out_image.set_data(final_image) 48 | 49 | 50 | def ice_grouper(img, x_patches, y_patches, num_of_segments): 51 | filter_mask = fd.lowpass(img, 0.85, 20, "cos", 50) 52 | lowpass, mag = fd.filtering(img, filter_mask) 53 | lowpass = cv2.GaussianBlur(lowpass, (45, 45), 0) 54 | lowpass12 = img 55 | rolled = wm.window(lowpass, x_patches, y_patches) 56 | rolled_resized = cv2.resize(rolled, (185, 190), interpolation=cv2.INTER_AREA) 57 | rolled_resized = cv2.GaussianBlur(rolled_resized, (5, 5), 0) 58 | KNNsegmented = KMeans_seg.segmenter(rolled_resized, num_of_segments) 59 | # upscaled_region = cv2.resize( 60 | # KNNsegmented, (lowpass.shape[1], lowpass.shape[0]), interpolation=cv2.INTER_AREA 61 | # ) 62 | 63 | regions_vals = np.unique(KNNsegmented) 64 | averaged_loc = np.zeros( 65 | (lowpass.shape[0], lowpass.shape[1], num_of_segments), np.float32 66 | ) 67 | res = np.zeros((lowpass.shape[0], lowpass.shape[1]), np.float32) 68 | for i in range(len(regions_vals)): 69 | averaged_loc[:, :, i] = omf.original_mask( 70 | lowpass, KNNsegmented, regions_vals[i], img, lowpass12 71 | ) 72 | res[:, :] += averaged_loc[:, :, i] 73 | 74 | final_image = res 75 | 76 | return final_image 77 | 78 | 79 | def main(indir, cpus): 80 | outdir = "grouped" 81 | path1 = os.path.join(indir, outdir) 82 | try: 83 | os.mkdir(path1) 84 | except OSError: 85 | print(f"Creation of the directory {path1} failed") 86 | else: 87 | print(f"Successfully created the directory {path1}") 88 | start_time = time.time() 89 | 90 | filelist = [] 91 | for filename in os.listdir(indir): 92 | if filename.endswith(".mrc"): 93 | finalpath = os.path.join(indir, filename) 94 | filelist.append(finalpath) 95 | # print(filelist) 96 | else: 97 | continue 98 | 99 | # cc = 0 100 | 101 | with Pool(cpus) as p: 102 | p.map(multigroup, filelist) 103 | 104 | print("------ %s sec------" % (time.time() - start_time)) 105 | 106 | return True 107 | 108 | 109 | if __name__ == "__main__": 110 | indir = sys.argv[1] 111 | batch_size = sys.argv[2] 112 | main(indir, batch_size) 113 | -------------------------------------------------------------------------------- /src/icebreaker/listdir.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import mrcfile 3 | 4 | import os 5 | 6 | x_patches = 40 7 | y_patches = 40 8 | num_of_segments = 16 9 | 10 | cc = 0 11 | 12 | 13 | def load_img(img_path): 14 | with mrcfile.open(img_path, "r", permissive=True) as mrc: 15 | img = mrc.data 16 | # plt.imshow(img, cmap='gray') 17 | # plt.show() 18 | return img 19 | 20 | 21 | outdir = "icegroups/" 22 | # define image path, number of patches, num of regions 23 | 24 | filelist = [] 25 | outputs = [] 26 | dir = sys.argv[1] 27 | path1 = dir + outdir 28 | try: 29 | os.mkdir(path1) 30 | except OSError: 31 | print("Creation of the directory %s failed" % path1) 32 | else: 33 | print("Successfully created the directory %s " % path1) 34 | 35 | for filename in os.listdir(dir): 36 | # if filename.endswith(".mrc"): #or filename.endswith(".py"): 37 | if filename.endswith(".mrc"): 38 | filelist.append(filename) 39 | print(filelist) 40 | # outputs.append(str(filename[:-4]) +'_'+str(x_patches)+'_' 41 | # +str(y_patches)+'_'+str(num_of_segments)+'.mrc') 42 | else: 43 | continue 44 | 45 | for filename in filelist: 46 | img = load_img(dir + filename) 47 | print(path1 + filename) 48 | with mrcfile.new( 49 | ( 50 | path1 51 | + str(filename[:-4]) 52 | + "_" 53 | + str(x_patches) 54 | + "x" 55 | + str(y_patches) 56 | + "x" 57 | + str(num_of_segments) 58 | + "_original_mean" 59 | + ".mrc" 60 | ), 61 | overwrite=True, 62 | ) as out_image: 63 | print(filename) 64 | out_image.set_data(img) 65 | -------------------------------------------------------------------------------- /src/icebreaker/local_mask.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | 5 | def local_mask(lowpass1, img2, val): 6 | reg3 = np.ones((img2.shape[0], img2.shape[1]), np.uint8) 7 | # coord0 = [] 8 | # coord1 = [] 9 | # k = np.zeros((img2.shape[0], img2.shape[1]), np.uint8) 10 | for u in range(img2.shape[0]): 11 | for v in range(img2.shape[1]): 12 | if img2[u, v] == val: 13 | reg3[u, v] = 1 14 | else: 15 | reg3[u, v] = 0 16 | 17 | up_reg3 = cv2.resize( 18 | reg3, (lowpass1.shape[1], lowpass1.shape[0]), interpolation=cv2.INTER_NEAREST 19 | ) 20 | # mean_vals = cv2.mean(lowpass1, up_reg3)[0] 21 | # up_cover4 = up_reg3 22 | up_reg3 = up_reg3 * lowpass1 23 | # Negative of ROI 24 | # cover3 = 1 - reg3 25 | # Upscaled negative 26 | # up_cover3 = cv2.resize(cover3, (lowpass1.shape[1], lowpass1.shape[0])) 27 | 28 | mask3 = cv2.equalizeHist(up_reg3) 29 | 30 | return mask3 31 | -------------------------------------------------------------------------------- /src/icebreaker/original_mask.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | 5 | def original_mask(lowpass1, img2, val, original, lowpass12): 6 | reg3 = np.ones((img2.shape[0], img2.shape[1]), np.uint8) 7 | orig_mean = np.zeros((lowpass12.shape[0], lowpass12.shape[1]), np.float32) 8 | coord0 = [] 9 | coord1 = [] 10 | # global_mean = np.mean(lowpass12) 11 | # k = np.zeros((img2.shape[0], img2.shape[1]), np.uint8) 12 | for u in range(img2.shape[0]): 13 | for v in range(img2.shape[1]): 14 | if img2[u, v] == val: 15 | reg3[u, v] = 1 16 | else: 17 | reg3[u, v] = 0 18 | 19 | up_reg3 = cv2.resize( 20 | reg3, (lowpass1.shape[1], lowpass1.shape[0]), interpolation=cv2.INTER_AREA 21 | ) 22 | # mean_vals = cv2.mean(lowpass1, up_reg3)[0] 23 | up_reg3 = up_reg3 * lowpass1 24 | x_mean = np.mean(original[np.nonzero(up_reg3)]) 25 | # x_min = np.min(original[np.nonzero(up_reg3)]) 26 | # x_max = np.max(original[np.nonzero(up_reg3)]) 27 | # x_mean2 = np.mean(lowpass12[np.nonzero(up_reg3)]) 28 | # x_min2 = np.min(lowpass12[np.nonzero(up_reg3)]) 29 | # x_max2 = np.max(lowpass12[np.nonzero(up_reg3)]) 30 | # x_mean1 = int(np.mean(lowpass1[np.nonzero(up_reg3)])) 31 | # x_min1 = np.min(lowpass1[np.nonzero(up_reg3)]) 32 | # x_max1 = np.max(lowpass1[np.nonzero(up_reg3)]) 33 | 34 | # Negative of ROI 35 | # cover3 = 1 - reg3 36 | # Upscaled negative 37 | # up_cover3 = cv2.resize(cover3, (lowpass1.shape[1], lowpass1.shape[0])) 38 | mask3 = up_reg3 39 | coord0 = np.nonzero(mask3)[0] 40 | coord1 = np.nonzero(mask3)[1] 41 | for i in range(len(coord1)): 42 | orig_mean[coord0[i]][coord1[i]] = x_mean # int(mean_vals) 43 | 44 | return orig_mean 45 | -------------------------------------------------------------------------------- /src/icebreaker/original_mask_fast.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | 4 | 5 | def original_mask(lowpass1, img2, val, original, lowpass12): 6 | reg3 = np.ones((img2.shape[0], img2.shape[1]), np.uint8) 7 | # orig_mean = np.zeros((lowpass12.shape[0], lowpass12.shape[1]), np.float32) 8 | # coord0 = [] 9 | # coord1 = [] 10 | # global_mean = np.mean(lowpass12) 11 | # k = np.zeros((img2.shape[0], img2.shape[1]), np.uint8) 12 | for u in range(img2.shape[0]): 13 | for v in range(img2.shape[1]): 14 | if img2[u, v] == val: 15 | reg3[u, v] = 1 16 | else: 17 | reg3[u, v] = 0 18 | 19 | up_reg3 = cv2.resize( 20 | reg3, (lowpass1.shape[1], lowpass1.shape[0]), interpolation=cv2.INTER_NEAREST 21 | ) 22 | # mean_vals = cv2.mean(lowpass1, up_reg3)[0] 23 | up_reg3 = up_reg3 * lowpass1 24 | x_mean = np.mean(original[np.nonzero(up_reg3)]) 25 | 26 | # Negative of ROI 27 | # cover3 = 1 - reg3 28 | 29 | # Upscaled negative 30 | up_reg3 = cv2.resize( 31 | reg3, (lowpass1.shape[1], lowpass1.shape[0]), interpolation=cv2.INTER_NEAREST 32 | ) 33 | orig_mean = up_reg3 * x_mean 34 | return orig_mean 35 | -------------------------------------------------------------------------------- /src/icebreaker/plot_boxplots.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import mrcfile 5 | import os 6 | import argparse 7 | 8 | 9 | def plot_boxes(indir): 10 | files = [] 11 | 12 | for filename in os.listdir(indir): 13 | if filename.endswith(".mrc"): 14 | finalpath = os.path.join(indir, filename) 15 | files.append(finalpath) 16 | else: 17 | continue 18 | files = sorted(files) 19 | vector_data = [] 20 | for img_path in files: 21 | with mrcfile.open(img_path, "r", permissive=True) as mrc: 22 | img = mrc.data 23 | a_vector = np.reshape(img, -1) 24 | vector_data.append(a_vector) 25 | labels = [] 26 | for label in files: 27 | splitpath = os.path.split(label) 28 | labels.append(splitpath[-1][:-4]) 29 | fig = plt.figure() 30 | ax = fig.add_subplot(111) 31 | ax.boxplot(vector_data, 0, "x") 32 | ax.set_xticklabels(labels, rotation=90, fontsize=8) 33 | 34 | plt.yticks(fontsize=10) 35 | plt.ylabel("Pixel intensity", fontsize=12) 36 | plt.xlabel("Micrographs", fontsize=12) 37 | plt.grid(axis="y") 38 | plt.tight_layout() 39 | plt.show() 40 | 41 | 42 | def main(): 43 | parser = argparse.ArgumentParser() 44 | parser.add_argument( 45 | "in_dir", help="Path to directory with micrographs_grouped", type=str 46 | ) 47 | args = parser.parse_args() 48 | indir = args.in_dir 49 | plot_boxes(indir) 50 | 51 | 52 | if __name__ == "__main__": 53 | sys.exit(main()) 54 | -------------------------------------------------------------------------------- /src/icebreaker/show_gradient.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | 4 | 5 | def show_gradient(image): 6 | xx, yy = np.mgrid[0 : image.shape[0], 0 : image.shape[1]] 7 | fig = plt.figure(figsize=(15, 15)) 8 | ax = fig.gca(projection="3d") 9 | ax.plot_surface(xx, yy, image, rstride=1, cstride=1, cmap=plt.cm.gray, linewidth=2) 10 | ax.view_init(80, 30) 11 | plt.show() 12 | -------------------------------------------------------------------------------- /src/icebreaker/star_appender.py: -------------------------------------------------------------------------------- 1 | import gemmi 2 | import os 3 | 4 | 5 | def update_star(starfile, ice_groups): 6 | new_document = gemmi.cif.Document() 7 | in_doc = gemmi.cif.read_file(starfile) 8 | block_optics = in_doc.find_block("optics") 9 | if block_optics is not None: 10 | block_new_optics = new_document.add_new_block("optics") 11 | 12 | for item in block_optics: 13 | if item.loop is not None: 14 | table_optics = item.loop 15 | 16 | tags_optics = table_optics.tags 17 | table_optics = block_optics.find(tags_optics) 18 | loop_optics = block_new_optics.init_loop("", tags_optics) # make temp new table 19 | 20 | for row in table_optics: 21 | loop_optics.add_row(list(row)) 22 | 23 | block = in_doc.find_block("particles") 24 | block_new_particles = new_document.add_new_block("particles") 25 | 26 | for item in block: 27 | if item.loop is not None: 28 | table = item.loop 29 | 30 | tags = table.tags 31 | table = block.find(tags) 32 | 33 | # add new tag: named this to work in relion, but really it is 34 | # ice group for each particle (_ibIceGroup) 35 | tags.append("_rlnHelicalTubeID") 36 | # tags.append('_ibIceGroup') 37 | loop = block_new_particles.init_loop("", tags) # make temp new table 38 | for i in range(len(table)): 39 | row = table[i] 40 | new_row = list(row) 41 | new_row.append(f"{ice_groups[i]}") 42 | loop.add_row(new_row) # update temp new table with all data 43 | 44 | new_document.write_file("particles.star") 45 | 46 | 47 | def mic_star(starfile, job, mode): 48 | in_doc = gemmi.cif.read_file(starfile) 49 | block = in_doc.find_block("micrographs") 50 | new_document = gemmi.cif.Document() 51 | block_one = new_document.add_new_block("micrographs") 52 | 53 | for item in block: 54 | if item.loop is not None: 55 | table = item.loop 56 | 57 | tags = table.tags 58 | tags.remove("_rlnMicrographName") 59 | table = block.find(tags) 60 | column = block.find(["_rlnMicrographName"]) 61 | 62 | tags.insert(0, "_rlnMicrographName") 63 | 64 | loop = block_one.init_loop("", tags) # make temp new table 65 | for i in range(len(table)): 66 | row = table[i] 67 | new_row = list(row) 68 | mic_name = list(column[i])[0] 69 | path_head, path_tail = os.path.split(mic_name) 70 | # mic_name2 = Path(mic_name).parts 71 | # mic_name_new = os.path.join(mic_name2[-2], mic_name2[-1][:-4] + f'_{mode}ed.mrc') 72 | 73 | mic_name_new = os.path.join(path_head, path_tail[:-4] + f"_{mode}ed.mrc") 74 | new_row.insert(0, os.path.join(job, mic_name_new)) 75 | 76 | loop.add_row(new_row) # update temp new table with all data 77 | 78 | new_document.write_file(f"{mode}ed_micrographs.star") 79 | 80 | 81 | if __name__ == "__main__": 82 | starfile = ( 83 | "/home/lexi/Documents/Diamond/ICEBREAKER/test_data/corrected_micrographs.star" 84 | ) 85 | job = "External" 86 | mode = "flatten" 87 | mic_star(starfile, job, mode) 88 | -------------------------------------------------------------------------------- /src/icebreaker/window_mean.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def window(img, x_patches, y_patches): 5 | rows, cols = img.shape 6 | x_patch_size = int(cols / x_patches) 7 | y_patch_size = int(rows / y_patches) 8 | mean_map = img.copy() 9 | for i in range(x_patches): 10 | for j in range(y_patches): 11 | if i == x_patches - 1 and j != y_patches - 1: 12 | mean_map[ 13 | (j * y_patch_size) : ((j + 1) * y_patch_size), (i * x_patch_size) : 14 | ] = np.mean( 15 | mean_map[ 16 | (j * y_patch_size) : ((j + 1) * y_patch_size), 17 | (i * x_patch_size) :, 18 | ] 19 | ) 20 | elif j == y_patches - 1 and i != x_patches - 1: 21 | mean_map[ 22 | (j * y_patch_size) :, (i * x_patch_size) : ((i + 1) * x_patch_size) 23 | ] = np.mean( 24 | mean_map[ 25 | (j * y_patch_size) :, 26 | (i * x_patch_size) : ((i + 1) * x_patch_size), 27 | ] 28 | ) 29 | elif j == y_patches - 1 and i == x_patches - 1: 30 | mean_map[(j * y_patch_size) :, (i * x_patch_size) :] = np.mean( 31 | mean_map[(j * y_patch_size) :, (i * x_patch_size) :] 32 | ) 33 | else: 34 | mean_map[ 35 | (j * y_patch_size) : ((j + 1) * y_patch_size), 36 | (i * x_patch_size) : ((i + 1) * x_patch_size), 37 | ] = np.mean( 38 | mean_map[ 39 | (j * y_patch_size) : ((j + 1) * y_patch_size), 40 | (i * x_patch_size) : ((i + 1) * x_patch_size), 41 | ] 42 | ) 43 | 44 | return mean_map 45 | --------------------------------------------------------------------------------