From 3aacf84df5209c65b945707000ab9f7e7a0210d1 Mon Sep 17 00:00:00 2001 From: Alex Liberzon Date: Fri, 29 Dec 2023 21:06:20 +0200 Subject: [PATCH] more numba --- openptv_python/trafo.py | 67 ++++++++++++++++++++++++----------------- tests/test_trafo.py | 63 ++++++++++++++++++++++++++------------ 2 files changed, 82 insertions(+), 48 deletions(-) diff --git a/openptv_python/trafo.py b/openptv_python/trafo.py index 981f37a..42b8dc3 100644 --- a/openptv_python/trafo.py +++ b/openptv_python/trafo.py @@ -1,9 +1,9 @@ """Module for coordinate transformations.""" -from math import cos, sin, sqrt from typing import Tuple import numpy as np -from numba import njit +from numba import float64, int32, njit +from numpy import cos, sin, sqrt from .calibration import Calibration, ap_52 from .parameters import ControlPar @@ -38,22 +38,29 @@ def fast_pixel_to_metric(x_pixel, y_pixel, imx, imy, pix_x, pix_y) -> Tuple[floa return (x_metric, y_metric) - -def arr_pixel_to_metric(pixel: np.ndarray, parameters: ControlPar) -> np.ndarray: +@njit(float64[:,:](int32[:,:],int32,int32,float64,float64)) +def arr_pixel_to_metric(pixel: np.ndarray, + imx: int, + imy: int, + pix_x: float, + pix_y: float) -> np.ndarray: """Convert pixel coordinates to metric coordinates. Arguments: --------- + imx (float): image width in pixels. + imy (float): image height in pixels. + pix_x (float): pixel size in x-direction. + pix_y (float): pixel size in y-direction. + + Returns + ------- metric (np.ndarray): output metric coordinates. - pixel (np.ndarray): input pixel coordinates. - parameters (ControlPar): control structure holding image and pixel sizes. + """ - pixel = np.atleast_2d(np.array(pixel)) - metric = np.empty_like(pixel) - metric[:, 0] = (pixel[:, 0] - float(parameters.imx) / - 2.0) * parameters.pix_x - metric[:, 1] = (float(parameters.imy) / 2.0 - - pixel[:, 1]) * parameters.pix_y + metric = np.empty_like(pixel, dtype=np.float64) + metric[:, 0] = (pixel[:, 0] - imx / 2.0) * pix_x + metric[:, 1] = (imy / 2.0 - pixel[:, 1]) * pix_y return metric @@ -81,6 +88,7 @@ def metric_to_pixel( parameters.pix_y ) + @njit def fast_metric_to_pixel( x_metric, @@ -89,7 +97,7 @@ def fast_metric_to_pixel( imy, pix_x, pix_y - ) -> Tuple[float, float]: +) -> Tuple[float, float]: """Convert metric coordinates to pixel coordinates.""" x_pixel = (x_metric / pix_x) + (float(imx) / 2.0) y_pixel = (float(imy) / 2.0) - (y_metric / pix_y) @@ -97,7 +105,8 @@ def fast_metric_to_pixel( return x_pixel, y_pixel -def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarray: +def arr_metric_to_pixel(metric: np.ndarray, + parameters: ControlPar) -> np.ndarray: """Convert an array of metric coordinates to pixel coordinates. Arguments: @@ -109,7 +118,7 @@ def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarra ------- pixel (np.ndarray): output array of pixel coordinates. """ - metric = np.atleast_2d(np.array(metric)) + metric = np.atleast_2d(metric) return fast_arr_metric_to_pixel( metric, @@ -119,18 +128,19 @@ def arr_metric_to_pixel(metric: np.ndarray, parameters: ControlPar) -> np.ndarra parameters.pix_y ) -@njit + +@njit(float64[:,:](float64[:,:],int32,int32,float64,float64)) def fast_arr_metric_to_pixel( - metric, - imx, - imy, - pix_x, - pix_y - ) -> np.ndarray: + metric: np.ndarray, + imx: int, + imy: int, + pix_x: float, + pix_y: float +) -> np.ndarray: """Convert an array of metric coordinates to pixel coordinates.""" pixel = np.zeros_like(metric) - pixel[:, 0] = (metric[:, 0] / pix_x) + (float(imx) / 2.0) - pixel[:, 1] = (float(imy) / 2.0) - (metric[:, 1] / pix_y) + pixel[:, 0] = (metric[:, 0] / pix_x) + (imx / 2.0) + pixel[:, 1] = (imy / 2.0) - (metric[:, 1] / pix_y) return pixel @@ -143,13 +153,14 @@ def distort_brown_affine(x: float, if x == 0 and y == 0: return 0, 0 - return fast_distort_brown_affine(x, y, ap.k1, ap.k2, ap.k3, + tmp = fast_distort_brown_affine(x, y, ap.k1, ap.k2, ap.k3, ap.p1, ap.p2, ap.she, ap.scx) + return tmp[0], tmp[1] # print(f"x {x}, y {y}") -@njit +@njit(float64[:](float64,float64,float64,float64,float64,float64,float64,float64,float64)) def fast_distort_brown_affine( x: float, y: float, @@ -160,7 +171,7 @@ def fast_distort_brown_affine( p2: float, she: float, scx: float, -) -> Tuple[float, float]: +) -> np.ndarray: """Distort a point using the Brown affine model.""" r = sqrt(x**2 + y**2) @@ -183,7 +194,7 @@ def fast_distort_brown_affine( # print(f"x1 {x1}, y1 {y1}") - return x1, y1 + return np.array([x1, y1]) def correct_brown_affine( diff --git a/tests/test_trafo.py b/tests/test_trafo.py index a10cb8f..956f16d 100644 --- a/tests/test_trafo.py +++ b/tests/test_trafo.py @@ -224,7 +224,13 @@ def test_pixel_to_metric(self): y_metric_expected = (float(parameters.imy) / 2.0 - y_pixel) * parameters.pix_y # call the function to get actual output metric coordinates - metric = arr_pixel_to_metric(np.array([x_pixel, y_pixel]), parameters) + metric = arr_pixel_to_metric( + np.atleast_2d(np.array([x_pixel, y_pixel], dtype=np.int32)), + parameters.imx, + parameters.imy, + parameters.pix_x, + parameters.pix_y + ) # check if the actual output matches the expected output assert metric[:, 0] == x_metric_expected @@ -251,23 +257,33 @@ def setUp(self): def test_transforms_regress(self): """Transformed values are as before.""" - input_pos = np.full((3, 2), 100.0) - - correct_output_pixel_to_metric = [ + correct_output_pixel_to_metric = np.array( + [ [-8181.0, 6657.92], [-8181.0, 6657.92], [-8181.0, 6657.92], - ] - correct_output_metric_to_pixel = [ + ]) + correct_output_metric_to_pixel = np.array([ [646.60066007, 505.81188119], [646.60066007, 505.81188119], [646.60066007, 505.81188119], - ] + ]) + + input_pos = np.full((3, 2), 100, dtype=np.int32) + # Test when passing a list - output = arr_pixel_to_metric(input_pos, self.control) + output = arr_pixel_to_metric( + input_pos, + self.control.imx, + self.control.imy, + self.control.pix_x, + self.control.pix_y + ) np.testing.assert_array_almost_equal(output, correct_output_pixel_to_metric) + + input_pos = np.full((3, 2), 100, dtype=np.float64) output = arr_metric_to_pixel(input_pos, self.control) np.testing.assert_array_almost_equal(output, correct_output_metric_to_pixel) @@ -278,21 +294,28 @@ def test_transforms(self): cpar.set_pixel_size((0.1, 0.1)) metric_pos = np.array([[1.0, 1.0], [-10.0, 15.0], [20.0, -30.0]]) - pixel_pos = np.array([[650.0, 490.0], [540.0, 350.0], [840.0, 800.0]]) + pixel_pos = np.array([[650, 490], [540, 350], [840, 800]],dtype=np.int32) np.testing.assert_array_almost_equal( - pixel_pos, arr_metric_to_pixel(metric_pos, cpar) + pixel_pos, + arr_metric_to_pixel( + metric_pos, + cpar) ) np.testing.assert_array_almost_equal( - metric_pos, arr_pixel_to_metric(pixel_pos, cpar) + metric_pos, arr_pixel_to_metric(pixel_pos, + cpar.imx, + cpar.imy, + cpar.pix_x, + cpar.pix_y) ) def test_brown_affine_regress(self): """Test that the brown affine transform gives the same results as before.""" input_vec = np.full((3, 2), 100.0) output = np.zeros((3, 2)) - correct_output_corr = [[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]] - correct_output_dist = [[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]] + correct_output_corr = np.array([[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]]) + correct_output_dist = np.array([[100.0, 100.0], [100.0, 100.0], [100.0, 100.0]]) # Test when passing an array for output for i in range(3): @@ -322,11 +345,11 @@ def test_brown_affine(self): cal = Calibration() cal.set_pos([0.0, 0.0, 40.0]) cal.set_angles([0.0, 0.0, 0.0]) - cal.set_primary_point([0.0, 0.0, 10.0]) + cal.set_primary_point(np.array([0.0, 0.0, 10.0])) cal.set_glass_vec(np.r_[0.0, 0.0, 20.0]) - cal.set_radial_distortion([0,0,0]) - cal.set_decentering([0,0]) - cal.set_affine_trans([1, 0]) + cal.set_radial_distortion(np.zeros(3,)) + cal.set_decentering(np.zeros(2,)) + cal.set_affine_trans(np.array([1,0])) # reference metric positions: ref_pos = np.array([[0.1, 0.1], [1.0, -1.0], [-10.0, 10.0]]) @@ -357,9 +380,9 @@ def test_full_correction(self): cal.set_angles(np.r_[0.0, 0.0, 0.0]) cal.set_primary_point(np.r_[0.0, 0.0, 10.0]) cal.set_glass_vec(np.r_[0.0, 0.0, 20.0]) - cal.set_radial_distortion([0,0,0]) - cal.set_decentering([0,0]) - cal.set_affine_trans([1, 0]) + cal.set_radial_distortion(np.zeros(3,)) + cal.set_decentering(np.zeros(2,)) + cal.set_affine_trans(np.array([1, 0])) # reference metric positions: # Note the last value is different than in test_brown_affine() because