Пример #1
0
def get_mask_from_roi(signal, roi, axes=None, gaussian=True):
    if axes is None and signal in roi.signal_map:
        axes = roi.signal_map[signal][1]
    else:
        axes = roi._parse_axes(axes, signal.axes_manager)

    # Needs to add support for other type of ROI
    if hasattr(roi, 'cx'):
        # CircleROI
        radius = roi.r
        cx = roi.cx
        cy = roi.cy
        r = np.linalg.norm([cx, cy]) * 0.8
        # The factor of 3 come from an estimate of how far the tail of the
        # Gaussian goes; to avoid getting the zero-frequency component in
        # the mask, we clip its radius_slice value
        radius_slice = np.clip(radius * 3, a_min=radius, a_max=r)
        ranges = [[cx - radius_slice, cx + radius_slice],
                  [cy - radius_slice, cy + radius_slice]]
    else:
        ranges = roi._get_ranges()

    if hasattr(roi, 'cx'):
        # The 'else' part is missing
        slices = roi._make_slices(axes, axes, ranges=ranges)

        if not gaussian:
            # in case of Bragg Filtering
            radius_slice = radius

        # Calculate a disk mask
        sig_axes = signal.axes_manager.signal_axes
        ir = [slices[sig_axes.index(axes[0])], slices[sig_axes.index(axes[1])]]
        vx = axes[0].axis[ir[0]] - cx
        vy = axes[1].axis[ir[1]] - cy
        gx, gy = np.meshgrid(vx, vy)
        gr = gx**2 + gy**2
        disk_mask = gr > radius_slice**2

        if gaussian:
            import hyperspy.api as hs
            mask = hs.signals.Signal2D(np.zeros(signal.data.shape))
            x = np.linspace(ranges[0][0], ranges[0][1], disk_mask.shape[1])
            y = np.linspace(ranges[1][0], ranges[1][1], disk_mask.shape[0])
            xx, yy = np.meshgrid(x, y)

            gaussian2d = hs.model.components2D.Gaussian2D(sigma_x=radius,
                                                          sigma_y=radius,
                                                          centre_x=cx,
                                                          centre_y=cy,
                                                          A=2 * np.pi *
                                                          radius**2)
            mask_circle = gaussian2d.function(xx, yy) * ~disk_mask
        else:
            mask = BaseSignal(np.full(signal.data.shape, True, dtype=bool))
            mask.axes_manager.set_signal_dimension(
                signal.axes_manager.signal_dimension)

            mask_circle = disk_mask

        mask.isig[slices] = mask_circle

        # If signal.data is cupy array, transfer the array to the GPU
        xp = get_array_module(signal.data)
        mask.data = xp.asarray(mask.data)

    return mask