import os, sys
import numpy as np
from numpy import cos, sin, arctan, sqrt
import cv2
import copy
import time

def Homography(image, img_points, world_width, world_height,
               interpolation=cv2.INTER_CUBIC, ratio_width=1.0, ratio_height=1.0):
    _points = np.array(img_points).reshape(-1, 2).astype(np.float32)

    expand_x = int(0.5 * world_width * (ratio_width - 1))
    expand_y = int(0.5 * world_height * (ratio_height - 1))

    pt_lefttop = [expand_x, expand_y]
    pt_righttop = [expand_x + world_width, expand_y]
    pt_leftbottom = [expand_x + world_width, expand_y + world_height]
    pt_rightbottom = [expand_x, expand_y + world_height]

    pts_std = np.float32([pt_lefttop, pt_righttop,
                          pt_leftbottom, pt_rightbottom])

    img_crop_width = int(world_width * ratio_width)
    img_crop_height = int(world_height * ratio_height)

    M = cv2.getPerspectiveTransform(_points, pts_std)

    dst_img = cv2.warpPerspective(
        image,
        M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_CONSTANT,  # BORDER_CONSTANT BORDER_REPLICATE
        flags=interpolation)

    return dst_img

class CurveTextRectifier:
    """
    spatial transformer via monocular vision
    """
    def __init__(self):
        self.get_virtual_camera_parameter()


    def get_virtual_camera_parameter(self):
        vcam_thz = 0
        vcam_thx1 = 180
        vcam_thy = 180
        vcam_thx2 = 0

        vcam_x = 0
        vcam_y = 0
        vcam_z = 100

        radian = np.pi / 180

        angle_z = radian * vcam_thz
        angle_x1 = radian * vcam_thx1
        angle_y = radian * vcam_thy
        angle_x2 = radian * vcam_thx2

        optic_x = vcam_x
        optic_y = vcam_y
        optic_z = vcam_z

        fu = 100
        fv = 100

        matT = np.zeros((4, 4))
        matT[0, 0] = cos(angle_z) * cos(angle_y) - sin(angle_z) * sin(angle_x1) * sin(angle_y)
        matT[0, 1] = cos(angle_z) * sin(angle_y) * sin(angle_x2) - sin(angle_z) * (
                    cos(angle_x1) * cos(angle_x2) - sin(angle_x1) * cos(angle_y) * sin(angle_x2))
        matT[0, 2] = cos(angle_z) * sin(angle_y) * cos(angle_x2) + sin(angle_z) * (
                    cos(angle_x1) * sin(angle_x2) + sin(angle_x1) * cos(angle_y) * cos(angle_x2))
        matT[0, 3] = optic_x
        matT[1, 0] = sin(angle_z) * cos(angle_y) + cos(angle_z) * sin(angle_x1) * sin(angle_y)
        matT[1, 1] = sin(angle_z) * sin(angle_y) * sin(angle_x2) + cos(angle_z) * (
                    cos(angle_x1) * cos(angle_x2) - sin(angle_x1) * cos(angle_y) * sin(angle_x2))
        matT[1, 2] = sin(angle_z) * sin(angle_y) * cos(angle_x2) - cos(angle_z) * (
                    cos(angle_x1) * sin(angle_x2) + sin(angle_x1) * cos(angle_y) * cos(angle_x2))
        matT[1, 3] = optic_y
        matT[2, 0] = -cos(angle_x1) * sin(angle_y)
        matT[2, 1] = cos(angle_x1) * cos(angle_y) * sin(angle_x2) + sin(angle_x1) * cos(angle_x2)
        matT[2, 2] = cos(angle_x1) * cos(angle_y) * cos(angle_x2) - sin(angle_x1) * sin(angle_x2)
        matT[2, 3] = optic_z
        matT[3, 0] = 0
        matT[3, 1] = 0
        matT[3, 2] = 0
        matT[3, 3] = 1

        matS = np.zeros((4, 4))
        matS[2, 3] = 0.5
        matS[3, 2] = 0.5

        self.ifu = 1 / fu
        self.ifv = 1 / fv

        self.matT = matT
        self.matS = matS
        self.K = np.dot(matT.T, matS)
        self.K = np.dot(self.K, matT)


    def vertical_text_process(self, points, org_size):
        """
        change sequence amd process
        :param points:
        :param org_size:
        :return:
        """
        org_w, org_h = org_size
        _points = np.array(points).reshape(-1).tolist()
        _points = np.array(_points[2:] + _points[:2]).reshape(-1, 2)

        # convert to horizontal points
        adjusted_points = np.zeros(_points.shape, dtype=np.float32)
        adjusted_points[:, 0] = _points[:, 1]
        adjusted_points[:, 1] = org_h - _points[:, 0] - 1

        _image_coord, _world_coord, _new_image_size = self.horizontal_text_process(adjusted_points)

        # # convert to vertical points back
        image_coord = _points.reshape(1, -1, 2)
        world_coord = np.zeros(_world_coord.shape, dtype=np.float32)
        world_coord[:, :, 0] = 0 - _world_coord[:, :, 1]
        world_coord[:, :, 1] = _world_coord[:, :, 0]
        world_coord[:, :, 2] = _world_coord[:, :, 2]
        new_image_size = (_new_image_size[1], _new_image_size[0])

        return image_coord, world_coord, new_image_size


    def horizontal_text_process(self, points):
        """
        get image coordinate and world coordinate
        :param points:
        :return:
        """
        poly = np.array(points).reshape(-1)

        dx_list = []
        dy_list = []
        for i in range(1, len(poly) // 2):
            xdx = poly[i * 2] - poly[(i - 1) * 2]
            xdy = poly[i * 2 + 1] - poly[(i - 1) * 2 + 1]
            d = sqrt(xdx ** 2 + xdy ** 2)
            dx_list.append(d)

        for i in range(0, len(poly) // 4):
            ydx = poly[i * 2] - poly[len(poly) - 1 - (i * 2 + 1)]
            ydy = poly[i * 2 + 1] - poly[len(poly) - 1 - (i * 2)]
            d = sqrt(ydx ** 2 + ydy ** 2)
            dy_list.append(d)

        dx_list = [(dx_list[i] + dx_list[len(dx_list) - 1 - i]) / 2 for i in range(len(dx_list) // 2)]

        height = np.around(np.mean(dy_list))

        rect_coord = [0, 0]
        for i in range(0, len(poly) // 4 - 1):
            x = rect_coord[-2]
            x += dx_list[i]
            y = 0
            rect_coord.append(x)
            rect_coord.append(y)

        rect_coord_half = copy.deepcopy(rect_coord)
        for i in range(0, len(poly) // 4):
            x = rect_coord_half[len(rect_coord_half) - 2 * i - 2]
            y = height
            rect_coord.append(x)
            rect_coord.append(y)

        np_rect_coord = np.array(rect_coord).reshape(-1, 2)
        x_min = np.min(np_rect_coord[:, 0])
        y_min = np.min(np_rect_coord[:, 1])
        x_max = np.max(np_rect_coord[:, 0])
        y_max = np.max(np_rect_coord[:, 1])
        new_image_size = (int(x_max - x_min + 0.5), int(y_max - y_min + 0.5))
        x_mean = (x_max - x_min) / 2
        y_mean = (y_max - y_min) / 2
        np_rect_coord[:, 0] -= x_mean
        np_rect_coord[:, 1] -= y_mean
        rect_coord = np_rect_coord.reshape(-1).tolist()

        rect_coord = np.array(rect_coord).reshape(-1, 2)
        world_coord = np.ones((len(rect_coord), 3)) * 0

        world_coord[:, :2] = rect_coord

        image_coord = np.array(poly).reshape(1, -1, 2)
        world_coord = world_coord.reshape(1, -1, 3)

        return image_coord, world_coord, new_image_size


    def horizontal_text_estimate(self, points):
        """
        horizontal or vertical text
        :param points:
        :return:
        """
        pts = np.array(points).reshape(-1, 2)
        x_min = int(np.min(pts[:, 0]))
        y_min = int(np.min(pts[:, 1]))
        x_max = int(np.max(pts[:, 0]))
        y_max = int(np.max(pts[:, 1]))
        x = x_max - x_min
        y = y_max - y_min
        is_horizontal_text = True
        if y / x > 1.5: # vertical text condition
            is_horizontal_text = False
        return is_horizontal_text


    def virtual_camera_to_world(self, size):
        ifu, ifv = self.ifu, self.ifv
        K, matT = self.K, self.matT

        ppu = size[0] / 2 + 1e-6
        ppv = size[1] / 2 + 1e-6

        P = np.zeros((size[1], size[0], 3))

        lu = np.array([i for i in range(size[0])])
        lv = np.array([i for i in range(size[1])])
        u, v = np.meshgrid(lu, lv)

        yp = (v - ppv) * ifv
        xp = (u - ppu) * ifu
        angle_a = arctan(sqrt(xp * xp + yp * yp))
        angle_b = arctan(yp / xp)

        D0 = sin(angle_a) * cos(angle_b)
        D1 = sin(angle_a) * sin(angle_b)
        D2 = cos(angle_a)

        D0[xp <= 0] = -D0[xp <= 0]
        D1[xp <= 0] = -D1[xp <= 0]

        ratio_a = K[0, 0] * D0 * D0 + K[1, 1] * D1 * D1 + K[2, 2] * D2 * D2 + \
                  (K[0, 1] + K[1, 0]) * D0 * D1 + (K[0, 2] + K[2, 0]) * D0 * D2 + (K[1, 2] + K[2, 1]) * D1 * D2
        ratio_b = (K[0, 3] + K[3, 0]) * D0 + (K[1, 3] + K[3, 1]) * D1 + (K[2, 3] + K[3, 2]) * D2
        ratio_c = K[3, 3] * np.ones(ratio_b.shape)

        delta = ratio_b * ratio_b - 4 * ratio_a * ratio_c
        t = np.zeros(delta.shape)
        t[ratio_a == 0] = -ratio_c[ratio_a == 0] / ratio_b[ratio_a == 0]
        t[ratio_a != 0] = (-ratio_b[ratio_a != 0] + sqrt(delta[ratio_a != 0])) / (2 * ratio_a[ratio_a != 0])
        t[delta < 0] = 0

        P[:, :, 0] = matT[0, 3] + t * (matT[0, 0] * D0 + matT[0, 1] * D1 + matT[0, 2] * D2)
        P[:, :, 1] = matT[1, 3] + t * (matT[1, 0] * D0 + matT[1, 1] * D1 + matT[1, 2] * D2)
        P[:, :, 2] = matT[2, 3] + t * (matT[2, 0] * D0 + matT[2, 1] * D1 + matT[2, 2] * D2)

        return P


    def world_to_image(self, image_size, world, intrinsic, distCoeffs, rotation, tvec):
        r11 = rotation[0, 0]
        r12 = rotation[0, 1]
        r13 = rotation[0, 2]
        r21 = rotation[1, 0]
        r22 = rotation[1, 1]
        r23 = rotation[1, 2]
        r31 = rotation[2, 0]
        r32 = rotation[2, 1]
        r33 = rotation[2, 2]

        t1 = tvec[0]
        t2 = tvec[1]
        t3 = tvec[2]

        k1 = distCoeffs[0]
        k2 = distCoeffs[1]
        p1 = distCoeffs[2]
        p2 = distCoeffs[3]
        k3 = distCoeffs[4]
        k4 = distCoeffs[5]
        k5 = distCoeffs[6]
        k6 = distCoeffs[7]

        if len(distCoeffs) > 8:
            s1 = distCoeffs[8]
            s2 = distCoeffs[9]
            s3 = distCoeffs[10]
            s4 = distCoeffs[11]
        else:
            s1 = s2 = s3 = s4 = 0

        if len(distCoeffs) > 12:
            tx = distCoeffs[12]
            ty = distCoeffs[13]
        else:
            tx = ty = 0

        fu = intrinsic[0, 0]
        fv = intrinsic[1, 1]
        ppu = intrinsic[0, 2]
        ppv = intrinsic[1, 2]

        cos_tx = cos(tx)
        cos_ty = cos(ty)
        sin_tx = sin(tx)
        sin_ty = sin(ty)

        tao11 = cos_ty * cos_tx * cos_ty + sin_ty * cos_tx * sin_ty
        tao12 = cos_ty * cos_tx * sin_ty * sin_tx - sin_ty * cos_tx * cos_ty * sin_tx
        tao13 = -cos_ty * cos_tx * sin_ty * cos_tx + sin_ty * cos_tx * cos_ty * cos_tx
        tao21 = -sin_tx * sin_ty
        tao22 = cos_ty * cos_tx * cos_tx + sin_tx * cos_ty * sin_tx
        tao23 = cos_ty * cos_tx * sin_tx - sin_tx * cos_ty * cos_tx

        P = np.zeros((image_size[1], image_size[0], 2))

        c3 = r31 * world[:, :, 0] + r32 * world[:, :, 1] + r33 * world[:, :, 2] + t3
        c1 = r11 * world[:, :, 0] + r12 * world[:, :, 1] + r13 * world[:, :, 2] + t1
        c2 = r21 * world[:, :, 0] + r22 * world[:, :, 1] + r23 * world[:, :, 2] + t2

        x1 = c1 / c3
        y1 = c2 / c3
        x12 = x1 * x1
        y12 = y1 * y1
        x1y1 = 2 * x1 * y1
        r2 = x12 + y12
        r4 = r2 * r2
        r6 = r2 * r4

        radial_distortion = (1 + k1 * r2 + k2 * r4 + k3 * r6) / (1 + k4 * r2 + k5 * r4 + k6 * r6)
        x2 = x1 * radial_distortion + p1 * x1y1 + p2 * (r2 + 2 * x12) + s1 * r2 + s2 * r4
        y2 = y1 * radial_distortion + p2 * x1y1 + p1 * (r2 + 2 * y12) + s3 * r2 + s4 * r4

        x3 = tao11 * x2 + tao12 * y2 + tao13
        y3 = tao21 * x2 + tao22 * y2 + tao23

        P[:, :, 0] = fu * x3 + ppu
        P[:, :, 1] = fv * y3 + ppv
        P[c3 <= 0] = 0

        return P


    def spatial_transform(self, image_data, new_image_size, mtx, dist, rvecs, tvecs, interpolation):
        rotation, _ = cv2.Rodrigues(rvecs)
        world_map = self.virtual_camera_to_world(new_image_size)
        image_map = self.world_to_image(new_image_size, world_map, mtx, dist, rotation, tvecs)
        image_map = image_map.astype(np.float32)
        dst = cv2.remap(image_data, image_map[:, :, 0], image_map[:, :, 1], interpolation)
        return dst


    def calibrate(self, org_size, image_coord, world_coord):
        """
        calibration
        :param org_size:
        :param image_coord:
        :param world_coord:
        :return:
        """
        # flag = cv2.CALIB_RATIONAL_MODEL | cv2.CALIB_TILTED_MODEL  | cv2.CALIB_THIN_PRISM_MODEL
        flag = cv2.CALIB_RATIONAL_MODEL
        flag2 = cv2.CALIB_RATIONAL_MODEL | cv2.CALIB_TILTED_MODEL
        flag3 = cv2.CALIB_RATIONAL_MODEL | cv2.CALIB_THIN_PRISM_MODEL
        flag4 = cv2.CALIB_RATIONAL_MODEL | cv2.CALIB_ZERO_TANGENT_DIST | cv2.CALIB_FIX_ASPECT_RATIO
        flag5 = cv2.CALIB_RATIONAL_MODEL | cv2.CALIB_TILTED_MODEL | cv2.CALIB_ZERO_TANGENT_DIST
        flag6 = cv2.CALIB_RATIONAL_MODEL | cv2.CALIB_FIX_ASPECT_RATIO
        flag_list = [flag2, flag3, flag4, flag5, flag6]

        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(world_coord.astype(np.float32),
                                                                image_coord.astype(np.float32),
                                                                org_size,
                                                                None,
                                                                None,
                                                                flags=flag)
        if ret > 2:
            # strategies
            min_ret = ret
            for i, flag in enumerate(flag_list):
                _ret, _mtx, _dist, _rvecs, _tvecs = cv2.calibrateCamera(world_coord.astype(np.float32),
                                                                   image_coord.astype(np.float32),
                                                                   org_size,
                                                                   None,
                                                                   None,
                                                                   flags=flag)
                if _ret < min_ret:
                    min_ret = _ret
                    ret, mtx, dist, rvecs, tvecs = _ret, _mtx, _dist, _rvecs, _tvecs

        return ret, mtx, dist, rvecs, tvecs


    def dc_homo(self, img, img_points, obj_points, is_horizontal_text, interpolation=cv2.INTER_LINEAR,
                ratio_width=1.0, ratio_height=1.0):
        """
        divide and conquer: homography
        # ratio_width and ratio_height must be 1.0 here
        """
        _img_points = img_points.reshape(-1, 2)
        _obj_points = obj_points.reshape(-1, 3)

        homo_img_list = []
        width_list = []
        height_list = []
        # divide and conquer
        for i in range(len(_img_points) // 2 - 1):
            new_img_points = np.zeros((4, 2)).astype(np.float32)
            new_obj_points = np.zeros((4, 2)).astype(np.float32)

            new_img_points[0:2, :] = _img_points[i:(i + 2), :2]
            new_img_points[2:4, :] = _img_points[::-1, :][i:(i + 2), :2][::-1, :]

            new_obj_points[0:2, :] = _obj_points[i:(i + 2), :2]
            new_obj_points[2:4, :] = _obj_points[::-1, :][i:(i + 2), :2][::-1, :]

            if is_horizontal_text:
                world_width = np.abs(new_obj_points[1, 0] - new_obj_points[0, 0])
                world_height = np.abs(new_obj_points[3, 1] - new_obj_points[0, 1])
            else:
                world_width = np.abs(new_obj_points[1, 1] - new_obj_points[0, 1])
                world_height = np.abs(new_obj_points[3, 0] - new_obj_points[0, 0])

            homo_img = Homography(img, new_img_points, world_width, world_height,
                                              interpolation=interpolation,
                                              ratio_width=ratio_width, ratio_height=ratio_height)

            homo_img_list.append(homo_img)
            _h, _w = homo_img.shape[:2]
            width_list.append(_w)
            height_list.append(_h)

        # stitching
        rectified_image = np.zeros((np.max(height_list), sum(width_list), 3)).astype(np.uint8)

        st = 0
        for (homo_img, w, h) in zip(homo_img_list, width_list, height_list):
            rectified_image[:h, st:st + w, :] = homo_img
            st += w

        if not is_horizontal_text:
            # vertical rotation
            rectified_image = np.rot90(rectified_image, 3)

        return rectified_image


    def __call__(self, image_data, points, interpolation=cv2.INTER_LINEAR, ratio_width=1.0, ratio_height=1.0, mode='calibration'):
        """
        spatial transform for a poly text
        :param image_data:
        :param points: [x1,y1,x2,y2,x3,y3,...], clockwise order, (x1,y1) must be the top-left of first char.
        :param interpolation: cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4
        :param ratio_width:  roi_image width expansion. It should not be smaller than 1.0
        :param ratio_height: roi_image height expansion. It should not be smaller than 1.0
        :param mode: 'calibration' or 'homography'. when homography, ratio_width and ratio_height must be 1.0
        :return:
        """
        org_h, org_w = image_data.shape[:2]
        org_size = (org_w, org_h)
        self.image = image_data

        is_horizontal_text = self.horizontal_text_estimate(points)
        if is_horizontal_text:
            image_coord, world_coord, new_image_size = self.horizontal_text_process(points)
        else:
            image_coord, world_coord, new_image_size = self.vertical_text_process(points, org_size)

        if mode.lower() == 'calibration':
            ret, mtx, dist, rvecs, tvecs = self.calibrate(org_size, image_coord, world_coord)

            st_size = (int(new_image_size[0]*ratio_width), int(new_image_size[1]*ratio_height))
            dst = self.spatial_transform(image_data, st_size, mtx, dist[0], rvecs[0], tvecs[0], interpolation)
        elif mode.lower() == 'homography':
            # ratio_width and ratio_height must be 1.0 here and ret set to 0.01 without loss manually
            ret = 0.01
            dst = self.dc_homo(image_data, image_coord, world_coord, is_horizontal_text,
                               interpolation=interpolation, ratio_width=1.0, ratio_height=1.0)
        else:
            raise ValueError('mode must be ["calibration", "homography"], but got {}'.format(mode))

        return dst, ret


class PlanB:
    def __call__(self, image, points, curveTextRectifier, interpolation=cv2.INTER_LINEAR,
                 ratio_width=1.0, ratio_height=1.0, loss_thresh=5.0, square=False):
        """
        Plan B using sub-image when it failed in original image
        :param image:
        :param points:
        :param curveTextRectifier: CurveTextRectifier
        :param interpolation: cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4
        :param ratio_width:  roi_image width expansion. It should not be smaller than 1.0
        :param ratio_height: roi_image height expansion. It should not be smaller than 1.0
        :param loss_thresh: if loss greater than loss_thresh --> get_rotate_crop_image
        :param square: crop square image or not. True or False. The default is False
        :return:
        """
        h, w = image.shape[:2]
        _points = np.array(points).reshape(-1, 2).astype(np.float32)
        x_min = int(np.min(_points[:, 0]))
        y_min = int(np.min(_points[:, 1]))
        x_max = int(np.max(_points[:, 0]))
        y_max = int(np.max(_points[:, 1]))
        dx = x_max - x_min
        dy = y_max - y_min
        max_d = max(dx, dy)
        mean_pt = np.mean(_points, 0)

        expand_x = (ratio_width - 1.0) * 0.5 * max_d
        expand_y = (ratio_height - 1.0) * 0.5 * max_d

        if square:
            x_min = np.clip(int(mean_pt[0] - max_d - expand_x), 0, w - 1)
            y_min = np.clip(int(mean_pt[1] - max_d - expand_y), 0, h - 1)
            x_max = np.clip(int(mean_pt[0] + max_d + expand_x), 0, w - 1)
            y_max = np.clip(int(mean_pt[1] + max_d + expand_y), 0, h - 1)
        else:
            x_min = np.clip(int(x_min - expand_x), 0, w - 1)
            y_min = np.clip(int(y_min - expand_y), 0, h - 1)
            x_max = np.clip(int(x_max + expand_x), 0, w - 1)
            y_max = np.clip(int(y_max + expand_y), 0, h - 1)

        new_image = image[y_min:y_max, x_min:x_max, :].copy()
        new_points = _points.copy()
        new_points[:, 0] -= x_min
        new_points[:, 1] -= y_min

        dst_img, loss = curveTextRectifier(new_image, new_points, interpolation, ratio_width, ratio_height, mode='calibration')

        return dst_img, loss


class AutoRectifier:
    def __init__(self):
        self.npoints = 10
        self.curveTextRectifier = CurveTextRectifier()

    @staticmethod
    def get_rotate_crop_image(img, points, interpolation=cv2.INTER_CUBIC, ratio_width=1.0, ratio_height=1.0):
        """
        crop or homography
        :param img:
        :param points:
        :param interpolation:
        :param ratio_width:
        :param ratio_height:
        :return:
        """
        h, w = img.shape[:2]
        _points = np.array(points).reshape(-1, 2).astype(np.float32)

        if len(_points) != 4:
            x_min = int(np.min(_points[:, 0]))
            y_min = int(np.min(_points[:, 1]))
            x_max = int(np.max(_points[:, 0]))
            y_max = int(np.max(_points[:, 1]))
            dx = x_max - x_min
            dy = y_max - y_min
            expand_x = int(0.5 * dx * (ratio_width - 1))
            expand_y = int(0.5 * dy * (ratio_height - 1))
            x_min = np.clip(int(x_min - expand_x), 0, w - 1)
            y_min = np.clip(int(y_min - expand_y), 0, h - 1)
            x_max = np.clip(int(x_max + expand_x), 0, w - 1)
            y_max = np.clip(int(y_max + expand_y), 0, h - 1)

            dst_img = img[y_min:y_max, x_min:x_max, :].copy()
        else:
            img_crop_width = int(
                max(
                    np.linalg.norm(_points[0] - _points[1]),
                    np.linalg.norm(_points[2] - _points[3])))
            img_crop_height = int(
                max(
                    np.linalg.norm(_points[0] - _points[3]),
                    np.linalg.norm(_points[1] - _points[2])))

            dst_img = Homography(img, _points, img_crop_width, img_crop_height, interpolation, ratio_width, ratio_height)

        return dst_img


    def visualize(self, image_data, points_list):
        visualization = image_data.copy()

        for box in points_list:
            box = np.array(box).reshape(-1, 2).astype(np.int32)
            cv2.drawContours(visualization, [np.array(box).reshape((-1, 1, 2))], -1, (0, 0, 255), 2)
            for i, p in enumerate(box):
                if i != 0:
                    cv2.circle(visualization, tuple(p), radius=1, color=(255, 0, 0), thickness=2)
                else:
                    cv2.circle(visualization, tuple(p), radius=1, color=(255, 255, 0), thickness=2)
        return visualization


    def __call__(self, image_data, points, interpolation=cv2.INTER_LINEAR,
                 ratio_width=1.0, ratio_height=1.0, loss_thresh=5.0, mode='calibration'):
        """
        rectification in strategies for a poly text
        :param image_data:
        :param points: [x1,y1,x2,y2,x3,y3,...], clockwise order, (x1,y1) must be the top-left of first char.
        :param interpolation: cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4
        :param ratio_width:  roi_image width expansion. It should not be smaller than 1.0
        :param ratio_height: roi_image height expansion. It should not be smaller than 1.0
        :param loss_thresh: if loss greater than loss_thresh --> get_rotate_crop_image
        :param mode: 'calibration' or 'homography'. when homography, ratio_width and ratio_height must be 1.0
        :return:
        """
        _points = np.array(points).reshape(-1,2)
        if len(_points) >= self.npoints and len(_points) % 2 == 0:
            try:
                curveTextRectifier = CurveTextRectifier()

                dst_img, loss = curveTextRectifier(image_data, points, interpolation, ratio_width, ratio_height, mode)
                if loss >= 2:
                    # for robust
                    # large loss means it cannot be reconstruct correctly, we must find other way to reconstruct
                    img_list, loss_list = [dst_img], [loss]
                    _dst_img, _loss = PlanB()(image_data, points, curveTextRectifier,
                                              interpolation, ratio_width, ratio_height,
                                              loss_thresh=loss_thresh,
                                              square=True)
                    img_list += [_dst_img]
                    loss_list += [_loss]

                    _dst_img, _loss = PlanB()(image_data, points, curveTextRectifier,
                                              interpolation, ratio_width, ratio_height,
                                              loss_thresh=loss_thresh, square=False)
                    img_list += [_dst_img]
                    loss_list += [_loss]

                    min_loss = min(loss_list)
                    dst_img = img_list[loss_list.index(min_loss)]

                    if min_loss >= loss_thresh:
                        print('calibration loss: {} is too large for spatial transformer. It is failed. Using get_rotate_crop_image'.format(loss))
                        dst_img = self.get_rotate_crop_image(image_data, points, interpolation, ratio_width, ratio_height)
                        print('here')
            except Exception as e:
                print(e)
                dst_img = self.get_rotate_crop_image(image_data, points, interpolation, ratio_width, ratio_height)
        else:
            dst_img = self.get_rotate_crop_image(image_data, _points, interpolation, ratio_width, ratio_height)

        return dst_img


    def run(self, image_data, points_list, interpolation=cv2.INTER_LINEAR,
            ratio_width=1.0, ratio_height=1.0, loss_thresh=5.0, mode='calibration'):
        """
        run for texts in an image
        :param image_data: numpy.ndarray. The shape is [h, w, 3]
        :param points_list: [[x1,y1,x2,y2,x3,y3,...], [x1,y1,x2,y2,x3,y3,...], ...], clockwise order, (x1,y1) must be the top-left of first char.
        :param interpolation: cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_AREA, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4
        :param ratio_width:  roi_image width expansion. It should not be smaller than 1.0
        :param ratio_height: roi_image height expansion. It should not be smaller than 1.0
        :param loss_thresh: if loss greater than loss_thresh --> get_rotate_crop_image
        :param mode: 'calibration' or 'homography'. when homography, ratio_width and ratio_height must be 1.0
        :return: res: roi-image list, visualized_image: draw polys in original image
        """
        if image_data is None:
            raise ValueError
        if not isinstance(points_list, list):
            raise ValueError
        for points in points_list:
            if not isinstance(points, list):
                raise ValueError

        if ratio_width < 1.0 or ratio_height < 1.0:
            raise ValueError('ratio_width and ratio_height cannot be smaller than 1, but got {}', (ratio_width, ratio_height))

        if mode.lower() != 'calibration' and mode.lower() != 'homography':
            raise ValueError('mode must be ["calibration", "homography"], but got {}'.format(mode))

        if mode.lower() == 'homography' and ratio_width != 1.0 and ratio_height != 1.0:
            raise ValueError('ratio_width and ratio_height must be 1.0 when mode is homography, but got mode:{}, ratio:({},{})'.format(mode, ratio_width, ratio_height))

        res = []
        for points in points_list:
            rectified_img = self(image_data, points, interpolation, ratio_width, ratio_height,
                                 loss_thresh=loss_thresh, mode=mode)
            res.append(rectified_img)

        # visualize
        visualized_image = self.visualize(image_data, points_list)

        return res, visualized_image



if __name__ == '__main__':
    # test for a poly text
    print('begin')

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', type=str, help='Assign the image path')
    parser.add_argument('--txt', type=str, help='Assign the path of .txt to get points',
                        default=None)
    parser.add_argument('--mode', type=str, help='Assign the mode: calibration or homography',
                        default='calibration')
    args = parser.parse_args()

    test_points = [106, 167, 147, 143, 264, 94, 319, 85, 443, 98, 492, 116, 557,
                     148, 539, 181, 474, 149, 435, 135, 319, 122, 274, 130, 166, 176,
                     125, 200]

    image_path = './data/test_images/arc_text.jpg'

    if not os.path.exists(image_path):
        raise FileNotFoundError

    with open(image_path, "rb") as f:
        image = np.array(bytearray(f.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

    if image is None:
        raise ValueError('the image of {} is None'.format(image_path))

    autoRectifier = AutoRectifier()
    res = autoRectifier(image, test_points, interpolation=cv2.INTER_LINEAR,
                        ratio_width=1.0, ratio_height=1.0,
                        loss_thresh=5.0, mode=args.mode)

    save_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results')
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    cv2.imwrite(os.path.join(save_path, 'OneBoxImageResult.jpg'), res)
    print('{} is saved.'.format(os.path.join(save_path, 'OneBoxImageResult.jpg')))
    print('done!')