Commit 081c272b authored by Nicola Vigano's avatar Nicola Vigano
Browse files

Alignment: added plotting and moved properties in base class


Signed-off-by: Nicola Vigano's avatarNicola VIGANÒ <nicola.vigano@esrf.fr>
parent dd0f5109
Pipeline #27921 passed with stages
in 3 minutes and 27 seconds
import numpy as np
import logging
from numpy.polynomial.polynomial import Polynomial
from numpy.polynomial.polynomial import Polynomial, polyval
from nabu.utils import previouspow2
from nabu.misc import fourier_filters
......@@ -15,8 +15,40 @@ except ImportError:
__have_scipy__ = False
try:
import matplotlib.pyplot as plt
__have_matplotlib__ = True
except ImportError:
logging.getLogger(__name__).warning("Matplotlib not available. Plotting disabled")
__have_matplotlib__ = False
class AlignmentBase(object):
def __init__(self, horz_fft_width=False, verbose=False):
"""
Alignment basic functions.
Parameters
----------
horz_fft_width: boolean, optional
If True, restrict the horizontal size to a power of 2:
>>> new_x_dim = 2 ** math.floor(math.log2(x_dim))
verbose: boolean, optional
When True it will produce verbose output, including plots.
"""
self._init_parameters(horz_fft_width, verbose)
def _init_parameters(self, horz_fft_width, verbose):
self.truncate_horz_pow2 = horz_fft_width
if verbose and not __have_matplotlib__:
logging.getLogger(__name__).warning("Matplotlib not available. Plotting disabled, despite being activated by user")
verbose = False
self.verbose = verbose
@staticmethod
def refine_max_position_2d(f_vals: np.ndarray, fy=None, fx=None):
"""Computes the sub-pixel max position of the given function sampling.
......@@ -83,7 +115,7 @@ class AlignmentBase(object):
if np.any(vertex_yx < vertex_min_yx) or np.any(vertex_yx > vertex_max_yx):
raise ValueError(
"Fitted (y: {}, x: {}) positions are outide the input margins y: [{}, {}], and x: [{}, {}]".format(
vertex_yx[0], vertex_yx[1], vertex_min_yx[0], vertex_max_yx[0], vertex_min_yx[1], vertex_max_yx[1]
vertex_yx[0], vertex_yx[1], vertex_min_yx[0], vertex_max_yx[0], vertex_min_yx[1], vertex_max_yx[1],
)
)
return vertex_yx
......@@ -150,7 +182,7 @@ class AlignmentBase(object):
)
else:
message = "Fitted positions outide the input margins [{}, {}]: %d below and %d above".format(
vertex_min_x, vertex_max_x, np.sum(1 - lower_bound_ok), np.sum(1 - upper_bound_ok)
vertex_min_x, vertex_max_x, np.sum(1 - lower_bound_ok), np.sum(1 - upper_bound_ok),
)
raise ValueError(message)
return vertex_x
......@@ -233,7 +265,7 @@ class AlignmentBase(object):
pix_max = np.argmax(cc, axis=axis)
# select a n neighborhood for the many 1D sub-pixel fittings (with wrapping)
p_ax_range = np.arange(- peak_radius, + peak_radius + 1)
p_ax_range = np.arange(-peak_radius, +peak_radius + 1)
p_ax = (pix_max[None, :] + p_ax_range[:, None]) % img_shape[axis]
p_ln = np.tile(np.arange(0, img_shape[axis])[None, :], [2 * peak_radius + 1, 1])
......@@ -249,13 +281,12 @@ class AlignmentBase(object):
return (f_vals, fc_ax)
@staticmethod
def _determine_roi(img_shape, roi_yxhw, do_truncate_horz_pow2):
def _determine_roi(self, img_shape, roi_yxhw):
if roi_yxhw is None:
# vertical window size is reduced to a power of 2 to accelerate fft
# same thing horizontal window - if requested. Default is not
roi_yxhw = previouspow2(img_shape)
if not do_truncate_horz_pow2:
if not self.truncate_horz_pow2:
roi_yxhw[1] = img_shape[1]
if len(roi_yxhw) == 2: # Convert centered 2-element roi into 4-element
......@@ -264,7 +295,9 @@ class AlignmentBase(object):
return roi_yxhw
@staticmethod
def _prepare_image(img, invalid_val=1e-5, roi_yxhw=None, median_filt_shape=None, low_pass=None, high_pass=None):
def _prepare_image(
img, invalid_val=1e-5, roi_yxhw=None, median_filt_shape=None, low_pass=None, high_pass=None,
):
"""
Prepare and returns a cropped and filtered image, or array of filtered images if the input is an array of images.
......@@ -367,23 +400,6 @@ class AlignmentBase(object):
class CenterOfRotation(AlignmentBase):
def __init__(self, horz_fft_width=False):
"""
Center of Rotation (CoR) computation object.
This class is used on radios.
Parameters
----------
horz_fft_width: boolean, optional
If True, restrict the horizontal size to a power of 2:
>>> new_x_dim = 2 ** math.floor(math.log2(x_dim))
"""
self._init_parameters(horz_fft_width)
def _init_parameters(self, horz_fft_width):
self.truncate_horz_pow2 = horz_fft_width
@staticmethod
def _check_img_sizes(img_1: np.ndarray, img_2: np.ndarray):
shape_1 = np.squeeze(img_1).shape
......@@ -492,7 +508,7 @@ class CenterOfRotation(AlignmentBase):
peak_fit_radius = 1
img_shape = img_2.shape
roi_yxhw = self._determine_roi(img_shape, roi_yxhw, self.truncate_horz_pow2)
roi_yxhw = self._determine_roi(img_shape, roi_yxhw)
img_1 = self._prepare_image(img_1, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape)
img_2 = self._prepare_image(img_2, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape)
......@@ -608,7 +624,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
num_imgs = img_stack.shape[0]
img_shape = img_stack.shape[-2:]
roi_yxhw = self._determine_roi(img_shape, roi_yxhw, False)
roi_yxhw = self._determine_roi(img_shape, roi_yxhw)
img_stack = self._prepare_image(img_stack, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape)
......@@ -641,6 +657,16 @@ class DetectorTranslationAlongBeam(AlignmentBase):
coeffs_v = Polynomial.fit(img_pos_increments, shifts_vh[:, 0], deg=1).convert().coef
coeffs_h = Polynomial.fit(img_pos_increments, shifts_vh[:, 1], deg=1).convert().coef
if self.verbose:
f, axs = plt.subplots(1, 2)
axs[0].scatter(img_pos_increments, shifts_vh[:, 0])
axs[0].plot(img_pos_increments, polyval(img_pos_increments, coeffs_v))
axs[0].set_title("Vertical shifts")
axs[1].scatter(img_pos_increments, shifts_vh[:, 1])
axs[1].plot(img_pos_increments, polyval(img_pos_increments, coeffs_h))
axs[1].set_title("Horizontal shifts")
plt.show(block=False)
if return_shifts:
return coeffs_v[1], coeffs_h[1], shifts_vh
else:
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment