class FileNameIteratorTest(unittest.TestCase):
    def setUp(self):
        self.app = QtGui.QApplication([])
        self.filename_iterator = FileNameIterator()

    def tearDown(self):
        del self.app

    def test_get_next_filename_with_existent_file(self):
        filename = 'image_001.tif'
        self.filename_iterator.update_filename(os.path.join(data_path, filename))
        new_filename = os.path.basename(self.filename_iterator.get_next_filename())
        self.assertEqual(new_filename, 'image_002.tif')

    def test_get_next_filename_with_non_existent_file(self):
        filename = 'image_002.tif'
        self.filename_iterator.update_filename(os.path.join(data_path, filename))
        self.assertEqual(self.filename_iterator.get_next_filename(), None)

    def test_get_next_filename_with_larger_Step(self):
        filename = 'image_000.tif'
        self.filename_iterator.update_filename(os.path.join(data_path, filename))
        new_filename = os.path.basename(self.filename_iterator.get_next_filename(step=2))
        self.assertEqual(new_filename, 'image_002.tif')

    def test_get_previous_filename_with_existent_file(self):
        filename = 'image_002.tif'
        self.filename_iterator.update_filename(os.path.join(data_path, filename))
        new_filename = os.path.basename(self.filename_iterator.get_previous_filename())
        self.assertEqual(new_filename, 'image_001.tif')

    def test_get_previous_filename_with_non_existent_file(self):
        filename = 'image_001.tif'
        self.filename_iterator.update_filename(os.path.join(data_path, filename))
        self.assertEqual(self.filename_iterator.get_previous_filename(), None)

    def test_get_previous_filename_with_larger_Step(self):
        filename = 'image_003.tif'
        self.filename_iterator.update_filename(os.path.join(data_path, filename))
        new_filename = os.path.basename(self.filename_iterator.get_previous_filename(step=2))
        self.assertEqual(new_filename, 'image_001.tif')
