ReconstructHelper

cdef class ReconstructHelper:
    cdef readonly double fc1, fc2, cc1, cc2
    cdef readonly double k1, k2, p1, p2
    cdef readonly double alpha_c

    def __init__(self, fc1, fc2, cc1, cc2, k1, k2, p1, p2, alpha_c=0 ):
        """create instance of ReconstructHelper

        ReconstructHelper(fc1, fc2, cc1, cc2, k1, k2, p1, p2 )
        where:
        fc - focal length
        cc - camera center
        k - radial distortion parameters (non-linear)
        p - tangential distortion parameters (non-linear)
        alpha_c - skew between X and Y pixel axes
        """
        self.fc1 = fc1
        self.fc2 = fc2
        self.cc1 = cc1
        self.cc2 = cc2
        self.k1 = k1
        self.k2 = k2
        self.p1 = p1
        self.p2 = p2
        self.alpha_c = alpha_c

    def __reduce__(self):
        """this allows ReconstructHelper to be pickled"""
        args = (self.fc1, self.fc2, self.cc1, self.cc2,
                self.k1, self.k2, self.p1, self.p2, self.alpha_c)
        return (make_ReconstructHelper, args)

    def add_element(self, parent):
        """add self as XML element to parent"""
        assert ET.iselement(parent)
        elem = ET.SubElement(parent,"non_linear_parameters")

        for name in ['fc1','fc2','cc1','cc2','k1','k2','p1','p2','alpha_c']:
            e = ET.SubElement(elem, name)
            e.text = repr( getattr(self,name) )

    def as_obj_for_json(self):
        result = {}
        for name in ['fc1','fc2','cc1','cc2','k1','k2','p1','p2','alpha_c']:
            result[name] = getattr(self,name)
        return result

    def save_to_rad_file( self, fname, comments=None ):
        rad_fd = open(fname,'w')
        K = self.get_K()
        nlparams = self.get_nlparams()
        k1, k2, p1, p2 = nlparams
        rad_fd.write('K11 = %s;\n'%repr(K[0,0]))
        rad_fd.write('K12 = %s;\n'%repr(K[0,1]))
        rad_fd.write('K13 = %s;\n'%repr(K[0,2]))
        rad_fd.write('K21 = %s;\n'%repr(K[1,0]))
        rad_fd.write('K22 = %s;\n'%repr(K[1,1]))
        rad_fd.write('K23 = %s;\n'%repr(K[1,2]))
        rad_fd.write('K31 = %s;\n'%repr(K[2,0]))
        rad_fd.write('K32 = %s;\n'%repr(K[2,1]))
        rad_fd.write('K33 = %s;\n'%repr(K[2,2]))
        rad_fd.write('\n')
        rad_fd.write('kc1 = %s;\n'%repr(k1))
        rad_fd.write('kc2 = %s;\n'%repr(k2))
        rad_fd.write('kc3 = %s;\n'%repr(p1))
        rad_fd.write('kc4 = %s;\n'%repr(p2))
        rad_fd.write('\n')
        if comments is not None:
            comments = str(comments)
        rad_fd.write("comments = '%s';\n"%comments)
        rad_fd.write('\n')
        rad_fd.close()

    def __richcmp__(self,other,op):
        # cmp op
        #  < 0
        # <= 1
        # == 2
        # != 3
        #  > 4
        # >= 5

        if isinstance(other, ReconstructHelper):
            isequal = (numpy.allclose(self.get_K(),other.get_K()) and
                      numpy.allclose(self.get_nlparams(),other.get_nlparams()) )
        else:
            isequal = False

        if op in [1,2,5]:
            result = isequal
        elif op == 3:
            result = not isequal
        else:
            result = False
        return result

    def get_K(self):
        # See
        # http://www.vision.caltech.edu/bouguetj/calib_doc/htmls/parameters.html
        K = numpy.array((( self.fc1, self.alpha_c*self.fc1, self.cc1),
                         ( 0,        self.fc2,              self.cc2),
                         ( 0,        0,                     1       )))
        return K

    def get_nlparams(self):
        return (self.k1, self.k2, self.p1, self.p2)

    def undistort(self, double x_kk, double y_kk):
        """undistort 2D coordinate pair

        Iteratively performs an undistortion using camera intrinsic
        parameters.

        Implementation translated from CalTechCal.

        See also the OpenCV reference manual, which has the equation
        used.
        """

        cdef double xl, yl

        cdef double xd, yd, x, y
        cdef double r_2, k_radial, delta_x, delta_y
        cdef int i

        # undoradial.m / CalTechCal/normalize.m

        xd = ( x_kk - self.cc1 ) / self.fc1
        yd = ( y_kk - self.cc2 ) / self.fc2

        xd = xd - self.alpha_c * yd

        # comp_distortion_oulu.m

        # initial guess
        x = xd
        y = yd

        for i from 0<=i<20:
            r_2 = x*x + y*y
            k_radial = 1.0 + (self.k1) * r_2 + (self.k2) * r_2*r_2
            delta_x = 2.0 * (self.p1)*x*y + (self.p2)*(r_2 + 2.0*x*x)
            delta_y = (self.p1) * (r_2 + 2.0*y*y)+2.0*(self.p2)*x*y
            x = (xd-delta_x)/k_radial
            y = (yd-delta_y)/k_radial

        # undoradial.m

        xl = (self.fc1)*x + (self.fc1*self.alpha_c)*y + (self.cc1)
        yl = (self.fc2)*y + (self.cc2)
        return (xl, yl)

    def distort(self, double xl, double yl):
        """distort 2D coordinate pair"""

        cdef double x, y, r_2, term1, xd, yd

        x = ( xl - self.cc1 ) / self.fc1
        y = ( yl - self.cc2 ) / self.fc2

        r_2 = x*x + y*y
        r_4 = r_2**2
        term1 = self.k1*r_2 + self.k2*r_4

        # OpenCV manual (chaper 6, "3D reconstruction", "camera
        # calibration" section) seems to disagree with
        # http://www.vision.caltech.edu/bouguetj/calib_doc/htmls/parameters.html

        # Furthermore, implementations in cvundistort.cpp seem still
        # unrelated.  Finally, project_points2.m in Bouguet's code is
        # consistent with his webpage and this below.

        xd = x + x*term1 + (2*self.p1*x*y + self.p2*(r_2+2*x**2))
        yd = y + y*term1 + (self.p1*(r_2+2*y**2) + 2*self.p2*x*y)

        xd = (self.fc1)*xd + (self.cc1)
        yd = (self.fc2)*yd + (self.cc2)

        return (xd, yd)

    def undistort_image(self, numpy_image ):
        assert len(numpy_image.shape)==2
        assert numpy_image.dtype == numpy.uint8

        f = self.fc1, self.fc2 # focal length
        c = self.cc1, self.cc2 # center
        k = self.k1, self.k2, self.p1, self.p2, 0  # NL terms: r^2, r^4, tan1, tan2, r^6

        imnew = flydra.undistort.rect(numpy_image, f=f, c=c, k=k) # perform the undistortion
        imnew = imnew.astype(numpy.uint8)
        return imnew