From 84f9c60bf4e1baa4a8340936249e0b08d57b789f Mon Sep 17 00:00:00 2001
From: sccolbert <sccolbert@gmail.com>
Date: Wed, 11 Nov 2009 14:49:23 +0100
Subject: [PATCH] added cvUndistort2 and its tests

---
 scikits/image/opencv/opencv_cv.pyx           | 63 ++++++++++++++++++++
 scikits/image/opencv/tests/test_opencv_cv.py | 16 +++++
 2 files changed, 79 insertions(+)

diff --git a/scikits/image/opencv/opencv_cv.pyx b/scikits/image/opencv/opencv_cv.pyx
index d02f1507..87a96f59 100644
--- a/scikits/image/opencv/opencv_cv.pyx
+++ b/scikits/image/opencv/opencv_cv.pyx
@@ -266,6 +266,11 @@ cdef cvCalibrateCamera2Ptr c_cvCalibrateCamera2
 c_cvCalibrateCamera2 = (<cvCalibrateCamera2Ptr*>
                         <size_t>ctypes.addressof(cv.cvCalibrateCamera2))[0]
 
+# cvUndistort2
+ctypedef void (*cvUndistort2Ptr)(IplImage*, IplImage*, CvMat*, CvMat*)
+cdef cvUndistort2Ptr c_cvUndistort2
+c_cvUndistort2 = (<cvUndistort2Ptr*><size_t>ctypes.addressof(cv.cvUndistort2))[0]
+
 # cvFindChessboardCorners
 ctypedef void (*cvFindChessboardCornersPtr)(IplImage*, CvSize, CvPoint2D32f*,
                                             int*, int)
@@ -2155,6 +2160,64 @@ def cvCalibrateCamera2(np.ndarray object_points, np.ndarray image_points,
     return intrinsics, distortion
 
 
+#-------------
+# cvUndistort2
+#-------------
+
+@cvdoc(package='cv', group='calibration', doc=\
+'''cvUndistort2(src, intrinsics, distortions)
+
+Undistorts an image given the camera intrinsics matrix and distortions vector.
+These values can be calculated using cvCalibrateCamera2.
+
+Parameters
+----------
+src : ndarray
+    The image to undistort
+intrinsics : ndarray, 3x3, dtype=float64
+    The camera intrinsics matrix.
+distortions : ndarray, 5-vector, dtype=float64
+    The camera distortion coefficients.
+
+Returns
+-------
+out : ndarray
+    The undistorted image the same size and dtype
+    as the source image.''')
+def cvUndistort2(src, intrinsics, distortions):
+    validate_array(src)
+    assert_dtype(intrinsics, [FLOAT64])
+    assert_dtype(distortions, [FLOAT64])
+    assert_ndims(intrinsics, [2])
+    assert_ndims(distortions, [1])
+
+    if intrinsics.shape[0] != 3 or intrinsics.shape[1] != 3:
+        raise ValueError('intrinsics must be 3x3')
+    if distortions.shape[0] != 5:
+        raise ValueError('distortions must be a 5-vector')
+
+    cdef np.ndarray out = new_array_like(src)
+
+    cdef IplImage srcimg
+    cdef IplImage outimg
+    cdef IplImage intrimg
+    cdef IplImage distimg
+
+    populate_iplimage(src, &srcimg)
+    populate_iplimage(out, &outimg)
+    populate_iplimage(intrinsics, &intrimg)
+    populate_iplimage(distortions, &distimg)
+
+    cdef CvMat* cvintr = cvmat_ptr_from_iplimage(&intrimg)
+    cdef CvMat* cvdist = cvmat_ptr_from_iplimage(&distimg)
+
+    c_cvUndistort2(&srcimg, &outimg, cvintr, cvdist)
+
+    PyMem_Free(cvintr)
+    PyMem_Free(cvdist)
+
+    return out
+
 #------------------------
 # cvFindChessboardCorners
 #------------------------
diff --git a/scikits/image/opencv/tests/test_opencv_cv.py b/scikits/image/opencv/tests/test_opencv_cv.py
index d647243c..cc73258d 100644
--- a/scikits/image/opencv/tests/test_opencv_cv.py
+++ b/scikits/image/opencv/tests/test_opencv_cv.py
@@ -287,6 +287,22 @@ class TestCalibrateCamera2(object):
         )
 
 
+class TestUndistort2(OpenCVTest):
+    @opencv_skip
+    def test_cvUndistort2(self):
+        intrinsics = np.array([[1, 0, 0],
+                               [0, 1, 0],
+                               [0, 0, 1]], dtype='float64')
+        distortions = np.array([0., 0., 0., 0., 0.], dtype='float64')
+
+        undist = cvUndistort2(self.lena_RGB_U8, intrinsics, distortions)
+        undistg = cvUndistort2(self.lena_GRAY_U8, intrinsics, distortions)
+
+        assert_array_almost_equal(undist, self.lena_RGB_U8)
+        assert_array_almost_equal(undistg, self.lena_GRAY_U8)
+
+
+
 
 if __name__ == '__main__':
     run_module_suite()