Example #2
0
class SpectrumModel(QtCore.QObject):
    """
    Main Spectrum handling class. Supporting several features:
      - loading spectra from any tabular source (readable by numpy)
      - having overlays
      - setting overlays as background
      - spectra and overlays can be scaled and have offset values

    all changes to the internal data throw pyqtSignals.
    """
    spectrum_changed = QtCore.pyqtSignal()
    overlay_changed = QtCore.pyqtSignal(int)  # changed index
    overlay_added = QtCore.pyqtSignal()
    overlay_removed = QtCore.pyqtSignal(int)  # removed index
    overlay_set_as_bkg = QtCore.pyqtSignal(int)  # index set as background
    overlay_unset_as_bkg = QtCore.pyqtSignal(int)  # index unset os background

    def __init__(self):
        super(SpectrumModel, self).__init__()
        self.spectrum = Spectrum()
        self.overlays = []
        self.phases = []

        self.file_iteration_mode = 'number'
        self.file_name_iterator = FileNameIterator()

        self.bkg_ind = -1
        self.spectrum_filename = ''

    def set_spectrum(self, x, y, filename='', unit=''):
        """
        set the current data spectrum.
        :param x: x-values
        :param y: y-values
        :param filename: name for the spectrum, defaults to ''
        :param unit: unit for the x values
        """
        self.spectrum_filename = filename
        self.spectrum.data = (x, y)
        self.spectrum.name = get_base_name(filename)
        self.unit = unit
        self.spectrum_changed.emit()

    def load_spectrum(self, filename):
        """
        Loads a spectrum from a tabular spectrum file (2 column txt file)
        :param filename: filename of the data file
        """
        logger.info("Load spectrum: {0}".format(filename))
        self.spectrum_filename = filename

        skiprows = 0
        if filename.endswith('.chi'):
            skiprows = 4
        self.spectrum.load(filename, skiprows)
        self.file_name_iterator.update_filename(filename)
        self.spectrum_changed.emit()

    def save_spectrum(self, filename, header=None, subtract_background=False):
        """
        Saves the current data spectrum.
        :param filename: where to save
        :param header: you can specify any specific header
        :param subtract_background: whether or not the background set will be used for saving or not
        """
        if subtract_background:
            x, y = self.spectrum.data
        else:
            x, y = self.spectrum._original_x, self.spectrum._original_y

        file_handle = open(filename, 'w')
        num_points = len(x)

        if filename.endswith('.chi'):
            if header is None or header == '':
                file_handle.write(filename + '\n')
                file_handle.write(self.unit + '\n\n')
                file_handle.write("       {0}\n".format(num_points))
            else:
                file_handle.write(header)
            for ind in xrange(num_points):
                file_handle.write(' {0:.7E}  {1:.7E}\n'.format(x[ind], y[ind]))
        else:
            if header is not None:
                file_handle.write(header)
                file_handle.write('\n')
            for ind in xrange(num_points):
                file_handle.write('{0:.9E}  {1:.9E}\n'.format(x[ind], y[ind]))
        file_handle.close()

    def get_spectrum(self):
        return self.spectrum

    def load_next_file(self, step=1):
        """
        Loads the next file from a sequel of filenames (e.g. *_001.xy --> *_002.xy)
        It assumes that the file numbers are at the end of the filename
        """
        next_file_name = self.file_name_iterator.get_next_filename(mode=self.file_iteration_mode, step=step)
        if next_file_name is not None:
            self.load_spectrum(next_file_name)
            return True
        return False

    def load_previous_file(self, step=1):
        """
        Loads the previous file from a sequel of filenames (e.g. *_002.xy --> *_001.xy)
        It assumes that the file numbers are at the end of the filename
        """
        next_file_name = self.file_name_iterator.get_previous_filename(mode=self.file_iteration_mode, step=step)
        if next_file_name is not None:
            self.load_spectrum(next_file_name)
            return True
        return False

    def set_file_iteration_mode(self, mode):
        if mode == 'number':
            self.file_iteration_mode = 'number'
            self.file_name_iterator.create_timed_file_list = False
        elif mode == 'time':
            self.file_iteration_mode = 'time'
            self.file_name_iterator.create_timed_file_list = True
            self.file_name_iterator.update_filename(self.filename)

    def add_overlay(self, x, y, name=''):
        """
        Adds an overlay to the list of overlays
        :param x: x-values
        :param y: y-values
        :param name: name of overlay to be used for displaying etc.
        """
        self.overlays.append(Spectrum(x, y, name))
        self.overlay_added.emit()

    def remove_overlay(self, ind):
        """
        Removes an overlay from the list of overlays
        :param ind: index of the overlay
        """
        if ind >= 0:
            del self.overlays[ind]
            if self.bkg_ind > ind:
                self.bkg_ind -= 1
            elif self.bkg_ind == ind:
                self.spectrum.unset_background_spectrum()
                self.bkg_ind = -1
                self.spectrum_changed.emit()
            self.overlay_removed.emit(ind)

    def get_overlay(self, ind):
        """
        :param ind: overlay ind
        :return: returns overlay if existent or None if it does not exist
        :type return: Spectrum
        """
        try:
            return self.overlays[ind]
        except IndexError:
            return None


    def add_spectrum_as_overlay(self):
        """
        Adds the current data spectrum as overlay to the list of overlays
        """
        current_spectrum = deepcopy(self.spectrum)
        overlay_spectrum = Spectrum(current_spectrum.x,
                                    current_spectrum.y,
                                    current_spectrum.name)
        self.overlays.append(overlay_spectrum)
        self.overlay_added.emit()

    def add_overlay_file(self, filename):
        """
        Reads a 2-column (x,y) text file and adds it as overlay to the list of overlays
        :param filename: path of the file to be loaded
        """
        self.overlays.append(Spectrum())
        self.overlays[-1].load(filename)
        self.overlay_added.emit()

    def get_overlay_name(self, ind):
        """
        :param ind: overlay index
        """
        return self.overlays[-1].name

    def set_overlay_scaling(self, ind, scaling):
        """
        Sets the scaling of the specified overlay
        :param ind: index of the overlay
        :param scaling: new scaling value
        """
        self.overlays[ind].scaling = scaling
        self.overlay_changed.emit(ind)
        if self.bkg_ind == ind:
            self.spectrum_changed.emit()

    def get_overlay_scaling(self, ind):
        """
        Returns the scaling of the specified overlay
        :param ind: index of the overlay
        :return: scaling value
        """
        return self.overlays[ind].scaling

    def set_overlay_offset(self, ind, offset):
        """
        Sets the offset of the specified overlay
        :param ind: index of the overlay
        :param offset: new offset value
        """
        self.overlays[ind].offset = offset
        self.overlay_changed.emit(ind)
        if self.bkg_ind == ind:
            self.spectrum_changed.emit()

    def get_overlay_offset(self, ind):
        """
        Return the offset of the specified overlay
        :param ind: index of the overlay
        :return: overlay value
        """
        return self.overlays[ind].offset

    def set_overlay_as_bkg(self, ind):
        """
        Sets an overlay as background for the data spectrum, and unsets any previously used background
        :param ind: index of the overlay
        """
        if self.bkg_ind >= 0:
            self.unset_overlay_as_bkg()
        self.bkg_ind = ind
        self.spectrum.background_spectrum = self.overlays[ind]
        self.spectrum_changed.emit()
        self.overlay_set_as_bkg.emit(ind)

    def set_spectrum_as_bkg(self):
        """
        Adds the current spectrum as Overlay and sets it as background spectrum and unsets any previously used
        background.
        """
        self.add_spectrum_as_overlay()
        self.set_overlay_as_bkg(len(self.overlays) - 1)

    def unset_overlay_as_bkg(self):
        """
        Unsets the currently used background overlay.
        """
        previous_bkg_ind = self.bkg_ind
        self.bkg_ind = -1
        self.spectrum.unset_background_spectrum()
        self.spectrum_changed.emit()
        self.overlay_unset_as_bkg.emit(previous_bkg_ind)

    def overlay_is_bkg(self, ind):
        """
        :param ind: overlay ind
        """
        return ind == self.bkg_ind and self.bkg_ind != -1

    def set_auto_background_subtraction(self, parameters, roi=None):
        """
        Enables auto background extraction and removal from the data spectrum
        :param parameters: array of parameters with [window_width, iterations, polynomial_order]
        :param roi: array of size two with [xmin, xmax] specifying the range for which the background subtraction
        will be performed
        """
        self.spectrum.set_auto_background_subtraction(parameters, roi)
        self.spectrum_changed.emit()

    def unset_auto_background_subtraction(self):
        """
        Disables auto background extraction and removal.
        """
        self.spectrum.unset_auto_background_subtraction()
        self.spectrum_changed.emit()
