Source code for mdt.gui.model_fit.tabs.generate_roi_mask_tab

import os
from textwrap import dedent

from mdt.lib.nifti import load_nifti
from PyQt5.QtCore import pyqtSlot, QObject, pyqtSignal
from PyQt5.QtWidgets import QFileDialog

from mdt.visualization.maps.base import SimpleDataInfo, MapPlotConfig
from mdt.gui.maps_visualizer.main import start_gui
from mdt.gui.model_fit.design.ui_generate_roi_mask_tab import Ui_GenerateROIMaskTabContent
from mdt.gui.utils import function_message_decorator, image_files_filters, MainTab, get_script_file_header_text
from mdt.utils import split_image_path, write_slice_roi

__author__ = 'Robbert Harms'
__date__ = "2016-06-26"
__maintainer__ = "Robbert Harms"
__email__ = "robbert@xkls.nl"


[docs]class GenerateROIMaskTab(MainTab, Ui_GenerateROIMaskTabContent, QObject): def __init__(self, shared_state, computations_thread): super().__init__() self._shared_state = shared_state self._computations_thread = computations_thread self._generate_mask_worker = GenerateROIMaskWorker()
[docs] def setupUi(self, tab_content): super().setupUi(tab_content) self.selectMaskButton.clicked.connect(lambda: self._select_mask()) self.selectOutputFileInput.clicked.connect(lambda: self._select_output_file()) self.viewButton.clicked.connect(self.view_mask) self.generateButton.clicked.connect(self.generate_roi_mask) self.selectedMaskText.textChanged.connect(self._check_enable_action_buttons) self.selectedOutputFileText.textChanged.connect(self._check_enable_action_buttons) self.selectedMaskText.textChanged.connect(self.mask_file_changed) self.dimensionInput.valueChanged.connect(self.update_dimension) self.sliceInput.valueChanged.connect(self.update_slice_index)
def _select_mask(self): initial_dir = self._shared_state.base_dir if self.selectedMaskText.text() != '': initial_dir = self.selectedMaskText.text() open_file, used_filter = QFileDialog().getOpenFileName( caption='Select the brain mask', directory=initial_dir, filter=';;'.join(image_files_filters)) if os.path.isfile(open_file): self.selectedMaskText.setText(open_file) self.mask_file_changed() self._shared_state.base_dir = os.path.dirname(open_file) def _select_output_file(self): output_file_name, used_filter = QFileDialog().getSaveFileName( caption='Select the output file', directory=self._shared_state.base_dir, filter=';;'.join(image_files_filters)) if output_file_name: self.selectedOutputFileText.setText(output_file_name) def _check_enable_action_buttons(self): self.generateButton.setEnabled(os.path.isfile(self.selectedMaskText.text()) and os.path.isdir(os.path.dirname(self.selectedOutputFileText.text()))) self.viewButton.setEnabled(os.path.isfile(self.selectedMaskText.text()) and os.path.isfile(self.selectedOutputFileText.text()))
[docs] def view_mask(self): data = SimpleDataInfo({'Original mask': load_nifti(self.selectedMaskText.text()).get_data(), 'Slice mask': load_nifti(self.selectedOutputFileText.text()).get_data()}) config = MapPlotConfig() config.dimension = self.dimensionInput.value() config.slice_index = self.sliceInput.value() config.maps_to_show = ['Original mask', 'Slice mask'] start_gui(data=data, config=config, app_exec=False)
[docs] def generate_roi_mask(self): kwargs = dict(mask=self.selectedMaskText.text(), output=self.selectedOutputFileText.text(), dimension=self.dimensionInput.value(), slice=self.sliceInput.value()) self._generate_mask_worker.set_args(**kwargs) self._computations_thread.start() self._generate_mask_worker.moveToThread(self._computations_thread) self._generate_mask_worker.starting.connect(self._computations_thread.starting) self._generate_mask_worker.finished.connect(self._computations_thread.finished) self._generate_mask_worker.starting.connect(lambda: self.generateButton.setEnabled(False)) self._generate_mask_worker.finished.connect(lambda: self.generateButton.setEnabled(True)) self._generate_mask_worker.finished.connect(lambda: self.viewButton.setEnabled(True)) image_path = split_image_path(self.selectedOutputFileText.text()) script_basename = os.path.join(image_path[0], 'scripts', 'create_roi_mask_' + image_path[1]) if not os.path.isdir(os.path.join(image_path[0], 'scripts')): os.makedirs(os.path.join(image_path[0], 'scripts')) self._generate_mask_worker.finished.connect( lambda: self._write_python_script_file(script_basename + '.py', **kwargs)) self._generate_mask_worker.finished.connect( lambda: self._write_bash_script_file(script_basename + '.sh', **kwargs)) self._generate_mask_worker.starting.emit()
[docs] @pyqtSlot(int) def update_dimension(self, value): if os.path.isfile(self.selectedMaskText.text()): self.update_slice_selector() self.update_output_file_text()
[docs] @pyqtSlot(int) def update_slice_index(self, value): self.update_output_file_text()
[docs] def mask_file_changed(self): self.dimensionInput.setValue(2) self.update_slice_selector() self.update_output_file_text()
[docs] def update_slice_selector(self): if os.path.isfile(self.selectedMaskText.text()): dimension_max = load_nifti(self.selectedMaskText.text()).shape[self.dimensionInput.value()] self.sliceInput.setMaximum(dimension_max) self.sliceInput.setValue(dimension_max // 2.0) self.maxSliceLabel.setText(str(dimension_max)) else: self.sliceInput.setValue(0) self.maxSliceLabel.setText('x')
[docs] def update_output_file_text(self): if os.path.isfile(self.selectedMaskText.text()): folder, basename, ext = split_image_path(self.selectedMaskText.text()) folder_base = os.path.join(folder, basename) if self.selectedOutputFileText.text() == '': self.selectedOutputFileText.setText('{}_{}_{}.nii.gz'.format(folder_base, self.dimensionInput.value(), self.sliceInput.value())) elif self.selectedOutputFileText.text()[0:len(folder_base)] == folder_base: self.selectedOutputFileText.setText('{}_{}_{}.nii.gz'.format(folder_base, self.dimensionInput.value(), self.sliceInput.value()))
def _write_python_script_file(self, output_file, **kwargs): with open(output_file, 'w') as f: f.write('#!/usr/bin/env python') f.write(dedent(''' {header} import mdt mdt.write_slice_roi( {mask!r}, {dimension!r}, {slice!r}, {output!r}, overwrite_if_exists=True) ''').format(header=get_script_file_header_text({'Purpose': 'Generated a slice ROI mask'}), **kwargs)) def _write_bash_script_file(self, output_file, **kwargs): with open(output_file, 'w') as f: f.write('#!/usr/bin/env bash') f.write(dedent(''' {header} mdt-create-roi-slice "{mask}" -d {dimension} -s {slice} -o "{output}" ''').format(header=get_script_file_header_text({'Purpose': 'Generated a slice ROI mask'}), **kwargs))
[docs]class GenerateROIMaskWorker(QObject): starting = pyqtSignal() finished = pyqtSignal() def __init__(self): super().__init__() self.starting.connect(self.run) self._args = [] self._kwargs = {}
[docs] def set_args(self, *args, **kwargs): self._args = args self._kwargs = kwargs
[docs] @function_message_decorator('Started with generating a slice ROI', 'Finished generating a slice ROI') def run(self): write_slice_roi(self._kwargs['mask'], self._kwargs['dimension'], self._kwargs['slice'], self._kwargs['output'], overwrite_if_exists=True) self.finished.emit()