Skip to content

Commit

Permalink
more numba
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexlib
alexlib committed Dec 29, 2023
1 parent ed9b0e9 commit 3aacf84
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 48 deletions.
67 changes: 39 additions & 28 deletions openptv_python/trafo.py
View file
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -81,6 +88,7 @@ def metric_to_pixel(
parameters.pix_y
)


@njit
def fast_metric_to_pixel(
x_metric,
Expand All @@ -89,15 +97,16 @@ 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)

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:
Expand All @@ -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,
Expand All @@ -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

Expand All @@ -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,
Expand All @@ -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)

Expand All @@ -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(
Expand Down
63 changes: 43 additions & 20 deletions tests/test_trafo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Expand All @@ -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):
Expand Down Expand Up @@ -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]])
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 3aacf84

Please sign in to comment.