Example #3
0
class ImgModel(Observable):
    """
    Main Image handling class. Supports several features:
        - loading image files in any format using fabio
        - iterating through files either by file number or time of creation
        - image transformations like rotating and flipping
        - setting a background image
        - setting an absorption correction (img_data is divided by this)
        - using supersampling (splitting each pixel into n**2 pixel with equal intensity)

    It inherits the Observable interface for implementing the observer pattern. To subscribe a function to changes in
    ImgData use:
        img_data = ImgData()
        img_data.subscribe(function)

    The function will be called every time the img_data has changed.
    """
    def __init__(self):
        """
        Defines all object variables and creates a dummy image.
        :return:
        """
        super(ImgModel, self).__init__()
        self.filename = ''
        self.img_transformations = []
        self.supersampling_factor = 1

        self.file_iteration_mode = 'number'
        self.file_name_iterator = FileNameIterator()

        self._img_data = None
        self._img_data_background_subtracted = None
        self._img_data_absorption_corrected = None
        self._img_data_background_subtracted_absorption_corrected = None

        self._img_data_supersampled = None
        self._img_data_supersampled_background_subtracted = None
        self._img_data_supersampled_absorption_corrected = None
        self._img_data_supersampled_background_subtracted_absorption_corrected = None

        self.background_filename = ''
        self._background_data = None
        self._background_scaling = 1
        self._background_offset = 0

        self._img_corrections = ImgCorrectionManager()

        self._create_dummy_img()

    def _create_dummy_img(self):
        self._img_data = np.zeros((2048, 2048))

    def load(self, filename):
        """
        Loads an image file in any format known by fabIO. Automatically performs all previous img transformations,
        performs supersampling and recalculates background subtracted and absorption corrected image data. Observers
        will be notified after the process.
        :param filename: path of the image file to be loaded
        """
        logger.info("Loading {0}.".format(filename))
        self.filename = filename
        try:
            self._img_data_fabio = fabio.open(filename)
            self._img_data = self._img_data_fabio.data[::-1]
        except AttributeError:
            self._img_data = np.array(Image.open(filename))[::-1]
        self.file_name_iterator.update_filename(filename)

        self._perform_img_transformations()
        self._calculate_img_data()
        self.notify()

    def save(self, filename):
        try:
            self._img_data_fabio.save(filename)
        except AttributeError:
            im_array = np.int32(np.copy(np.flipud(self._img_data)))
            im = Image.fromarray(im_array)
            im.save(filename)

    def load_background(self, filename):
        """
        Loads an image file as background in any format known by fabIO. Automatically performs all previous img
        transformations, supersampling and recalculates background subtracted and absorption corrected image data.
        Observers will be notified after the process.
        :param filename: path of the image file to be loaded
        """
        self.background_filename = filename
        try:
            self._background_data_fabio = fabio.open(filename)
            self._background_data = self._background_data_fabio.data[::-1].astype(float)
        except AttributeError:
            self._background_data = np.array(Image.open(filename))[::-1].astype(float)

        self._perform_background_transformations()
        self._calculate_img_data()
        self.notify()

    def _image_and_background_shape_equal(self):
        """
        Tests if the original image and original background image have the same shape
        :return: Boolean
        """
        if self._background_data is None:
            return True
        if self._background_data.shape == self._img_data.shape:
            return True
        return False

    def _reset_background(self):
        """
        Resets the background data to None
        """
        self.background_filename = None
        self._background_data = None
        self._background_data_fabio = None
        self._calculate_img_data()

    def reset_background(self):
        self._reset_background()
        self.notify()

    def has_background(self):
        return self._background_data is not None

    def set_background_scaling(self, value):
        self._background_scaling = value
        self._calculate_img_data()
        self.notify()

    def set_background_offset(self, value):
        self._background_offset = value
        self._calculate_img_data()
        self.notify()

    def load_next_file(self, step=1):
        next_file_name = self.file_name_iterator.get_next_filename(mode=self.file_iteration_mode, step=step)
        if next_file_name is not None:
            self.load(next_file_name)

    def load_previous_file(self, step=1):
        previous_file_name = self.file_name_iterator.get_previous_filename(mode=self.file_iteration_mode, step=step)
        if previous_file_name is not None:
            self.load(previous_file_name)

    def set_file_iteration_mode(self, mode):
        if mode == 'number':
            self.file_iteration_mode = 'number'
            self.file_name_iterator.create_timed_file_list = False
        elif mode == 'time':
            self.file_iteration_mode = 'time'
            self.file_name_iterator.create_timed_file_list = True
            self.file_name_iterator.update_filename(self.filename)

    def get_img_data(self):
        return self.img_data

    def get_img(self):
        if self._background_data is not None:
            return self._img_data_background_subtracted
        else:
            return self._img_data

    def _calculate_img_data(self):
        """
        Calculates compound img_data based on the state of the object. This function is used internally to not compute
        those img arrays every time somebody requests the image data by get_img_data() and img_data.
        """

        #check that all data has the same dimensions
        if self._background_data is not None:
            if self._img_data.shape != self._background_data.shape:
                self._background_data = None
        if self._img_corrections.has_items():
            self._img_corrections.set_shape(self._img_data.shape)

        #calculate the current _img_data
        if self._background_data is not None and not self._img_corrections.has_items():
            self._img_data_background_subtracted = self._img_data - (self._background_scaling *
                                                                     self._background_data +
                                                                     self._background_offset)
        elif self._background_data is None and self._img_corrections.has_items():
            self._img_data_absorption_corrected = self._img_data / self._img_corrections.get_data()

        elif self._background_data is not None and self._img_corrections.has_items():
            self._img_data_background_subtracted_absorption_corrected = (self._img_data - (
                self._background_scaling * self._background_data + self._background_offset)) / \
                                                                        self._img_corrections.get_data()

        # supersample the current image data
        if self.supersampling_factor > 1:
            if self._background_data is None and not self._img_corrections.has_items():
                self._img_data_supersampled = self.supersample_data(self._img_data, self.supersampling_factor)

            if self._background_data is not None and not self._img_corrections.has_items():
                self._img_data_supersampled_background_subtracted = \
                    self.supersample_data(self._img_data_background_subtracted, self.supersampling_factor)

            elif self._background_data is None and self._img_corrections.has_items():
                self._img_data_supersampled_absorption_corrected = \
                    self.supersample_data(self._img_data_absorption_corrected, self.supersampling_factor)

            elif self._background_data is not None and self._img_corrections.has_items():
                self._img_data_supersampled_background_subtracted_absorption_corrected = \
                    self.supersample_data(self._img_data_background_subtracted_absorption_corrected,
                                          self.supersampling_factor)


    @property
    def img_data(self):
        """
        :return:
            The image based on the current state of the ImgData object. If supersampling is set it will return a
            supersampled image array if background_data is set it will return a background_subtracted array and so on.
            It also works for combinations of all these options.
        """
        if self.supersampling_factor == 1:
            if self._background_data is None and not self._img_corrections.has_items():
                return self._img_data

            elif self._background_data is not None and not self._img_corrections.has_items():
                return self._img_data_background_subtracted

            elif self._background_data is None and self._img_corrections.has_items():
                return self._img_data_absorption_corrected

            elif self._background_data is not None and self._img_corrections.has_items():
                return self._img_data_background_subtracted_absorption_corrected

        else:
            if self._background_data is None and not self._img_corrections.has_items():
                return self._img_data_supersampled

            elif self._background_data is not None and not self._img_corrections.has_items():
                return self._img_data_supersampled_background_subtracted

            elif self._background_data is None and self._img_corrections.has_items():
                return self._img_data_supersampled_absorption_corrected

            elif self._background_data is not None and self._img_corrections.has_items():
                return self._img_data_supersampled_background_subtracted_absorption_corrected

    def rotate_img_p90(self):
        """
        Rotates the image by 90 degree and updates the background accordingly (does not effect absorption correction).
        The transformation is saved and applied to every new image and background image loaded.
        Notifies observers.
        """
        self._img_data = rotate_matrix_p90(self._img_data)

        if self._background_data is not None:
            self._background_data = rotate_matrix_p90(self._background_data)

        self.img_transformations.append(rotate_matrix_p90)

        self._calculate_img_data()
        self.notify()

    def rotate_img_m90(self):
        """
        Rotates the image by -90 degree and updates the background accordingly (does not effect absorption correction).
        The transformation is saved and applied to every new image and background image loaded.
        Notifies observers.
        """
        self._img_data = rotate_matrix_m90(self._img_data)
        if self._background_data is not None:
            self._background_data = rotate_matrix_m90(self._background_data)
        self.img_transformations.append(rotate_matrix_m90)

        self._calculate_img_data()
        self.notify()

    def flip_img_horizontally(self):
        """
        Flips image about a horizontal axis and updates the background accordingly (does not effect absorption
        correction). The transformation is saved and applied to every new image and background image loaded.
        Notifies observers.
        """
        self._img_data = np.fliplr(self._img_data)
        if self._background_data is not None:
            self._background_data = np.fliplr(self._background_data)
        self.img_transformations.append(np.fliplr)

        self._calculate_img_data()
        self.notify()

    def flip_img_vertically(self):
        """
        Flips image about a vertical axis and updates the background accordingly (does not effect absorption
        correction). The transformation is saved and applied to every new image and background image loaded.
        Notifies observers.
        """
        self._img_data = np.flipud(self._img_data)
        if self._background_data is not None:
            self._background_data = np.flipud(self._background_data)
        self.img_transformations.append(np.flipud)

        self._calculate_img_data()
        self.notify()

    def reset_img_transformations(self):
        """
        Reverts all image transformations and resets the transformation stack.
        Notifies observers.
        """
        for transformation in reversed(self.img_transformations):
            if transformation == rotate_matrix_p90:
                self._img_data = rotate_matrix_m90(self._img_data)
                if self._background_data is not None:
                    self._background_data = rotate_matrix_m90(self._background_data)
            elif transformation == rotate_matrix_m90:
                self._img_data = rotate_matrix_p90(self._img_data)
                if self._background_data is not None:
                    self._background_data = rotate_matrix_p90(self._background_data)
            else:
                self._img_data = transformation(self._img_data)
                if self._background_data is not None:
                    self._background_data = transformation(self._background_data)
        self.img_transformations = []
        self._calculate_img_data()
        self.notify()

    def _perform_img_transformations(self):
        """
        Performs all saved image transformation on original image.
        """
        for transformation in self.img_transformations:
            self._img_data = transformation(self._img_data)

    def _perform_background_transformations(self):
        """
        Performs all saved image transformation on background image.
        """
        if self._background_data is not None:
            for transformation in self.img_transformations:
                self._background_data = transformation(self._background_data)


    def set_supersampling(self, factor=None):
        """
        Stores the supersampling factor and calculates supersampled original and background image arrays.
        Updates all data calculations according to current ImgData object state.
        Does not notify Observers!
        :param factor: int - supersampling factor
        """
        self.supersampling_factor = factor
        self._calculate_img_data()

    def supersample_data(self, img_data, factor):
        """
        Creates a supersampled array from img_data.
        :param img_data: image array
        :param factor: int - supersampling factor
        :return:
        """
        if factor > 1:
            img_data_supersampled = np.zeros((img_data.shape[0] * factor,
                                              img_data.shape[1] * factor))
            for row in range(factor):
                for col in range(factor):
                    img_data_supersampled[row::factor, col::factor] = img_data

            return img_data_supersampled
        else:
            return img_data

    def add_img_correction(self, correction, name=None):
        self._img_corrections.add(correction, name)
        self._calculate_img_data()
        self.notify()

    def get_img_correction(self, name):
        return self._img_corrections.get_correction(name)

    def delete_img_correction(self, name=None):
        self._img_corrections.delete(name)
        self._calculate_img_data()
        self.notify()

    def has_corrections(self):
        """
        :return: Whether the ImgData object has active absorption corrections or not
        """
        return self._img_corrections.has_items()